from __future__ import annotations

import asyncio
import hashlib
import json
import secrets
from dataclasses import dataclass, field
from typing import Any, AsyncIterator

from .models import MediaJobRequest, MediaJobView, StoredEvent, utc_now


TERMINAL = {"completed", "failed", "cancelled"}


class JobCapacityExceeded(RuntimeError):
    pass


class IdempotencyConflict(RuntimeError):
    pass


class JobNotFound(KeyError):
    pass


class ReplayGap(RuntimeError):
    pass


@dataclass(slots=True)
class JobRecord:
    request: MediaJobRequest
    view: MediaJobView
    idempotency_key: str
    request_hash: str
    lease_id: str
    endpoint: str
    events: list[StoredEvent] = field(default_factory=list)
    next_seq: int = 1
    task: asyncio.Task[None] | None = None
    downstream_id: str | None = None


class JobStore:
    def __init__(
        self,
        *,
        max_active: int,
        max_records: int,
        max_events: int,
        max_subscribers: int,
    ) -> None:
        self.max_active = max_active
        self.max_records = max_records
        self.max_events = max_events
        self.max_subscribers = max_subscribers
        self._jobs: dict[str, JobRecord] = {}
        self._idempotency: dict[str, tuple[str, str]] = {}
        self._conditions: dict[str, asyncio.Condition] = {}
        self._subscribers = 0
        self._lock = asyncio.Lock()

    @staticmethod
    def digest(request: MediaJobRequest, lease_id: str) -> str:
        canonical = json.dumps(
            {"request": request.model_dump(mode="json"), "lease_id": lease_id},
            ensure_ascii=True,
            sort_keys=True,
            separators=(",", ":"),
        )
        return hashlib.sha256(canonical.encode("utf-8")).hexdigest()

    def _event_locked(self, record: JobRecord, event_type: str, payload: dict[str, Any]) -> None:
        terminal = event_type.rsplit(".", 1)[-1] in TERMINAL
        if len(record.events) >= self.max_events:
            if not terminal:
                return
            record.events.pop(0)
        seq = record.next_seq
        record.next_seq += 1
        record.events.append(
            StoredEvent(
                seq=seq,
                event_id=str(seq),
                type=event_type,
                payload=payload,
            )
        )

    def _prune_locked(self) -> None:
        if len(self._jobs) < self.max_records:
            return
        terminal = sorted(
            (item for item in self._jobs.values() if item.view.status in TERMINAL),
            key=lambda item: item.view.updated_at,
        )
        while terminal and len(self._jobs) >= self.max_records:
            record = terminal.pop(0)
            self._jobs.pop(record.view.job_id, None)
            self._conditions.pop(record.view.job_id, None)
            self._idempotency.pop(record.idempotency_key, None)

    async def create(
        self,
        request: MediaJobRequest,
        *,
        lease_id: str,
        endpoint: str,
        idempotency_key: str,
    ) -> tuple[MediaJobView, bool]:
        digest = self.digest(request, lease_id)
        async with self._lock:
            existing = self._idempotency.get(idempotency_key)
            if existing:
                if existing[0] != digest:
                    raise IdempotencyConflict("idempotency key reused with another request")
                record = self._jobs.get(existing[1])
                if record is None:
                    raise IdempotencyConflict("idempotent job record is no longer retained")
                return record.view.model_copy(deep=True), False
            active = sum(1 for item in self._jobs.values() if item.view.status not in TERMINAL)
            if active >= self.max_active:
                raise JobCapacityExceeded("active media job capacity reached")
            self._prune_locked()
            if len(self._jobs) >= self.max_records:
                raise JobCapacityExceeded("media job record capacity reached")
            now = utc_now()
            view = MediaJobView(
                job_id=f"job_{secrets.token_hex(16)}",
                kind=request.kind,
                status="queued",
                created_at=now,
                updated_at=now,
            )
            record = JobRecord(
                request=request.model_copy(deep=True),
                view=view,
                idempotency_key=idempotency_key,
                request_hash=digest,
                lease_id=lease_id,
                endpoint=endpoint,
            )
            self._jobs[view.job_id] = record
            self._idempotency[idempotency_key] = (digest, view.job_id)
            self._conditions[view.job_id] = asyncio.Condition(self._lock)
            common = {"job_id": view.job_id, "kind": view.kind}
            self._event_locked(record, "job.created", common)
            self._event_locked(record, "job.queued", common)
            return view.model_copy(deep=True), True

    async def replay(
        self,
        request: MediaJobRequest,
        *,
        lease_id: str,
        idempotency_key: str,
    ) -> MediaJobView | None:
        """Resolve a retry before consulting an already-released lease."""
        digest = self.digest(request, lease_id)
        async with self._lock:
            existing = self._idempotency.get(idempotency_key)
            if not existing:
                return None
            if existing[0] != digest:
                raise IdempotencyConflict("idempotency key reused with another request")
            record = self._jobs.get(existing[1])
            if record is None:
                raise IdempotencyConflict("idempotent job record is no longer retained")
            return record.view.model_copy(deep=True)

    async def set_task(self, job_id: str, task: asyncio.Task[None]) -> None:
        async with self._lock:
            self._jobs[job_id].task = task

    async def set_downstream(self, job_id: str, downstream_id: str) -> None:
        async with self._lock:
            self._jobs[job_id].downstream_id = downstream_id

    async def runtime(self, job_id: str) -> tuple[MediaJobRequest, str, str, str | None, asyncio.Task[None] | None]:
        async with self._lock:
            record = self._jobs.get(job_id)
            if record is None:
                raise JobNotFound(job_id)
            return (
                record.request.model_copy(deep=True),
                record.lease_id,
                record.endpoint,
                record.downstream_id,
                record.task,
            )

    async def view(self, job_id: str) -> MediaJobView:
        async with self._lock:
            record = self._jobs.get(job_id)
            if record is None:
                raise JobNotFound(job_id)
            return record.view.model_copy(deep=True)

    async def progress(
        self,
        job_id: str,
        event_type: str,
        *,
        progress: float | None = None,
        stage: str | None = None,
    ) -> bool:
        async with self._lock:
            record = self._jobs.get(job_id)
            if record is None:
                raise JobNotFound(job_id)
            if record.view.status in TERMINAL:
                return False
            suffix = event_type.rsplit(".", 1)[-1]
            if suffix == "started":
                record.view.status = "started"
            elif suffix == "cancelling":
                record.view.status = "cancelling"
            if progress is not None:
                record.view.progress = max(0.0, min(1.0, progress))
            record.view.updated_at = utc_now()
            payload: dict[str, Any] = {
                "type": event_type,
                "job_id": job_id,
                "kind": record.view.kind,
            }
            if progress is not None:
                payload["progress"] = record.view.progress
            if stage:
                payload["stage"] = stage[:160]
            self._event_locked(record, event_type, payload)
            self._conditions[job_id].notify_all()
            return True

    async def terminal(
        self,
        job_id: str,
        status: str,
        *,
        artifacts: list[dict[str, Any]] | None = None,
        error_code: str | None = None,
        summary: str | None = None,
    ) -> bool:
        if status not in TERMINAL:
            raise ValueError("terminal status required")
        async with self._lock:
            record = self._jobs.get(job_id)
            if record is None:
                raise JobNotFound(job_id)
            if record.view.status in TERMINAL:
                return False
            record.view.status = status  # type: ignore[assignment]
            record.view.updated_at = utc_now()
            record.view.progress = 1.0 if status == "completed" else record.view.progress
            if artifacts is not None:
                from .models import MediaArtifact

                record.view.artifacts = [MediaArtifact.model_validate(item) for item in artifacts]
            if error_code:
                record.view.error = {"code": error_code[:120]}
            payload: dict[str, Any] = {
                "type": f"job.{status}",
                "job_id": job_id,
                "kind": record.view.kind,
                "artifacts": [item.model_dump(mode="json") for item in record.view.artifacts],
            }
            if error_code:
                payload["error"] = {"code": error_code[:120]}
            if summary:
                payload["summary"] = summary[:4000]
            self._event_locked(record, f"job.{status}", payload)
            self._conditions[job_id].notify_all()
            return True

    async def validate_after(self, job_id: str, after: int) -> None:
        async with self._lock:
            record = self._jobs.get(job_id)
            if record is None:
                raise JobNotFound(job_id)
            if record.events and after and after < record.events[0].seq - 1:
                raise ReplayGap("requested media event is no longer retained")

    async def stream(self, job_id: str, after: int, heartbeat_s: float) -> AsyncIterator[StoredEvent | None]:
        async with self._lock:
            if self._subscribers >= self.max_subscribers:
                raise JobCapacityExceeded("media event subscriber capacity reached")
            if job_id not in self._jobs:
                raise JobNotFound(job_id)
            self._subscribers += 1
        cursor = after
        try:
            while True:
                async with self._lock:
                    record = self._jobs.get(job_id)
                    if record is None:
                        raise JobNotFound(job_id)
                    batch = [item.model_copy(deep=True) for item in record.events if item.seq > cursor]
                    terminal = record.view.status in TERMINAL
                    if not batch and not terminal:
                        try:
                            await asyncio.wait_for(
                                self._conditions[job_id].wait(), timeout=heartbeat_s
                            )
                        except TimeoutError:
                            pass
                        batch = [
                            item.model_copy(deep=True)
                            for item in record.events
                            if item.seq > cursor
                        ]
                        terminal = record.view.status in TERMINAL
                if batch:
                    for item in batch:
                        cursor = item.seq
                        yield item
                elif terminal:
                    return
                else:
                    yield None
        finally:
            async with self._lock:
                self._subscribers = max(0, self._subscribers - 1)

    async def cancel_all(self) -> None:
        async with self._lock:
            tasks = [
                record.task
                for record in self._jobs.values()
                if record.task is not None and not record.task.done()
            ]
        for task in tasks:
            task.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)