"""Async jobs — the backbone for heavy work (sources / separate / loops) (#25).

The API enqueues a job (status=pending) and hands back a `job_id`; a single
worker daemon (worker.py) atomically claims pending jobs FIFO + priority, runs
the handler registered for the job's `type`, and writes progress + a terminal
state (done / error). The API never imports engine code — submit stays light,
and the worker is the only place the heavy deps load.

Design choices (see task #25):
- SINGLE serialized worker — one job at a time. erable is one box, 2 GB, no GPU;
  running two demucs separations at once would OOM. Parallelism is a future lane
  split, not now. The atomic claim is still race-safe for N workers regardless.
- Identity is the gateway tenant (X-Tenant), stored as `account` text. Ownership
  (account match) gates poll/cancel; central auth already gated the scope.
- Idempotency: the submit path checks the cache first and, on a hit, creates the
  job already `done` (enqueue(..., result=...)) so a repeat returns instantly.

Handlers register by type:

    @jobs.handler("separate")
    def run_separate(job, progress):
        progress(0.5)
        return {"stems": ...}          # -> stored as the job result

A handler returns a JSON-able dict (inline result) and/or sets `job["result_ref"]`
via the 2-tuple return `(result, ref)`; raising marks the job error (the worker
applies the attempt cap / requeue policy).
"""
from __future__ import annotations

import json
import secrets
import time
from typing import Callable

import db

# job type -> handler(job: dict, progress: Callable[[float], None]) -> dict | (dict, str)
HANDLERS: dict[str, Callable] = {}

TERMINAL = ("done", "error", "cancelled")
MAX_ATTEMPTS = 3


def handler(job_type: str):
    """Register a worker handler for a job type. The worker dispatches on `type`."""
    def deco(fn: Callable) -> Callable:
        HANDLERS[job_type] = fn
        return fn
    return deco


def _new_id() -> str:
    return "j_" + secrets.token_hex(8)


def _row_to_public(row) -> dict:
    """The API-facing view of a job (no internal columns leak meaningfully)."""
    return {
        "job_id": row["id"],
        "type": row["type"],
        "status": row["status"],
        "progress": round(row["progress"], 4),
        "result": json.loads(row["result_json"]) if row["result_json"] else None,
        "result_ref": row["result_ref"],
        "error": row["error"],
        "created": row["created"],
        "started": row["started"],
        "finished": row["finished"],
    }


# ── enqueue ─────────────────────────────────────────────────────────────────────
def enqueue(account: str, job_type: str, params: dict | None = None, *,
            priority: int = 0, webhook_url: str | None = None,
            result: dict | None = None) -> dict:
    """Create a job and return its public view. If `result` is given (the caller
    already had a cache hit), the job is born `done` — the idempotent fast path."""
    jid = _new_id()
    now = time.time()
    done = result is not None
    with db.cursor() as c:
        c.execute(
            "INSERT INTO jobs(id, account, type, params_json, status, priority, "
            "progress, result_json, webhook_url, attempts, created, started, finished) "
            "VALUES(?,?,?,?,?,?,?,?,?,0,?,?,?)",
            (jid, account, job_type, json.dumps(params or {}),
             "done" if done else "pending", priority,
             1.0 if done else 0.0,
             json.dumps(result) if done else None,
             webhook_url, now, now if done else None, now if done else None))
        row = c.execute("SELECT * FROM jobs WHERE id=?", (jid,)).fetchone()
    return _row_to_public(row)


# ── read / ownership ──────────────────────────────────────────────────────────
def get(job_id: str, account: str | None = None) -> dict | None:
    """Public view of a job, or None if missing — or owned by another account
    when `account` is given (callers should treat that as a 404, not a 403, so
    job existence doesn't leak across tenants)."""
    with db.cursor() as c:
        row = c.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()
    if not row:
        return None
    if account is not None and row["account"] != account:
        return None
    return _row_to_public(row)


def cancel(job_id: str, account: str | None = None) -> str | None:
    """Cancel a pending/running job. Returns the new status, or None if missing/
    not owned. A running job is cancelled cooperatively — the handler sees it via
    is_cancelled() between progress steps."""
    with db.cursor() as c:
        row = c.execute("SELECT account, status FROM jobs WHERE id=?", (job_id,)).fetchone()
        if not row or (account is not None and row["account"] != account):
            return None
        if row["status"] in TERMINAL:
            return row["status"]          # already finished — idempotent no-op
        c.execute("UPDATE jobs SET status='cancelled', finished=? WHERE id=?",
                  (time.time(), job_id))
    return "cancelled"


def is_cancelled(job_id: str) -> bool:
    with db.cursor() as c:
        row = c.execute("SELECT status FROM jobs WHERE id=?", (job_id,)).fetchone()
    return bool(row) and row["status"] == "cancelled"


# ── worker-side lifecycle ───────────────────────────────────────────────────────
def claim() -> dict | None:
    """Atomically claim the next pending job → running. Returns the raw row dict,
    or None if the queue is empty.

    Race-safety: SQLite (WAL) serializes writers, so two workers can't both
    transition the same row — the second's `WHERE status='pending'` no longer
    matches. We pick a candidate, then UPDATE guarded by id+status; rowcount 0
    means someone beat us, so we retry until we win one or the queue drains."""
    while True:
        with db.cursor() as c:
            cand = c.execute(
                "SELECT id FROM jobs WHERE status='pending' "
                "ORDER BY priority DESC, created ASC LIMIT 1").fetchone()
            if not cand:
                return None
            now = time.time()
            cur = c.execute(
                "UPDATE jobs SET status='running', started=?, attempts=attempts+1 "
                "WHERE id=? AND status='pending'", (now, cand["id"]))
            if cur.rowcount == 1:
                row = c.execute("SELECT * FROM jobs WHERE id=?", (cand["id"],)).fetchone()
                return dict(row)
        # lost the race for this candidate — loop and try the next pending one


def set_progress(job_id: str, progress: float) -> None:
    with db.cursor() as c:
        c.execute("UPDATE jobs SET progress=? WHERE id=? AND status='running'",
                  (max(0.0, min(1.0, progress)), job_id))


def finish(job_id: str, result: dict | None = None, result_ref: str | None = None) -> None:
    with db.cursor() as c:
        c.execute(
            "UPDATE jobs SET status='done', progress=1.0, result_json=?, result_ref=?, "
            "finished=? WHERE id=? AND status='running'",
            (json.dumps(result) if result is not None else None, result_ref,
             time.time(), job_id))


def fail(job_id: str, error: str, *, requeue: bool) -> str:
    """Mark a job failed. If `requeue` and under the attempt cap, send it back to
    pending for another worker pass; else terminal error. Returns the new status."""
    with db.cursor() as c:
        row = c.execute("SELECT attempts FROM jobs WHERE id=?", (job_id,)).fetchone()
        if requeue and row and row["attempts"] < MAX_ATTEMPTS:
            c.execute("UPDATE jobs SET status='pending', started=NULL, error=? WHERE id=?",
                      (error, job_id))
            return "pending"
        c.execute("UPDATE jobs SET status='error', error=?, finished=? WHERE id=?",
                  (error, time.time(), job_id))
    return "error"


def recover() -> int:
    """Crash recovery, run once at worker boot: any job left 'running' (the worker
    died mid-flight) goes back to 'pending' if under the attempt cap, else 'error'.
    Returns how many were requeued."""
    now = time.time()
    with db.cursor() as c:
        requeued = c.execute(
            "UPDATE jobs SET status='pending', started=NULL "
            "WHERE status='running' AND attempts < ?", (MAX_ATTEMPTS,)).rowcount
        c.execute(
            "UPDATE jobs SET status='error', error='exceeded attempts (crash loop)', "
            "finished=? WHERE status='running' AND attempts >= ?", (now, MAX_ATTEMPTS))
    return requeued
