"""Small asyncio job system (downloads, scans, benchmarks, harvests)."""

from __future__ import annotations

import asyncio
import json
import logging
import time
import uuid
from typing import Any, Awaitable, Callable

from .db import Database
from .events import bus

log = logging.getLogger("llm_api.jobs")

QUEUED, RUNNING, COMPLETED, FAILED, CANCELLED = "queued", "running", "completed", "failed", "cancelled"


class Job:
    def __init__(self, id: str, kind: str, title: str, payload: dict):
        self.id = id
        self.kind = kind
        self.title = title
        self.payload = payload
        self.status = QUEUED
        self.progress = 0.0
        self.detail: dict[str, Any] = {}
        self.result: Any = None
        self.error: str | None = None
        self.created_at = time.time()
        self.started_at: float | None = None
        self.finished_at: float | None = None
        self.cancel_event = asyncio.Event()
        self.task: asyncio.Task | None = None
        self._last_flush = 0.0

    def to_dict(self) -> dict:
        return {
            "id": self.id, "kind": self.kind, "title": self.title, "payload": self.payload, "status": self.status,
            "progress": round(self.progress, 4), "detail": self.detail, "result": self.result, "error": self.error,
            "created_at": self.created_at, "started_at": self.started_at, "finished_at": self.finished_at,
        }

    @property
    def cancelled(self) -> bool:
        return self.cancel_event.is_set()


class JobRunner:
    def __init__(self, db: Database, max_concurrent: int = 3):
        self.db = db
        self.jobs: dict[str, Job] = {}
        self.sem = asyncio.Semaphore(max_concurrent)
        self._download_lock = asyncio.Semaphore(1)  # downloads are serialized (SSD + bandwidth)

    async def start(self) -> None:
        # Mark jobs left running by a crashed server as failed
        await self.db.execute("UPDATE jobs SET status=?, error=?, finished_at=? WHERE status IN (?, ?)",
                              (FAILED, "server restarted", time.time(), QUEUED, RUNNING))

    def submit(self, kind: str, title: str, payload: dict,
               fn: Callable[[Job], Awaitable[Any]], *, exclusive_download: bool = False) -> Job:
        job = Job(uuid.uuid4().hex[:12], kind, title, payload)
        self.jobs[job.id] = job
        asyncio.create_task(self._persist(job))

        async def _run():
            async with self.sem:
                if exclusive_download:
                    async with self._download_lock:
                        await self._execute(job, fn)
                else:
                    await self._execute(job, fn)

        job.task = asyncio.create_task(_run(), name=f"job-{kind}-{job.id}")
        bus.publish("job", job.to_dict())
        return job

    async def _execute(self, job: Job, fn) -> None:
        if job.cancelled:
            job.status = CANCELLED
            job.finished_at = time.time()
            await self._persist(job)
            return
        job.status = RUNNING
        job.started_at = time.time()
        await self._persist(job)
        bus.publish("job", job.to_dict())
        try:
            job.result = await fn(job)
            job.status = CANCELLED if job.cancelled else COMPLETED
            job.progress = 1.0 if job.status == COMPLETED else job.progress
        except asyncio.CancelledError:
            job.status = CANCELLED
        except Exception as e:
            log.exception("job %s failed", job.id)
            job.status = FAILED
            job.error = f"{type(e).__name__}: {e}"
        job.finished_at = time.time()
        await self._persist(job)
        bus.publish("job", job.to_dict())

    def update(self, job: Job, progress: float | None = None, **detail: Any) -> None:
        if progress is not None:
            job.progress = max(0.0, min(1.0, progress))
        if detail:
            job.detail.update(detail)
        now = time.time()
        if now - job._last_flush > 0.5:
            job._last_flush = now
            bus.publish("job", job.to_dict())
            asyncio.create_task(self._persist(job))

    async def _persist(self, job: Job) -> None:
        try:
            await self.db.execute(
                "INSERT INTO jobs(id, kind, status, title, payload, progress, detail, result, error, created_at, started_at, finished_at) "
                "VALUES(?,?,?,?,?,?,?,?,?,?,?,?) ON CONFLICT(id) DO UPDATE SET status=excluded.status, progress=excluded.progress, "
                "detail=excluded.detail, result=excluded.result, error=excluded.error, started_at=excluded.started_at, finished_at=excluded.finished_at",
                (job.id, job.kind, job.status, job.title, json.dumps(job.payload), job.progress, json.dumps(job.detail),
                 json.dumps(job.result, default=str) if job.result is not None else None, job.error, job.created_at,
                 job.started_at, job.finished_at))
        except Exception:
            log.exception("persist job failed")

    async def cancel(self, job_id: str) -> bool:
        job = self.jobs.get(job_id)
        if not job:
            return False
        if job.status in (COMPLETED, FAILED, CANCELLED):
            return False
        job.cancel_event.set()
        if job.status == QUEUED and job.task:
            job.task.cancel()
        return True

    def list(self, kinds: set[str] | None = None, limit: int = 100) -> list[dict]:
        items = sorted(self.jobs.values(), key=lambda j: -j.created_at)
        if kinds:
            items = [j for j in items if j.kind in kinds]
        return [j.to_dict() for j in items[:limit]]

    async def history(self, kinds: set[str] | None = None, limit: int = 100) -> list[dict]:
        rows = await self.db.fetchall("SELECT * FROM jobs ORDER BY created_at DESC LIMIT ?", (limit,))
        out = []
        for r in rows:
            if kinds and r["kind"] not in kinds:
                continue
            for k in ("payload", "detail", "result"):
                if r.get(k):
                    try:
                        r[k] = json.loads(r[k])
                    except Exception:
                        pass
            # prefer live state
            live = self.jobs.get(r["id"])
            out.append(live.to_dict() if live else r)
        return out

    def get(self, job_id: str) -> Job | None:
        return self.jobs.get(job_id)