"""
Background analysis jobs: a fair queue in front of the CPU, real per-step
progress, cancellation, and a small result cache.

POST /api/analyze/jobs starts a job and returns its id; the client polls
GET /api/analyze/jobs/{id}. The Space has two vCPUs and one analysis
already keeps both busy, so jobs take the CPU one at a time (CPU_SLOTS) in
arrival order and everyone else sees their place in the queue. Network
waits (YouTube download, the external FST model) happen outside that slot.

A job belongs to whoever polls it. One that nobody has polled for a while
when its turn comes (the tab was closed) is dropped instead of analysed.
A reloaded page resumes polling the same id, so it isn't dropped.

State lives in process memory: the Space runs one replica, and a job only
has to outlive the request that started it, not a restart.
"""

from __future__ import annotations

import asyncio
import os
import time
import uuid
from collections import OrderedDict
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Awaitable, Callable, Dict, List, Optional

JOB_TTL_SEC = 15 * 60
MAX_JOBS = 200
# Jobs waiting or running at once; beyond this, new ones get 503.
MAX_PENDING = int(os.getenv("AURIS_MAX_PENDING", "24"))
CPU_SLOTS = int(os.getenv("AURIS_CPU_SLOTS", "1"))
# Unpolled this long when its turn comes: the visitor has gone.
ABANDON_SEC = 180.0
CACHE_SIZE = 64
CACHE_TTL_SEC = 6 * 3600

STEP_PENDING = "pending"
STEP_RUNNING = "running"
STEP_DONE = "done"
STEP_FAILED = "failed"
STEP_SKIPPED = "skipped"

FILE_STEPS = ["features", "vocals", "wav2vec2", "clap", "timeline", "fst", "xai", "meta"]
URL_STEPS = ["download", *FILE_STEPS]


class JobCancelled(Exception):
    """Raised inside a pipeline when its job was cancelled or abandoned."""

    def __init__(self, code: str = "cancelled") -> None:
        super().__init__(code)
        self.code = code


class Progress:
    """Ordered step states for one analysis, plus its cancel flag."""

    def __init__(self, steps: List[str]) -> None:
        self._steps: Dict[str, str] = {s: STEP_PENDING for s in steps}
        self._started: Dict[str, float] = {}
        self._took: Dict[str, float] = {}
        # Steps made of many parts (the timeline scan) report how far along they are.
        self._fraction: Dict[str, float] = {}
        self.cancelled = False
        # queued: waiting for the CPU; running: downloading or analysing.
        self.phase = "queued"
        # Returns True when nobody is waiting for the result any more.
        self.gone: Callable[[], Awaitable[bool]] = _never

    def set(self, step: str, state: str) -> None:
        if step not in self._steps:
            return
        now = time.monotonic()
        if state == STEP_RUNNING:
            self._started.setdefault(step, now)
        elif step in self._started and step not in self._took:
            self._took[step] = round(now - self._started[step], 2)
        self._steps[step] = state

    def advance(self, step: str, done: int, total: int) -> None:
        if step in self._steps and total > 0:
            self._fraction[step] = round(min(1.0, done / total), 3)

    def running(self, *steps: str) -> None:
        for s in steps:
            self.set(s, STEP_RUNNING)

    def skip_pending(self) -> None:
        for step, state in self._steps.items():
            if state in (STEP_PENDING, STEP_RUNNING):
                self._steps[step] = STEP_SKIPPED

    def check(self) -> None:
        if self.cancelled:
            raise JobCancelled("cancelled")

    def snapshot(self) -> List[Dict[str, Any]]:
        out: List[Dict[str, Any]] = []
        for step, state in self._steps.items():
            item: Dict[str, Any] = {"id": step, "state": state}
            if step in self._took:
                item["seconds"] = self._took[step]
            if step in self._fraction and state == STEP_RUNNING:
                item["progress"] = self._fraction[step]
            out.append(item)
        return out


async def _never() -> bool:
    return False


class CpuGate:
    """First come, first served access to the CPU-heavy stage."""

    def __init__(self, slots: int) -> None:
        self._slots = max(1, slots)
        self._busy = 0
        self._waiting: List[str] = []
        self._events: Dict[str, asyncio.Event] = {}
        # Mean seconds a job holds the CPU; None until one has finished.
        self._mean: Optional[float] = None

    def position(self, ticket: str) -> Optional[int]:
        """1-based place in the queue, or None when not waiting."""
        try:
            return self._waiting.index(ticket) + 1
        except ValueError:
            return None

    def estimated_wait(self, position: int) -> Optional[float]:
        if self._mean is None:
            return None
        # Everyone ahead, plus whatever is on the CPU now (half done on average).
        return round(((position - 1) / self._slots + (0.5 if self._busy else 0)) * self._mean, 1)

    @property
    def waiting(self) -> int:
        return len(self._waiting)

    @property
    def busy(self) -> int:
        return self._busy

    def _wake(self) -> None:
        while self._waiting and self._busy < self._slots:
            ticket = self._waiting.pop(0)
            self._busy += 1
            self._events.pop(ticket).set()

    @asynccontextmanager
    async def slot(self, ticket: str, progress: Optional[Progress] = None) -> AsyncIterator[None]:
        if progress is not None:
            progress.phase = "queued"
        event = asyncio.Event()
        self._events[ticket] = event
        self._waiting.append(ticket)
        self._wake()
        try:
            await event.wait()
        except BaseException:
            if ticket in self._waiting:
                self._waiting.remove(ticket)
                self._events.pop(ticket, None)
            else:
                self._release()
            raise
        started = time.monotonic()
        try:
            if progress is not None:
                progress.phase = "running"
                progress.check()
                if await progress.gone():
                    raise JobCancelled("job_abandoned")
            yield
        finally:
            held = time.monotonic() - started
            if held > 1.0:
                self._mean = held if self._mean is None else self._mean * 0.7 + held * 0.3
            self._release()

    def _release(self) -> None:
        self._busy = max(0, self._busy - 1)
        self._wake()


cpu_gate = CpuGate(CPU_SLOTS)


class ResultCache:
    """Recent successful responses by content key (same file, same link)."""

    def __init__(self, size: int, ttl: float) -> None:
        self._size = size
        self._ttl = ttl
        self._items: "OrderedDict[str, tuple[float, Dict[str, Any]]]" = OrderedDict()

    def get(self, key: Optional[str]) -> Optional[Dict[str, Any]]:
        if not key or key not in self._items:
            return None
        stored, value = self._items[key]
        if time.time() - stored > self._ttl:
            del self._items[key]
            return None
        self._items.move_to_end(key)
        return value

    def put(self, key: Optional[str], value: Dict[str, Any]) -> None:
        if not key or value.get("errors") or not value.get("result"):
            return
        self._items[key] = (time.time(), value)
        self._items.move_to_end(key)
        while len(self._items) > self._size:
            self._items.popitem(last=False)


result_cache = ResultCache(CACHE_SIZE, CACHE_TTL_SEC)


@dataclass
class Job:
    id: str
    source_type: str
    progress: Progress
    key: Optional[str] = None
    created: float = field(default_factory=time.time)
    updated: float = field(default_factory=time.time)
    last_seen: float = field(default_factory=time.monotonic)
    status: str = "active"          # active | done | error
    cached: bool = False
    response: Optional[Dict[str, Any]] = None
    task: Optional[asyncio.Task] = None

    @property
    def finished(self) -> bool:
        return self.status in ("done", "error")

    def to_dict(self) -> Dict[str, Any]:
        data: Dict[str, Any] = {
            "jobId": self.id,
            "sourceType": self.source_type,
            # "running" until settled, whether waiting or working, so clients
            # that predate the queue keep polling; `phase` tells them apart.
            "status": self.status if self.finished else "running",
            "phase": self.status if self.finished else self.progress.phase,
            "steps": self.progress.snapshot(),
            "elapsedSec": round((self.updated if self.finished else time.time()) - self.created, 2),
            "response": self.response,
        }
        position = cpu_gate.position(self.id)
        if not self.finished and position is not None:
            data["queue"] = {
                "position": position,
                # Everyone in line before you, plus whoever is on the CPU now.
                "ahead": position - 1 + cpu_gate.busy,
                "estimatedWaitSec": cpu_gate.estimated_wait(position),
            }
        if self.cached:
            data["cached"] = True
        return data


class QueueFull(Exception):
    pass


Work = Callable[[Job], Awaitable[Dict[str, Any]]]


class JobStore:
    def __init__(self) -> None:
        self._jobs: Dict[str, Job] = {}
        self._inflight: Dict[str, str] = {}

    def _prune(self) -> None:
        cutoff = time.time() - JOB_TTL_SEC
        for job_id in [j.id for j in self._jobs.values() if j.finished and j.updated < cutoff]:
            del self._jobs[job_id]
        if len(self._jobs) > MAX_JOBS:
            finished = sorted((j for j in self._jobs.values() if j.finished), key=lambda j: j.updated)
            for job in finished[: len(self._jobs) - MAX_JOBS]:
                del self._jobs[job.id]

    def pending_count(self) -> int:
        return sum(1 for j in self._jobs.values() if not j.finished)

    def running_count(self) -> int:
        return sum(1 for j in self._jobs.values() if not j.finished and j.progress.phase == "running")

    def start(self, source_type: str, steps: List[str], work: Work, key: Optional[str] = None) -> Job:
        """Queue a job; the same file or link already in flight is shared, a recent one is served from cache."""
        self._prune()
        if key and key in self._inflight:
            shared = self._jobs.get(self._inflight[key])
            if shared and not shared.finished:
                shared.last_seen = time.monotonic()
                return shared

        job = Job(id=uuid.uuid4().hex, source_type=source_type, progress=Progress(steps), key=key)
        hit = result_cache.get(key)
        if hit is not None:
            job.status, job.cached, job.response = "done", True, hit
            job.progress.skip_pending()
            for step in steps:
                job.progress.set(step, STEP_DONE)
            self._jobs[job.id] = job
            return job

        if self.pending_count() >= MAX_PENDING:
            raise QueueFull()

        async def gone() -> bool:
            return time.monotonic() - job.last_seen > ABANDON_SEC

        job.progress.gone = gone
        self._jobs[job.id] = job
        if key:
            self._inflight[key] = job.id

        async def runner() -> None:
            try:
                job.response = await work(job)
                failed = bool(job.response.get("errors")) and not job.response.get("result")
                job.status = "error" if failed else "done"
                result_cache.put(key, job.response)
            except JobCancelled as stop:
                job.progress.skip_pending()
                job.response = {"result": None, "warnings": [], "errors": [stop.code]}
                job.status = "error"
            except asyncio.CancelledError:
                # Cancelled while still waiting for the CPU.
                job.progress.skip_pending()
                job.response = {"result": None, "warnings": [], "errors": ["cancelled"]}
                job.status = "error"
            except Exception:  # noqa: BLE001 — the job must always settle
                job.progress.skip_pending()
                job.response = {"result": None, "warnings": [], "errors": ["internal_error"]}
                job.status = "error"
            finally:
                job.updated = time.time()
                if key and self._inflight.get(key) == job.id:
                    del self._inflight[key]
                job.task = None

        # Keep a reference: asyncio only holds weak references to tasks.
        job.task = asyncio.create_task(runner())
        return job

    def get(self, job_id: str, touch: bool = True) -> Optional[Job]:
        self._prune()
        job = self._jobs.get(job_id)
        if job is not None and touch:
            job.last_seen = time.monotonic()
        return job

    def cancel(self, job_id: str) -> Optional[Job]:
        job = self._jobs.get(job_id)
        if job is None:
            return None
        if not job.finished:
            job.progress.cancelled = True
            # A job still waiting for the CPU never gets there.
            if cpu_gate.position(job.id) is not None and job.task is not None:
                job.task.cancel()
        return job


job_store = JobStore()