from __future__ import annotations

import asyncio
import importlib
from collections.abc import AsyncIterator
from datetime import UTC, datetime, timedelta
from pathlib import Path

import pytest
import pytest_asyncio
from alembic import command
from backup_tool.config import Settings
from backup_tool.db.engine import create_engine
from backup_tool.db.models import Execution, Job, Repository, Source
from backup_tool.execution import (
    EnqueueError,
    claim,
    complete_cancellation,
    enqueue,
    heartbeat,
    request_cancellation,
    retry,
)
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker

cli = importlib.import_module("backup_tool.cli")


@pytest_asyncio.fixture
async def database(
    tmp_path: Path,
) -> AsyncIterator[tuple[async_sessionmaker[AsyncSession], AsyncEngine]]:
    key = tmp_path / "master.key"
    key.write_bytes(b"m5-test-master-key-material-32-bytes-minimum")
    key.chmod(0o600)
    data_dir = tmp_path / "data"
    repositories = tmp_path / "repositories"
    sources = tmp_path / "sources"
    restores = tmp_path / "restores"
    for directory in (data_dir, repositories, sources, restores):
        directory.mkdir()
    settings = Settings(
        data_dir=data_dir,
        database_url=f"sqlite+aiosqlite:///{data_dir / 'metadata.db'}",
        repository_roots=(repositories,),
        local_source_roots=(sources,),
        restore_roots=(restores,),
        master_key_file=key,
    )
    command.upgrade(cli.build_alembic_config(settings), "head")
    engine = create_engine(settings)
    yield async_sessionmaker(engine, expire_on_commit=False), engine
    await engine.dispose()


async def create_job(
    db: AsyncSession, suffix: str, *, enabled: bool = True, state: str = "active"
) -> str:
    repository = Repository(
        name=f"repository-{suffix}",
        root=f"/repositories/{suffix}",
        format_version=1,
        compression="none",
        encryption="none",
    )
    source = Source(
        name=f"source-{suffix}",
        kind="local",
        public_config={"root": f"/sources/{suffix}"},
        secret_refs=[],
    )
    db.add_all([repository, source])
    await db.flush()
    job = Job(
        name=f"job-{suffix}",
        source_id=source.id,
        repository_id=repository.id,
        requested_mode="full",
        exclusions=[],
        retention={},
        enabled=enabled,
        state=state,
    )
    db.add(job)
    await db.commit()
    return job.id


@pytest.mark.asyncio
async def test_concurrent_enqueue_allows_exactly_one_active_execution(
    database: tuple[async_sessionmaker[AsyncSession], AsyncEngine],
) -> None:
    sessions, _ = database
    async with sessions() as db:
        job_id = await create_job(db, "concurrent")

    async def start() -> Execution | EnqueueError:
        async with sessions() as db:
            try:
                return await enqueue(db, job_id)
            except EnqueueError as error:
                return error

    first, second = await asyncio.gather(start(), start())
    results = [first, second]
    successes = [result for result in results if isinstance(result, Execution)]
    failures = [result for result in results if isinstance(result, EnqueueError)]

    assert len(successes) == 1
    assert len(failures) == 1
    assert failures[0].code == "execution_active"
    assert failures[0].active_execution_id == successes[0].id
    async with sessions() as db:
        executions = list(await db.scalars(select(Execution).where(Execution.job_id == job_id)))
    assert [execution.id for execution in executions] == [successes[0].id]


@pytest.mark.asyncio
@pytest.mark.parametrize(("enabled", "state"), [(False, "active"), (True, "archived")])
async def test_enqueue_rejects_disabled_or_archived_jobs(
    database: tuple[async_sessionmaker[AsyncSession], AsyncEngine],
    enabled: bool,
    state: str,
) -> None:
    sessions, _ = database
    async with sessions() as db:
        job_id = await create_job(
            db, f"unavailable-{enabled}-{state}", enabled=enabled, state=state
        )
        with pytest.raises(EnqueueError) as raised:
            await enqueue(db, job_id)

    assert raised.value.code == "job_disabled"


@pytest.mark.asyncio
async def test_reclaimed_lease_fences_the_previous_worker(
    database: tuple[async_sessionmaker[AsyncSession], AsyncEngine],
) -> None:
    sessions, _ = database
    async with sessions() as db:
        execution = await enqueue(db, await create_job(db, "leases"))
        assert await claim(db, execution.id, "worker-a") is not None
        assert await heartbeat(db, execution.id, "worker-a")
        persisted = await db.get(Execution, execution.id)
        assert persisted is not None
        persisted.lease_expires_at = datetime.now(UTC) - timedelta(seconds=1)
        await db.commit()

    async with sessions() as db:
        assert await claim(db, execution.id, "worker-b") is not None
        assert not await heartbeat(db, execution.id, "worker-a")
        assert await request_cancellation(db, execution.id) is not None
        assert not await complete_cancellation(db, execution.id, "worker-a")
        assert await complete_cancellation(db, execution.id, "worker-b")

    async with sessions() as db:
        persisted = await db.get(Execution, execution.id)
    assert persisted is not None
    assert persisted.state == "cancelled"
    assert persisted.lease_owner is None


@pytest.mark.asyncio
async def test_retry_reuses_execution_and_rejects_non_transient_failures(
    database: tuple[async_sessionmaker[AsyncSession], AsyncEngine],
) -> None:
    sessions, _ = database
    async with sessions() as db:
        execution = await enqueue(db, await create_job(db, "retry"))
        execution.state = "failed"
        execution.reason_code = "transient_io"
        execution.operator_message = "temporary failure"
        execution.lease_owner = "worker-a"
        execution.lease_expires_at = datetime.now(UTC) + timedelta(seconds=60)
        await db.commit()

        retried = await retry(db, execution.id)
        assert retried is not None
        assert retried.id == execution.id
        assert retried.state == "queued"
        assert retried.attempt == 2
        assert retried.reason_code is None
        assert retried.operator_message is None
        assert retried.lease_owner is None
        assert retried.lease_expires_at is None

        retried.state = "failed"
        retried.reason_code = "integrity_failure"
        await db.commit()
        with pytest.raises(EnqueueError) as raised:
            await retry(db, execution.id)

    assert raised.value.code == "retry_not_allowed"