"""Job runner: executes operators as subprocesses across concurrency lanes.

Operator contract:
  <python> run.py --input <path> [--input <path> ...] --outdir <dir> --params '<json>'
  - inputs are passed in asset order (single-input ops get one --input)
  - stdout/stderr are streamed into the job log (secret values redacted)
  - operator writes outputs into outdir and (optionally) outdir/result.json:
      {"outputs": [{"path": "rel/or/abs", "name": "...", "meta": {...}}, ...],
       "summary": {...}}
  - if result.json is absent, every top-level file in outdir is registered
Concurrency lanes (from manifest "resources"): gpu and cpu are a node pool
(distributed across the primary + remote workers — gpu 1/node, cpu N/node);
net=6 stays primary-only. Env from settings is injected; "requires_env" gates.
"""
import asyncio
import json
import os
import signal
import sys
from pathlib import Path

from . import db, remote, settings as settings_mod, store

JOBS_DIR = db.DATA / "jobs"
# gpu AND cpu are a NODE POOL (this Mac + remote workers): gpu runs 1 job/node
# (one Metal device each), cpu runs per-node slots — LANE_LIMITS["cpu"] on the
# primary, fewer on helpers (nodes.json "cpu_slots" overrides). net stays a
# primary-only semaphore because the cloud API keys live only on the primary.
LANE_LIMITS = {"cpu": 3, "net": 6}


class Runner:
    def __init__(self):
        self.queue: asyncio.Queue[str] = asyncio.Queue()
        self.subscribers: dict = {}    # websocket -> user dict (for per-user scoping)
        self.operators: dict[str, dict] = {}
        self.lanes: dict[str, asyncio.Semaphore] = {}
        self.gpu_nodes: list[dict] = []   # [{local}, {ssh,root,operators}, ...]
        self.procs: dict = {}          # job_id -> subprocess
        self.cancelled: set[str] = set()
        self._tasks: set = set()
        self.get_settings = lambda: {}  # set by main.py

    # -- pubsub -------------------------------------------------------------
    async def broadcast(self, message: dict):
        # "job" events carry a payload injected straight into the client's list, so
        # scope them: a guest only receives their own jobs (owner receives all).
        # "assets_changed" just triggers a per-user re-fetch, so send to everyone.
        job_uid = message.get("job", {}).get("user_id") if message.get("type") == "job" else None
        dead = []
        for ws, u in list(self.subscribers.items()):
            if job_uid is not None and u.get("role") != "owner" and job_uid != u.get("id"):
                continue
            try:
                await ws.send_json(message)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.subscribers.pop(ws, None)

    def lane_of(self, operator: str) -> str:
        op = self.operators.get(operator, {})
        lane = op.get("resources", "cpu")
        # gpu is a valid lane (the node pool) even though it isn't in LANE_LIMITS,
        # which holds only the semaphore lanes (cpu/net).
        return lane if lane in LANE_LIMITS or lane == "gpu" else "cpu"

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

    def active_job_count(self, con, user_id: str) -> int:
        row = con.execute(
            "SELECT COUNT(*) c FROM jobs WHERE user_id = ? AND status IN ('queued','running')",
            (user_id,)).fetchone()
        return row["c"]

    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]

    def input_paths(self, con, job: dict) -> list[str]:
        ids = job.get("asset_ids") or ([job["asset_id"]] if job.get("asset_id") else [])
        paths = []
        for aid in ids:
            asset = store.get_asset(con, aid)
            if asset:
                paths.append(asset["path"])
        return paths

    async def cancel(self, con, job_id: str) -> bool:
        job = self.get_job(con, job_id)
        if not job or job["status"] in ("done", "error", "cancelled"):
            return False
        self.cancelled.add(job_id)
        proc = self.procs.get(job_id)
        if proc and proc.returncode is None:
            try:
                os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
            except (ProcessLookupError, PermissionError):
                pass
        else:  # queued but not yet started
            await self._update(con, job_id, status="cancelled", finished_at=db.now())
        return True

    def delete_job(self, con, job_id: str) -> bool:
        job = self.get_job(con, job_id)
        if not job or job_id in self.procs:
            return False
        import shutil
        if job.get("outdir"):
            shutil.rmtree(job["outdir"], ignore_errors=True)
        con.execute("DELETE FROM jobs WHERE id = ?", (job_id,))
        con.commit()
        return True

    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()
        for key in LANE_LIMITS:
            self.lanes[key] = asyncio.Semaphore(LANE_LIMITS[key])
        self.gpu_nodes = remote.load_gpu_nodes()
        for n in self.gpu_nodes:
            n.setdefault("inflight", {"gpu": 0, "cpu": 0})
        names = ", ".join(n.get("name", n.get("ssh", "?")) for n in self.gpu_nodes)
        print(f"[runner] node pool (gpu 1/node + cpu distributed): {names}", flush=True)
        # re-queue jobs 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()
            task = asyncio.create_task(self._run_lane(con, job_id))
            self._tasks.add(task)
            task.add_done_callback(self._tasks.discard)

    async def _run_lane(self, con, job_id: str):
        job = self.get_job(con, job_id)
        if not job:  # deleted while queued
            self.cancelled.discard(job_id)
            return
        if job_id in self.cancelled:
            self.cancelled.discard(job_id)
            await self._update(con, job_id, status="cancelled", finished_at=db.now())
            return
        lane = self.lane_of(job["operator"])
        if lane in ("gpu", "cpu"):  # pooled lanes — distribute across the node pool
            node = await self._acquire_node(job["operator"], lane)
            try:
                if job_id in self.cancelled:
                    self.cancelled.discard(job_id)
                    await self._update(con, job_id, status="cancelled", finished_at=db.now())
                    return
                await self._run_job(con, job_id, node=node)
            except Exception as e:
                await self._update(con, job_id, status="error", error=str(e), finished_at=db.now())
            finally:
                self._release_node(node, lane)
            return
        async with self.lanes[lane]:  # net: cloud keys live on the primary — not pooled
            if job_id in self.cancelled:
                self.cancelled.discard(job_id)
                await self._update(con, job_id, status="cancelled", finished_at=db.now())
                return
            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())

    def _node_capacity(self, node: dict, lane: str) -> int:
        """Max concurrent jobs of this lane a node may run. gpu=1 (one Metal
        device per Mac); cpu=per-node slots (small helper nodes get fewer),
        overridable per node in nodes.json via "cpu_slots"; net is not pooled
        (cloud keys live only on the primary)."""
        if lane == "gpu":
            return 1
        if lane == "cpu":
            default = LANE_LIMITS["cpu"] if node.get("local") else 2
            return int(node.get("cpu_slots", default))
        return 0

    async def _acquire_node(self, op_id: str, lane: str) -> dict:
        """Return the first node with free capacity in this lane that supports the
        operator (local always does; a remote must list it + be reachable). Polls."""
        while True:
            for n in self.gpu_nodes:
                if not remote.node_supports(n, op_id):
                    continue
                cap = self._node_capacity(n, lane)
                if cap <= 0 or n["inflight"].get(lane, 0) >= cap:
                    continue
                n["inflight"][lane] = n["inflight"].get(lane, 0) + 1  # reserve before the await
                if not n.get("local") and not await remote.healthy(n):
                    n["inflight"][lane] -= 1  # unreachable right now; release and keep looking
                    continue
                return n
            await asyncio.sleep(0.25)

    def _release_node(self, node: dict, lane: str) -> None:
        if node is not None and "inflight" in node:
            node["inflight"][lane] = max(0, node["inflight"].get(lane, 0) - 1)

    async def _run_job(self, con, job_id: str, node: dict | None = None):
        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

        cur_settings = self.get_settings()
        env = os.environ.copy()
        # least privilege: only this operator's declared paid keys (+ HF/infra) reach
        # the subprocess, so guest local jobs never get the owner's fal/OpenRouter keys
        env.update(settings_mod.env_for_operator(cur_settings, op.get("requires_env")))
        secrets = settings_mod.secret_values(cur_settings)

        missing = [e for e in op.get("requires_env", []) if not env.get(e)]
        if missing:
            keys = ", ".join(sorted({k for k, v in settings_mod.ENV_MAP.items() if v in missing}))
            await self._update(con, job_id, status="error",
                               error=f"missing setting(s): {keys} — add them in Settings",
                               finished_at=db.now())
            return

        outdir = Path(job["outdir"])
        params_json = json.dumps(job["params"])
        local_inputs = self.input_paths(con, job)
        entry = Path(op["dir"]) / op.get("entry", "run.py")
        # manifest "python" may be repo-relative (portable) or absolute (legacy)
        op_py_rel = op.get("python")
        op_py_abs = (str(db.ROOT / op_py_rel)
                     if op_py_rel and not Path(op_py_rel).is_absolute() else op_py_rel)

        node = node or {"local": True, "name": "local"}
        if node.get("local"):
            python = op_py_abs or sys.executable
            if op_py_abs and not Path(python).exists():
                await self._update(con, job_id, status="error", finished_at=db.now(),
                                   error=f"operator not installed — {python} missing. "
                                         f"Run its install script under scripts/.")
                return
            cmd = [python, str(entry), "--outdir", str(outdir), "--params", params_json]
            for p in local_inputs:
                cmd += ["--input", p]
            run_env = env
        else:  # remote node: rsync inputs out, run over ssh, rsync outputs back
            if op_py_rel and Path(op_py_rel).is_absolute():
                await self._update(con, job_id, status="error", finished_at=db.now(),
                                   error="operator has an absolute python path — can't run remotely")
                return
            # python-less ops (ffmpeg/ffprobe wrappers) run on the node's system
            # python3, resolved via the remote login shell (bash -lc in remote_cmd).
            remote_python = op_py_rel or "python3"
            try:
                remote_inputs = await remote.prepare(node, job_id, local_inputs)
            except Exception as e:
                await self._update(con, job_id, status="error", finished_at=db.now(),
                                   error=f"remote prepare failed: {e}")
                return
            entry_rel = os.path.relpath(str(entry), str(db.ROOT))
            cmd = remote.remote_cmd(node, op["id"], remote_python, entry_rel,
                                    job_id, params_json, remote_inputs)
            run_env = os.environ.copy()

        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,
            start_new_session=True, env=run_env)
        self.procs[job_id] = proc
        log_lines: list[str] = [f"[node: {node.get('name', 'local')}]\n"]
        assert proc.stdout
        try:
            async for raw in proc.stdout:
                log_lines.append(raw.decode(errors="replace"))
                if len(log_lines) % 5 == 0:
                    log = settings_mod.redact("".join(log_lines)[-100_000:], secrets)
                    await self._update(con, job_id, log=log)
            code = await proc.wait()
        finally:
            self.procs.pop(job_id, None)
        log = settings_mod.redact("".join(log_lines)[-100_000:], secrets)

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

        if not node.get("local"):  # pull the remote outputs into the local job dir
            try:
                await remote.collect(node, job_id, outdir)
            except Exception as e:
                await self._update(con, job_id, status="error", log=log,
                                   error=f"remote collect failed: {e}", finished_at=db.now())
                return
        registered = self._register_outputs(con, outdir, job_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})

    def _register_outputs(self, con, outdir: Path, job_id: str) -> list[str]:
        result_path = outdir / "result.json"
        registered = []
        # outputs belong to whoever ran the job — guests see only their own
        # assets, so registering without user_id hides a guest's results from them
        row = con.execute("SELECT user_id FROM jobs WHERE id = ?", (job_id,)).fetchone()
        owner_id = row["user_id"] if row else None
        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"),
                                            user_id=owner_id)
                    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, user_id=owner_id)
                registered.append(a["id"])
        return registered


runner = Runner()