"""E72 durable long-horizon mission controller.

State lives in SQLite (WAL). The controller can be killed and restarted: RUNNING tasks are returned to READY,
completed tasks and checkpoints survive, and the mission continues against the same deadline and budget.

This module deliberately separates three numbers that are easy to conflate:
  * mission elapsed time      wall-clock since creation (can be hours)
  * task execution time       only the intervals a task was actually being executed/verified
  * idle time                 waiting for a scheduled sample or for a dependency

A one-hour elapsed mission is not one productive worker-hour.
"""
from __future__ import annotations

import hashlib
import json
import os
import platform
import sqlite3
import threading
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence

try:
    import psutil
except Exception:  # noqa: BLE001
    psutil = None  # type: ignore

SCHEMA = "cain72.longrun.v1"


def _iso(ts: Optional[float] = None) -> str:
    return datetime.fromtimestamp(time.time() if ts is None else ts, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


class LongRunStore:
    def __init__(self, path: str):
        self.path = str(path)
        Path(self.path).parent.mkdir(parents=True, exist_ok=True)
        self.lock = threading.RLock()
        self.db = sqlite3.connect(self.path, check_same_thread=False)
        self.db.row_factory = sqlite3.Row
        self.db.execute("PRAGMA journal_mode=WAL")
        self.db.execute("PRAGMA synchronous=FULL")
        self._init()

    def _init(self) -> None:
        with self.lock:
            self.db.executescript("""
            CREATE TABLE IF NOT EXISTS missions(
              mission_id TEXT PRIMARY KEY, objective TEXT, state TEXT, spec TEXT, budget TEXT,
              created_at REAL, deadline REAL, updated_at REAL, last_progress REAL, finished_at REAL,
              stall_s REAL, reason TEXT);
            CREATE TABLE IF NOT EXISTS tasks(
              task_id TEXT PRIMARY KEY, mission_id TEXT, seq INTEGER, kind TEXT, operation TEXT,
              params TEXT, depends TEXT, state TEXT, attempt INTEGER, max_attempts INTEGER,
              due_at REAL, created_at REAL, started_at REAL, finished_at REAL, result TEXT,
              result_digest TEXT, verifier_result TEXT, verifier_digest TEXT, error TEXT);
            CREATE TABLE IF NOT EXISTS checkpoints(
              seq INTEGER PRIMARY KEY AUTOINCREMENT, mission_id TEXT, at REAL, state_digest TEXT,
              prev_hash TEXT, hash TEXT, body TEXT);
            CREATE TABLE IF NOT EXISTS events(
              seq INTEGER PRIMARY KEY AUTOINCREMENT, at REAL, mission_id TEXT, type TEXT, body TEXT);
            """)
            self.db.commit()

    def q(self, sql: str, args: Sequence[Any] = ()) -> List[Dict[str, Any]]:
        with self.lock:
            cur = self.db.execute(sql, tuple(args))
            return [dict(r) for r in cur.fetchall()]

    def one(self, sql: str, args: Sequence[Any] = ()) -> Optional[Dict[str, Any]]:
        rows = self.q(sql, args)
        return rows[0] if rows else None

    def x(self, sql: str, args: Sequence[Any] = ()) -> None:
        with self.lock:
            self.db.execute(sql, tuple(args))
            self.db.commit()

    def event(self, mission_id: str, type_: str, body: Mapping[str, Any]) -> None:
        self.x("INSERT INTO events(at,mission_id,type,body) VALUES(?,?,?,?)",
               (time.time(), mission_id, type_, json.dumps(dict(body), sort_keys=True)))

    def checkpoint(self, mission_id: str, body: Mapping[str, Any]) -> Dict[str, Any]:
        with self.lock:
            prev = self.one("SELECT hash FROM checkpoints WHERE mission_id=? ORDER BY seq DESC LIMIT 1", (mission_id,))
            prev_hash = prev["hash"] if prev else ""
            payload = json.dumps(dict(body), sort_keys=True)
            h = hashlib.sha256((prev_hash + payload).encode()).hexdigest()
            self.db.execute("INSERT INTO checkpoints(mission_id,at,state_digest,prev_hash,hash,body)"
                            " VALUES(?,?,?,?,?,?)",
                            (mission_id, time.time(), hashlib.sha256(payload.encode()).hexdigest(),
                             prev_hash, h, payload))
            self.db.commit()
            return {"hash": h, "prev_hash": prev_hash}

    def checkpoints(self, mission_id: str) -> List[Dict[str, Any]]:
        return self.q("SELECT * FROM checkpoints WHERE mission_id=? ORDER BY seq", (mission_id,))

    def verify_checkpoints(self, mission_id: str) -> Dict[str, Any]:
        prev = ""
        bad = []
        for r in self.checkpoints(mission_id):
            h = hashlib.sha256((prev + r["body"]).encode()).hexdigest()
            if h != r["hash"] or r["prev_hash"] != prev:
                bad.append(r["seq"])
            prev = r["hash"]
        return {"ok": not bad, "checkpoints": len(self.checkpoints(mission_id)), "bad": bad, "head": prev}


def _host_sample() -> Dict[str, Any]:
    data: Dict[str, Any] = {"at": time.time(), "hostname": platform.node(),
                            "boot_id": Path("/proc/sys/kernel/random/boot_id").read_text().strip(),
                            "pid": os.getpid()}
    if psutil:
        data.update({"cpu_percent": psutil.cpu_percent(interval=0.05),
                     "load1": os.getloadavg()[0],
                     "mem_available_mb": round(psutil.virtual_memory().available / (1024 * 1024), 2),
                     "disk_free_gb": round(psutil.disk_usage("/").free / (1024 ** 3), 3),
                     "processes": len(psutil.pids())})
    else:
        data["load1"] = os.getloadavg()[0]
    return data


def _repo_audit(params: Mapping[str, Any]) -> Dict[str, Any]:
    root = Path(str(params.get("root") or "/root/cain-studio-platform"))
    paths = [str(p) for p in (params.get("paths") or [])]
    rows = []
    for rel in paths:
        p = root / rel
        try:
            raw = p.read_bytes()
            text = raw.decode("utf-8", errors="ignore")
            rows.append({"path": rel, "bytes": len(raw), "sha256": hashlib.sha256(raw).hexdigest(),
                         "lines": text.count("\n") + 1, "imports": text.count("import "),
                         "todo": text.upper().count("TODO"), "parses": _py_parses(text) if rel.endswith(".py") else None})
        except Exception as e:  # noqa: BLE001
            rows.append({"path": rel, "error": f"{type(e).__name__}: {e}"})
    out = {"files": len(rows), "total_bytes": sum(int(r.get("bytes") or 0) for r in rows), "rows": rows}
    out["digest"] = hashlib.sha256(json.dumps(rows, sort_keys=True).encode()).hexdigest()
    return out


def _py_parses(text: str) -> bool:
    import ast
    try:
        ast.parse(text)
        return True
    except Exception:  # noqa: BLE001
        return False


class LongRunController:
    def __init__(self, path: str):
        self.store = LongRunStore(path)
        self.stop = threading.Event()

    # ------------------------------------------------------------------ mission
    def create(self, *, mission_id: str, objective: str, tasks: Sequence[Mapping[str, Any]],
               budget: Optional[Mapping[str, Any]] = None, deadline_s: float = 3600,
               stall_s: float = 300) -> Dict[str, Any]:
        now = time.time()
        spec = {"objective": objective, "tasks": list(tasks), "budget": dict(budget or {})}
        self.store.x("INSERT INTO missions(mission_id,objective,state,spec,budget,created_at,deadline,updated_at,"
                     "last_progress,stall_s) VALUES(?,?,?,?,?,?,?,?,?,?)",
                     (mission_id, objective, "RUNNING", json.dumps(spec, sort_keys=True),
                      json.dumps(dict(budget or {}), sort_keys=True), now, now + float(deadline_s), now, now,
                      float(stall_s)))
        for i, t in enumerate(tasks):
            tid = str(t.get("task_id") or f"t-{i:05d}")
            due = float(t.get("due_at") or 0)
            self.store.x("INSERT INTO tasks(task_id,mission_id,seq,kind,operation,params,depends,state,attempt,"
                         "max_attempts,due_at,created_at) VALUES(?,?,?,?,?,?,?,?,?,?,?,?)",
                         (tid, mission_id, i, str(t.get("kind") or "EXECUTE"), str(t.get("operation")),
                          json.dumps(dict(t.get("params") or {}), sort_keys=True),
                          json.dumps(list(t.get("depends") or [])), "READY", 0,
                          int(t.get("max_attempts") or 2), due, now))
        self.store.event(mission_id, "MISSION_CREATED", {"tasks": len(tasks), "deadline_s": deadline_s})
        return {"mission_id": mission_id, "tasks": len(tasks)}

    def mission(self, mission_id: str) -> Optional[Dict[str, Any]]:
        m = self.store.one("SELECT * FROM missions WHERE mission_id=?", (mission_id,))
        if not m:
            return None
        counts = self.store.q("SELECT state,COUNT(*) n FROM tasks WHERE mission_id=? GROUP BY state", (mission_id,))
        m["task_states"] = {r["state"]: int(r["n"]) for r in counts}
        m["tasks"] = sum(m["task_states"].values())
        m["verified"] = m["task_states"].get("VERIFIED", 0)
        m["elapsed_s"] = round(time.time() - float(m["created_at"]), 2)
        return m

    def resume(self, mission_id: str) -> Dict[str, Any]:
        """After a controller crash: any task left RUNNING is returned to READY. Completed work is untouched."""
        n = self.store.one("SELECT COUNT(*) n FROM tasks WHERE mission_id=? AND state='RUNNING'", (mission_id,))["n"]
        self.store.x("UPDATE tasks SET state='READY', error='RECOVERED_AFTER_CONTROLLER_RESTART'"
                     " WHERE mission_id=? AND state='RUNNING'", (mission_id,))
        self.store.x("UPDATE missions SET state='RUNNING', updated_at=?, last_progress=? WHERE mission_id=?"
                     " AND state='STALLED'", (time.time(), time.time(), mission_id))
        self.store.event(mission_id, "CONTROLLER_RESUME", {"recovered_running_tasks": int(n)})
        return {"mission_id": mission_id, "recovered": int(n)}

    # ------------------------------------------------------------------ execution
    def _promote(self, mission_id: str) -> int:
        pending = self.store.q("SELECT * FROM tasks WHERE mission_id=? AND state='BLOCKED_ON_DEPS'", (mission_id,))
        n = 0
        for t in pending:
            deps = json.loads(t["depends"] or "[]")
            states = [self.store.one("SELECT state FROM tasks WHERE task_id=?", (d,)) for d in deps]
            if any(x is None or x["state"] == "FAILED" for x in states):
                self.store.x("UPDATE tasks SET state='FAILED', error='DEPENDENCY_FAILED' WHERE task_id=?", (t["task_id"],))
                continue
            if all(x["state"] == "VERIFIED" for x in states):
                self.store.x("UPDATE tasks SET state='READY' WHERE task_id=?", (t["task_id"],))
                n += 1
        return n

    def tick(self, mission_id: str, *, executors: Mapping[str, Callable], verifiers: Mapping[str, Callable]) -> int:
        done = 0
        now = time.time()
        m = self.store.one("SELECT * FROM missions WHERE mission_id=?", (mission_id,))
        if not m or m["state"] != "RUNNING":
            return 0
        self._promote(mission_id)
        task = self.store.one("SELECT * FROM tasks WHERE mission_id=? AND state='READY' AND due_at<=?"
                              " ORDER BY seq LIMIT 1", (mission_id, now))
        if not task:
            # mark tasks waiting on dependencies
            for t in self.store.q("SELECT * FROM tasks WHERE mission_id=? AND state='READY'", (mission_id,)):
                deps = json.loads(t["depends"] or "[]")
                if deps:
                    self.store.x("UPDATE tasks SET state='BLOCKED_ON_DEPS' WHERE task_id=?", (t["task_id"],))
            return 0
        self.store.x("UPDATE tasks SET state='RUNNING', attempt=attempt+1, started_at=?, error=NULL"
                     " WHERE task_id=?", (now, task["task_id"]))
        try:
            fn = executors[task["operation"]]
            params = json.loads(task["params"] or "{}")
            t0 = time.time()
            result = fn(params)
            t1 = time.time()
            digest = hashlib.sha256(json.dumps(result, sort_keys=True).encode()).hexdigest()
            self.store.x("UPDATE tasks SET state='VERIFY_READY', finished_at=?, result=?, result_digest=?"
                         " WHERE task_id=?", (t1, json.dumps(result, sort_keys=True), digest, task["task_id"]))
            self.store.event(mission_id, "TASK_EXECUTED", {"task_id": task["task_id"], "seconds": round(t1 - t0, 4),
                                                           "digest": digest})
        except Exception as e:  # noqa: BLE001
            self.store.x("UPDATE tasks SET state=CASE WHEN attempt>=max_attempts THEN 'FAILED' ELSE 'READY' END,"
                         " error=? WHERE task_id=?", (f"{type(e).__name__}: {e}", task["task_id"]))
            self.store.event(mission_id, "TASK_ERROR", {"task_id": task["task_id"], "error": str(e)})
            return 1
        # verify immediately with a different function instance (the verifier recomputes from params)
        vtask = self.store.one("SELECT * FROM tasks WHERE task_id=?", (task["task_id"],))
        try:
            vfn = verifiers[vtask["operation"]]
            params = json.loads(vtask["params"] or "{}")
            t0 = time.time()
            vres = vfn(params)
            t1 = time.time()
            vdigest = hashlib.sha256(json.dumps(vres, sort_keys=True).encode()).hexdigest()
            structural = bool(vres.get("__structural__")) and bool(vres.get("ok"))
            if vdigest == vtask["result_digest"] or structural:
                self.store.x("UPDATE tasks SET state='VERIFIED', verifier_result=?, verifier_digest=?"
                             " WHERE task_id=?", (json.dumps(vres, sort_keys=True), vdigest, task["task_id"]))
                self.store.event(mission_id, "TASK_VERIFIED", {"task_id": task["task_id"],
                                                               "seconds": round(t1 - t0, 4), "digest": vdigest,
                                                               "method": "STRUCTURAL" if structural else "RECOMPUTE"})
                done = 1
            else:
                self.store.x("UPDATE tasks SET state='VERIFY_MISMATCH', error='DIGEST_MISMATCH' WHERE task_id=?",
                             (task["task_id"],))
                self.store.event(mission_id, "TASK_VERIFY_MISMATCH", {"task_id": task["task_id"]})
        except Exception as e:  # noqa: BLE001
            self.store.x("UPDATE tasks SET state='VERIFY_ERROR', error=? WHERE task_id=?",
                         (f"{type(e).__name__}: {e}", task["task_id"]))
        self.store.x("UPDATE missions SET last_progress=?, updated_at=? WHERE mission_id=?",
                     (time.time(), time.time(), mission_id))
        self._maybe_complete(mission_id)
        return done

    def _maybe_complete(self, mission_id: str) -> None:
        m = self.mission(mission_id)
        if m and m["state"] == "RUNNING" and m["task_states"] and set(m["task_states"]) <= {"VERIFIED"}:
            self.store.x("UPDATE missions SET state='COMPLETED', finished_at=?, updated_at=? WHERE mission_id=?",
                         (time.time(), time.time(), mission_id))
            self.store.event(mission_id, "MISSION_COMPLETED", {"tasks": m["tasks"]})

    def supervise(self, mission_id: str) -> Dict[str, Any]:
        """Deadline, budget and stall checks. Returns the current state (never silently continues a stalled mission)."""
        m = self.store.one("SELECT * FROM missions WHERE mission_id=?", (mission_id,))
        if not m or m["state"] != "RUNNING":
            return {"state": m["state"] if m else "MISSING"}
        now = time.time()
        budget = json.loads(m["budget"] or "{}")
        counts = self.mission(mission_id)
        if now > float(m["deadline"]):
            self.store.x("UPDATE missions SET state='DEADLINE_EXCEEDED', updated_at=? WHERE mission_id=?",
                         (now, mission_id))
            self.store.event(mission_id, "DEADLINE_EXCEEDED", {})
        elif int(budget.get("max_tasks") or 10 ** 9) < counts["tasks"]:
            self.store.x("UPDATE missions SET state='BUDGET_EXCEEDED', updated_at=? WHERE mission_id=?",
                         (now, mission_id))
        elif now - float(m["last_progress"]) > float(m["stall_s"]):
            self.store.x("UPDATE missions SET state='STALLED', updated_at=? WHERE mission_id=?", (now, mission_id))
            self.store.event(mission_id, "STALL_DETECTED", {"idle_s": round(now - float(m["last_progress"]), 1)})
        return self.mission(mission_id)

    def checkpoint(self, mission_id: str) -> Dict[str, Any]:
        m = self.mission(mission_id)
        tasks = self.store.q("SELECT task_id,state,attempt,result_digest,verifier_digest FROM tasks"
                            " WHERE mission_id=? ORDER BY seq", (mission_id,))
        body = {"mission": {k: m[k] for k in ("mission_id", "state", "tasks", "verified", "elapsed_s")},
                "tasks": tasks}
        cp = self.store.checkpoint(mission_id, body)
        self.store.event(mission_id, "CHECKPOINT", cp)
        return cp

    def accounting(self, mission_id: str) -> Dict[str, Any]:
        tasks = self.store.q("SELECT * FROM tasks WHERE mission_id=?", (mission_id,))
        events = self.store.q("SELECT type,body FROM events WHERE mission_id=? ORDER BY seq", (mission_id,))
        exec_s = 0.0
        verify_s = 0.0
        for e in events:
            try:
                sec = float(json.loads(e["body"]).get("seconds") or 0)
            except Exception:  # noqa: BLE001
                sec = 0.0
            if e["type"] == "TASK_EXECUTED":
                exec_s += sec
            elif e["type"] == "TASK_VERIFIED":
                verify_s += sec
        m = self.one("SELECT created_at,finished_at,deadline FROM missions WHERE mission_id=?", (mission_id,))
        elapsed = (float(m["finished_at"]) if m["finished_at"] else time.time()) - float(m["created_at"])
        return {"schema": SCHEMA, "mission_id": mission_id, "elapsed_s": round(elapsed, 2),
                "tasks": len(tasks), "verified": sum(1 for t in tasks if t["state"] == "VERIFIED"),
                "execute_seconds": round(exec_s, 3), "verification_seconds": round(verify_s, 3),
                "productive_seconds": round(exec_s + verify_s, 3),
                "productive_worker_hours": round((exec_s + verify_s) / 3600.0, 6),
                "checkpoints": self.store.verify_checkpoints(mission_id),
                "note": "elapsed wall-clock is not productive time; idle/sleep is excluded"}