from __future__ import annotations

import asyncio
import logging
import signal
import threading
import uuid
from collections.abc import Sequence
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import cast
from unittest.mock import MagicMock

import pytest

from app.core.config import Settings
from app.persistence.job_queue import (
    BackgroundJob,
    JobLease,
    JobState,
    LeaseHeartbeat,
    LeaseLostError,
)
from app.worker import Worker, WorkerConfig, build_default_handlers, install_signal_handlers

NOW = datetime(2026, 7, 13, 8, 0, tzinfo=UTC)
JOB_ID = uuid.UUID("40000000-0000-0000-0000-000000000004")
RESOURCE_ID = uuid.UUID("50000000-0000-0000-0000-000000000005")
LEASE_TOKEN = uuid.UUID("60000000-0000-0000-0000-000000000006")


def _job(job_type: str = "KNOWN_JOB") -> BackgroundJob:
    lease = JobLease(JOB_ID, "worker-a", LEASE_TOKEN)
    return BackgroundJob(
        id=JOB_ID,
        job_type=job_type,
        required_capability="embedding",
        resource_type="document_version",
        resource_id=RESOURCE_ID,
        idempotency_key="job:one",
        payload={"resource_id": str(RESOURCE_ID)},
        stage="PROCESSING",
        progress=0,
        priority=0,
        attempt=1,
        max_attempts=3,
        run_after=NOW,
        lease_until=NOW + timedelta(seconds=60),
        created_at=NOW,
        updated_at=NOW,
        lease=lease,
    )


def _state(status: str) -> JobState:
    return JobState(
        job_id=JOB_ID,
        status=status,
        attempt=1,
        max_attempts=3,
        finished_at=NOW if status in {"SUCCEEDED", "FAILED"} else None,
    )


class FakeQueue:
    def __init__(self, job: BackgroundJob | None = None) -> None:
        self.job = job
        self.claim_count = 0
        self.claim_called = threading.Event()
        self.in_claim = False
        self.heartbeat_count = 0
        self.complete_leases: list[JobLease] = []
        self.failures: list[tuple[JobLease, str, str, int]] = []
        self.reap_count = 0
        self.heartbeat_error: Exception | None = None
        self.complete_error: Exception | None = None

    def claim(
        self,
        *,
        worker_id: str,
        worker_capabilities: Sequence[str],
        lease_seconds: int,
    ) -> BackgroundJob | None:
        assert worker_id == "worker-a"
        assert tuple(worker_capabilities) == ("embedding",)
        assert lease_seconds == 1
        self.claim_count += 1
        self.claim_called.set()
        self.in_claim = True
        try:
            result = self.job
            self.job = None
            return result
        finally:
            self.in_claim = False

    def heartbeat(self, lease: JobLease, *, lease_seconds: int) -> LeaseHeartbeat:
        assert lease == JobLease(JOB_ID, "worker-a", LEASE_TOKEN)
        assert lease_seconds == 1
        self.heartbeat_count += 1
        if self.heartbeat_error is not None:
            raise self.heartbeat_error
        return LeaseHeartbeat(JOB_ID, NOW + timedelta(seconds=1))

    def complete(self, lease: JobLease) -> JobState:
        self.complete_leases.append(lease)
        if self.complete_error is not None:
            raise self.complete_error
        return _state("SUCCEEDED")

    def fail_or_retry(
        self,
        lease: JobLease,
        *,
        error_code: str,
        error_message: str,
        retry_delay_seconds: int,
    ) -> JobState:
        self.failures.append((lease, error_code, error_message, retry_delay_seconds))
        return _state("QUEUED")

    def reap_expired(
        self,
        *,
        lock_key: int,
        batch_size: int = 100,
    ) -> tuple[JobState, ...]:
        assert isinstance(lock_key, int)
        assert batch_size == 10
        self.reap_count += 1
        return ()


def _config(
    *,
    capabilities: tuple[str, ...] = ("embedding",),
    heartbeat_seconds: float = 0.01,
    poll_seconds: float = 0.01,
    reaper_batch_size: int = 10,
) -> WorkerConfig:
    return WorkerConfig(
        worker_id="worker-a",
        capabilities=capabilities,
        lease_seconds=1,
        heartbeat_seconds=heartbeat_seconds,
        poll_seconds=poll_seconds,
        retry_delay_seconds=7,
        reaper_interval_seconds=30.0,
        reaper_batch_size=reaper_batch_size,
        reaper_lock_key=42,
    )


@pytest.mark.asyncio
async def test_handler_runs_after_claim_transaction_and_completes_with_same_lease() -> None:
    queue = FakeQueue(_job())
    observed_payload: dict[str, object] = {}

    async def handler(job: BackgroundJob) -> None:
        assert queue.in_claim is False
        observed_payload.update(job.payload)

    worker = Worker(queue, _config(), handlers={"KNOWN_JOB": handler})

    worked = await worker.run_once()

    assert worked is True
    assert observed_payload == {"resource_id": str(RESOURCE_ID)}
    assert queue.complete_leases == [JobLease(JOB_ID, "worker-a", LEASE_TOKEN)]
    assert queue.failures == []


@pytest.mark.asyncio
async def test_long_handler_is_heartbeated_before_completion() -> None:
    queue = FakeQueue(_job())

    async def handler(_job: BackgroundJob) -> None:
        await asyncio.sleep(0.035)

    worker = Worker(queue, _config(), handlers={"KNOWN_JOB": handler})

    await worker.run_once()

    assert queue.heartbeat_count >= 2
    assert queue.complete_leases == [JobLease(JOB_ID, "worker-a", LEASE_TOKEN)]


@pytest.mark.asyncio
async def test_heartbeat_lease_loss_cancels_handler_without_terminal_write() -> None:
    queue = FakeQueue(_job())
    queue.heartbeat_error = LeaseLostError("lease moved")
    handler_cancelled = asyncio.Event()

    async def handler(_job: BackgroundJob) -> None:
        try:
            await asyncio.Event().wait()
        finally:
            handler_cancelled.set()

    worker = Worker(queue, _config(), handlers={"KNOWN_JOB": handler})

    await worker.run_once()

    assert handler_cancelled.is_set()
    assert queue.complete_leases == []
    assert queue.failures == []


@pytest.mark.asyncio
async def test_unknown_job_is_safely_failed_without_handler_execution() -> None:
    job = _job("UNREGISTERED_JOB")
    queue = FakeQueue(job)
    worker = Worker(queue, _config(), handlers={})

    await worker.run_once()

    assert queue.complete_leases == []
    assert queue.failures == [
        (
            job.lease,
            "UNKNOWN_JOB_TYPE",
            "No registered handler exists for this job type.",
            7,
        )
    ]


@pytest.mark.asyncio
async def test_handler_exception_is_redacted_before_queue_failure(
    caplog: pytest.LogCaptureFixture,
) -> None:
    job = _job()
    queue = FakeQueue(job)

    async def handler(_job: BackgroundJob) -> None:
        raise RuntimeError("private-document-text database-password")

    worker = Worker(queue, _config(), handlers={"KNOWN_JOB": handler})

    caplog.set_level(logging.ERROR, logger="geological_rag.worker")
    await worker.run_once()

    assert len(queue.failures) == 1
    lease, code, message, delay = queue.failures[0]
    assert lease == job.lease
    assert code == "JOB_HANDLER_FAILED"
    assert message == "Registered job handler failed."
    assert "private-document-text" not in message
    assert delay == 7
    assert [getattr(record, "error_type", None) for record in caplog.records] == ["RuntimeError"]
    assert "private-document-text" not in caplog.text


@pytest.mark.asyncio
async def test_completion_fence_rejection_does_not_retry_completed_handler() -> None:
    queue = FakeQueue(_job())
    queue.complete_error = LeaseLostError("lease moved")
    calls = 0

    async def handler(_job: BackgroundJob) -> None:
        nonlocal calls
        calls += 1

    worker = Worker(queue, _config(), handlers={"KNOWN_JOB": handler})

    await worker.run_once()

    assert calls == 1
    assert queue.failures == []


@pytest.mark.asyncio
async def test_reaper_is_rate_limited_by_monotonic_schedule() -> None:
    queue = FakeQueue()
    clock_value = 100.0
    worker = Worker(
        queue,
        _config(),
        handlers={},
        monotonic=lambda: clock_value,
    )

    assert await worker.run_once() is False
    assert await worker.run_once() is False
    assert queue.reap_count == 1

    clock_value = 131.0
    assert await worker.run_once() is False
    assert queue.reap_count == 2


@pytest.mark.asyncio
async def test_stop_wakes_poll_and_prevents_new_claims() -> None:
    queue = FakeQueue()
    worker = Worker(queue, _config(poll_seconds=30.0), handlers={})
    task = asyncio.create_task(worker.run())

    assert await asyncio.to_thread(queue.claim_called.wait, 0.2)
    worker.request_stop()
    await asyncio.wait_for(task, timeout=0.2)

    assert worker.stopping is True
    assert queue.claim_count == 1


def test_sigterm_callback_requests_graceful_stop() -> None:
    worker = Worker(FakeQueue(), _config(), handlers={})
    loop = MagicMock(spec=asyncio.AbstractEventLoop)

    installed = install_signal_handlers(worker, loop)

    assert signal.SIGTERM in installed
    sigterm_call = next(
        call for call in loop.add_signal_handler.call_args_list if call.args[0] == signal.SIGTERM
    )
    sigterm_call.args[1]()
    assert worker.stopping is True


@pytest.mark.parametrize(
    ("field", "value"),
    [
        ("capabilities", ()),
        ("capabilities", ("embedding", "embedding")),
        ("heartbeat_seconds", 0.34),
        ("heartbeat_seconds", 1.0),
        ("reaper_batch_size", 0),
    ],
)
def test_worker_config_rejects_unsafe_lease_or_routing_values(
    field: str,
    value: object,
) -> None:
    with pytest.raises(ValueError):
        if field == "capabilities":
            assert isinstance(value, tuple)
            _config(capabilities=cast(tuple[str, ...], value))
        elif field == "heartbeat_seconds":
            assert isinstance(value, float)
            _config(heartbeat_seconds=value)
        else:
            assert isinstance(value, int)
            _config(reaper_batch_size=value)


def test_default_handler_registry_is_capability_isolated(tmp_path: Path) -> None:
    settings = Settings(upload_root=tmp_path.resolve())

    local_handlers = build_default_handlers(settings, ("document_parse",))

    assert set(local_handlers) == {"PARSE_DOCUMENT"}


def test_default_handler_registry_rejects_unimplemented_capabilities(tmp_path: Path) -> None:
    settings = Settings(upload_root=tmp_path.resolve())

    with pytest.raises(RuntimeError, match="unsupported worker capabilities: evaluation"):
        build_default_handlers(settings, ("evaluation",))