"""Job runner: executes operators as subprocesses, one heavy job at a time.

Operator contract:
  <python> run.py --input <asset path> --outdir <job outdir> --params '<json>'
  - stdout/stderr are streamed into the job log
  - operator writes outputs into outdir and (optionally) outdir/result.json:
      {"outputs": [{"path": "relative/or/absolute", "name": "...", "meta": {...}}, ...],
       "summary": {...}}
  - if result.json is absent, every top-level file in outdir is registered as an asset
"""
import asyncio
import json
import sys
from pathlib import Path

from . import db, store

JOBS_DIR = db.DATA / "jobs"


class Runner:
    def __init__(self):
        self.queue: asyncio.Queue[str] = asyncio.Queue()
        self.subscribers: set = set()  # websockets
        self.operators: dict[str, dict] = {}

    # -- pubsub -------------------------------------------------------------
    async def broadcast(self, message: dict):
        dead = []
        for ws in self.subscribers:
            try:
                await ws.send_json(message)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.subscribers.discard(ws)

    # -- job lifecycle --------------------------------------------------------
    def create_job(self, con, operator: str, asset_id: str | None, params: dict) -> dict:
        job_id = db.new_id()
        outdir = JOBS_DIR / job_id
        outdir.mkdir(parents=True, exist_ok=True)
        con.execute(
            "INSERT INTO jobs (id, operator, status, asset_id, params, outdir, created_at) "
            "VALUES (?, ?, 'queued', ?, ?, ?, ?)",
            (job_id, operator, asset_id, json.dumps(params), str(outdir), db.now()),
        )
        con.commit()
        self.queue.put_nowait(job_id)
        return self.get_job(con, job_id)

    def get_job(self, con, job_id: str) -> dict | None:
        row = con.execute("SELECT * FROM jobs WHERE id = ?", (job_id,)).fetchone()
        return db.row_to_dict(row) if row else None

    def list_jobs(self, con) -> list[dict]:
        rows = con.execute("SELECT * FROM jobs ORDER BY created_at DESC LIMIT 200").fetchall()
        return [db.row_to_dict(r) for r in rows]

    async def _update(self, con, job_id: str, **fields):
        sets = ", ".join(f"{k} = ?" for k in fields)
        con.execute(f"UPDATE jobs SET {sets} WHERE id = ?", (*fields.values(), job_id))
        con.commit()
        job = self.get_job(con, job_id)
        await self.broadcast({"type": "job", "job": job})

    # -- worker -----------------------------------------------------------------
    async def worker(self):
        con = db.connect()
        # re-queue jobs that were left running/queued by a previous process
        for row in con.execute("SELECT id FROM jobs WHERE status IN ('queued','running')"):
            con.execute("UPDATE jobs SET status='queued' WHERE id = ?", (row["id"],))
            self.queue.put_nowait(row["id"])
        con.commit()
        while True:
            job_id = await self.queue.get()
            try:
                await self._run_job(con, job_id)
            except Exception as e:
                await self._update(con, job_id, status="error", error=str(e), finished_at=db.now())

    async def _run_job(self, con, job_id: str):
        job = self.get_job(con, job_id)
        if not job:
            return
        op = self.operators.get(job["operator"])
        if not op:
            await self._update(con, job_id, status="error", error=f"unknown operator {job['operator']}",
                               finished_at=db.now())
            return
        asset = store.get_asset(con, job["asset_id"]) if job["asset_id"] else None
        outdir = Path(job["outdir"])
        entry = Path(op["dir"]) / op.get("entry", "run.py")
        python = op.get("python") or sys.executable
        cmd = [python, str(entry), "--outdir", str(outdir),
               "--params", json.dumps(job["params"])]
        if asset:
            cmd += ["--input", asset["path"]]

        await self._update(con, job_id, status="running", started_at=db.now())
        proc = await asyncio.create_subprocess_exec(
            *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT)
        log_lines: list[str] = []
        assert proc.stdout
        async for raw in proc.stdout:
            line = raw.decode(errors="replace")
            log_lines.append(line)
            if len(log_lines) % 5 == 0:  # don't hammer the DB on chatty tools
                await self._update(con, job_id, log="".join(log_lines)[-100_000:])
        code = await proc.wait()
        log = "".join(log_lines)[-100_000:]

        if code != 0:
            await self._update(con, job_id, status="error", log=log,
                               error=f"exit code {code}", finished_at=db.now())
            return

        # register outputs
        result_path = outdir / "result.json"
        registered = []
        if result_path.exists():
            result = json.loads(result_path.read_text())
            for out in result.get("outputs", []):
                p = Path(out["path"])
                if not p.is_absolute():
                    p = outdir / p
                if p.exists():
                    a = store.register_file(con, p, name=out.get("name"),
                                            parent_job=job_id, move=True,
                                            meta=out.get("meta"))
                    registered.append(a["id"])
        else:
            for p in sorted(outdir.iterdir()):
                if p.name == "result.json" or p.name.startswith("."):
                    continue
                a = store.register_file(con, p, parent_job=job_id, move=True)
                registered.append(a["id"])

        await self._update(con, job_id, status="done", log=log, finished_at=db.now())
        await self.broadcast({"type": "assets_changed", "job_id": job_id, "assets": registered})


runner = Runner()