"""Every generation request/response pair, stored in one SQLite file for Stage 2 evaluation.

One row per request: what the tester asked, what was retrieved, every model
attempt with its validation result, the final output, plus two review
channels: the tester's own feedback (did it work when they ran it?) and the
Orcanion reviewer's verdict (``python -m api.review``).
"""

from __future__ import annotations

import datetime as dt
import json
import sqlite3
import uuid
from contextlib import contextmanager
from pathlib import Path

SCHEMA = """
CREATE TABLE IF NOT EXISTS generations (
    id              TEXT PRIMARY KEY,
    created_at      TEXT NOT NULL,
    tester          TEXT NOT NULL,
    kind            TEXT NOT NULL,          -- 'simulate' | 'code'
    request         TEXT NOT NULL,          -- JSON of the request body
    model           TEXT,
    sources         TEXT,                   -- JSON list of retrieved chunk sources
    attempts        TEXT,                   -- JSON list of {raw_output, validation, ...}
    output          TEXT,                   -- final MJCF, or JSON {path: content} for code
    valid           INTEGER,                -- automatic validity check of the final output
    validation      TEXT,                   -- JSON of the final validation result
    error           TEXT,                   -- server-side failure (LLM down, etc.)
    duration_ms     INTEGER,
    tester_rating   TEXT,                   -- 'works' | 'partly' | 'broken'
    tester_note     TEXT,
    review_verdict  TEXT,                   -- 'correct' | 'partial' | 'incorrect'
    review_note     TEXT,
    reviewed_at     TEXT
);
CREATE INDEX IF NOT EXISTS generations_created ON generations(created_at);

-- Every backend job: simulation runs, sweeps, ROS 2 checks, video reviews, dataset exports, Isaac renders
-- (usage visibility: who ran what, when, and what it cost).
CREATE TABLE IF NOT EXISTS runs (
    id              TEXT PRIMARY KEY,
    created_at      TEXT NOT NULL,
    tester          TEXT NOT NULL,
    kind            TEXT NOT NULL DEFAULT 'sim',  -- sim | sweep | ros_check | review | export | isaac
    parent_id       TEXT,                   -- the run a sweep/review/render was made from
    hardware        TEXT,                   -- price key: t4 | l4 | cpu | runpod
    params          TEXT,                   -- JSON: kind-specific request parameters
    cost_usd        REAL,                   -- estimated, set when the job ends
    generation_id   TEXT,                   -- source generation, if the scene came from /simulate/generate
    robot           TEXT,
    duration_s      REAL,
    controller      TEXT,
    backend         TEXT,                   -- modal | runpod
    job_id          TEXT,                   -- the backend's id for the job
    status          TEXT NOT NULL,          -- SUBMITTING | IN_QUEUE | IN_PROGRESS | COMPLETED | FAILED | ...
    error           TEXT,
    output          TEXT,                   -- JSON job output: summary + R2 keys
    execution_ms    INTEGER,                -- GPU execution time (RunPod's figure, else the worker's own)
    delay_ms        INTEGER,                -- queue + cold start (RunPod only)
    updated_at      TEXT
);
CREATE INDEX IF NOT EXISTS runs_tester_created ON runs(tester, created_at);

-- Every AI model call, for the rupee budget (api/budget.py).
CREATE TABLE IF NOT EXISTS usage (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    created_at      TEXT NOT NULL,
    tester          TEXT NOT NULL,
    generation_id   TEXT,
    category        TEXT NOT NULL,          -- coding | simulation_ai | debugging | testing
    model           TEXT,
    tier            TEXT,
    provider        TEXT,                   -- who billed it: openrouter, ...
    input_tokens    INTEGER,
    output_tokens   INTEGER,
    cost_usd        REAL,
    cost_inr        REAL,
    billable        INTEGER NOT NULL,       -- 0: OpenRouter-exempt, recorded but not deducted
    exempt_reason   TEXT,
    estimated       INTEGER NOT NULL DEFAULT 0,  -- 1: cost computed from list prices, not reported by the provider
    estimate_inr    REAL                    -- what the plan expected before running
);
CREATE INDEX IF NOT EXISTS usage_tester ON usage(tester, created_at);

-- Lightweight feedback on anything: generations, runs, checks, routing choices.
CREATE TABLE IF NOT EXISTS feedback (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    created_at      TEXT NOT NULL,
    tester          TEXT NOT NULL,
    target_type     TEXT NOT NULL,          -- generation | job
    target_id       TEXT NOT NULL,
    rating          TEXT NOT NULL,          -- helpful | not_helpful
    note            TEXT,
    context         TEXT                    -- JSON: task, routing tier/model, simulator, result state, ...
);
"""

JSON_FIELDS = ("request", "sources", "attempts", "validation")


def _decode_run(d: dict) -> dict:
    d["output"] = json.loads(d["output"]) if d.get("output") else None
    d["params"] = json.loads(d["params"]) if d.get("params") else {}
    return d


def _now() -> str:
    return dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds")


class LogStore:
    def __init__(self, path: Path):
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with self._conn() as c:
            c.executescript(SCHEMA)
            self._migrate(c)

    @staticmethod
    def _migrate(c) -> None:
        """Bring run tables created by earlier versions up to the current schema."""
        cols = {r[1] for r in c.execute("PRAGMA table_info(runs)")}
        if "runpod_job_id" in cols:  # before Modal support
            c.execute("ALTER TABLE runs RENAME COLUMN runpod_job_id TO job_id")
        added = {
            "backend": "TEXT",
            "kind": "TEXT NOT NULL DEFAULT 'sim'",
            "parent_id": "TEXT",
            "hardware": "TEXT",
            "params": "TEXT",
            "cost_usd": "REAL",
        }
        for name, decl in added.items():
            if name not in cols:
                c.execute(f"ALTER TABLE runs ADD COLUMN {name} {decl}")
        if "runpod_job_id" in cols and "backend" not in cols:
            c.execute("UPDATE runs SET backend = 'runpod' WHERE job_id IS NOT NULL")
        c.execute("CREATE INDEX IF NOT EXISTS runs_created ON runs(created_at)")

    @contextmanager
    def _conn(self):
        conn = sqlite3.connect(self.path, timeout=30)
        conn.row_factory = sqlite3.Row
        try:
            conn.execute("PRAGMA journal_mode=WAL")
            yield conn
            conn.commit()
        finally:
            conn.close()

    def record(
        self,
        *,
        tester: str,
        kind: str,
        request: dict,
        model: str | None = None,
        sources: list | None = None,
        attempts: list | None = None,
        output: str | None = None,
        valid: bool | None = None,
        validation: dict | None = None,
        error: str | None = None,
        duration_ms: int | None = None,
    ) -> str:
        gen_id = uuid.uuid4().hex[:12]
        with self._conn() as c:
            c.execute(
                "INSERT INTO generations (id, created_at, tester, kind, request, model, sources, attempts, output, "
                "valid, validation, error, duration_ms) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)",
                (
                    gen_id,
                    _now(),
                    tester,
                    kind,
                    json.dumps(request),
                    model,
                    json.dumps(sources or []),
                    json.dumps(attempts or []),
                    output,
                    None if valid is None else int(valid),
                    json.dumps(validation or {}),
                    error,
                    duration_ms,
                ),
            )
        return gen_id

    def get(self, gen_id: str) -> dict | None:
        with self._conn() as c:
            row = c.execute("SELECT * FROM generations WHERE id = ?", (gen_id,)).fetchone()
        return self._decode(row) if row else None

    def list(
        self, kind: str | None = None, tester: str | None = None, unreviewed: bool = False, limit: int = 50
    ) -> list[dict]:
        sql, args = "SELECT * FROM generations WHERE 1=1", []
        if kind:
            sql += " AND kind = ?"
            args.append(kind)
        if tester:
            sql += " AND tester = ?"
            args.append(tester)
        if unreviewed:
            sql += " AND review_verdict IS NULL"
        sql += " ORDER BY created_at DESC LIMIT ?"
        args.append(limit)
        with self._conn() as c:
            return [self._decode(r) for r in c.execute(sql, args).fetchall()]

    def all(self) -> list[dict]:
        with self._conn() as c:
            return [self._decode(r) for r in c.execute("SELECT * FROM generations ORDER BY created_at").fetchall()]

    def set_feedback(self, gen_id: str, tester: str, rating: str, note: str) -> bool:
        with self._conn() as c:
            cur = c.execute(
                "UPDATE generations SET tester_rating = ?, tester_note = ? WHERE id = ? AND tester = ?",
                (rating, note, gen_id, tester),
            )
        return cur.rowcount == 1

    def set_review(self, gen_id: str, verdict: str, note: str) -> bool:
        with self._conn() as c:
            cur = c.execute(
                "UPDATE generations SET review_verdict = ?, review_note = ?, reviewed_at = ? WHERE id = ?",
                (verdict, note, _now(), gen_id),
            )
        return cur.rowcount == 1

    def summary(self) -> list[dict]:
        sql = """SELECT kind, COUNT(*) AS total, SUM(valid = 1) AS auto_valid,
                        SUM(tester_rating = 'works') AS tester_works, SUM(tester_rating IS NOT NULL) AS tester_rated,
                        SUM(review_verdict = 'correct') AS reviewed_correct,
                        SUM(review_verdict IS NOT NULL) AS reviewed
                 FROM generations GROUP BY kind ORDER BY kind"""
        with self._conn() as c:
            return [dict(r) for r in c.execute(sql).fetchall()]

    # ------------------------------------------------------------------ AI usage and feedback

    def record_usage(self, **row) -> None:
        allowed = {"created_at", "tester", "generation_id", "category", "model", "tier",
                   "provider", "input_tokens", "output_tokens", "cost_usd", "cost_inr",
                   "billable", "exempt_reason", "estimated", "estimate_inr"}
        row = {k: v for k, v in row.items() if k in allowed}
        row = {"created_at": _now(), **row}
        cols = ", ".join(row)
        with self._conn() as c:
            c.execute(f"INSERT INTO usage ({cols}) VALUES ({', '.join('?' * len(row))})", tuple(row.values()))

    def usage_rows(self, tester: str | None = None, since_iso: str | None = None) -> list[dict]:
        sql, args = "SELECT * FROM usage WHERE 1=1", []
        if tester:
            sql += " AND tester = ?"
            args.append(tester)
        if since_iso:
            sql += " AND created_at >= ?"
            args.append(since_iso)
        with self._conn() as c:
            return [dict(r) for r in c.execute(sql + " ORDER BY id", args).fetchall()]

    def record_feedback(self, tester: str, target_type: str, target_id: str, rating: str, note: str, context: dict):
        with self._conn() as c:
            c.execute(
                "INSERT INTO feedback (created_at, tester, target_type, target_id, rating, note, context) "
                "VALUES (?,?,?,?,?,?,?)",
                (_now(), tester, target_type, target_id, rating, note, json.dumps(context)),
            )

    def feedback_rows(self) -> list[dict]:
        with self._conn() as c:
            rows = [dict(r) for r in c.execute("SELECT * FROM feedback ORDER BY id").fetchall()]
        for r in rows:
            r["context"] = json.loads(r["context"]) if r.get("context") else {}
        return rows

    # ------------------------------------------------------------------ runs

    def record_run(
        self,
        *,
        run_id: str,
        tester: str,
        generation_id: str | None,
        robot: str | None,
        duration_s: float | None,
        controller: str | None,
        kind: str = "sim",
        parent_id: str | None = None,
        hardware: str | None = None,
        params: dict | None = None,
    ) -> None:
        now = _now()
        with self._conn() as c:
            c.execute(
                "INSERT INTO runs (id, created_at, tester, kind, parent_id, hardware, params, generation_id, robot, "
                "duration_s, controller, status, updated_at) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)",
                (
                    run_id,
                    now,
                    tester,
                    kind,
                    parent_id,
                    hardware,
                    json.dumps(params) if params is not None else None,
                    generation_id,
                    robot,
                    duration_s,
                    controller,
                    "SUBMITTING",
                    now,
                ),
            )

    def update_run(self, run_id: str, **fields) -> None:
        allowed = {"backend", "job_id", "status", "error", "output", "execution_ms", "delay_ms", "cost_usd"}
        if set(fields) - allowed:
            raise ValueError(f"unknown run fields {set(fields) - allowed}")
        if "output" in fields and not isinstance(fields["output"], (str, type(None))):
            fields["output"] = json.dumps(fields["output"])
        fields["updated_at"] = _now()
        cols = ", ".join(f"{k} = ?" for k in fields)
        with self._conn() as c:
            c.execute(f"UPDATE runs SET {cols} WHERE id = ?", (*fields.values(), run_id))

    def get_run(self, run_id: str) -> dict | None:
        with self._conn() as c:
            row = c.execute("SELECT * FROM runs WHERE id = ?", (run_id,)).fetchone()
        if not row:
            return None
        return _decode_run(dict(row))

    def list_runs(self, tester: str | None = None, limit: int = 50) -> list[dict]:
        sql, args = "SELECT * FROM runs", []
        if tester:
            sql += " WHERE tester = ?"
            args.append(tester)
        sql += " ORDER BY created_at DESC LIMIT ?"
        args.append(limit)
        with self._conn() as c:
            rows = [dict(r) for r in c.execute(sql, args).fetchall()]
        return [_decode_run(r) for r in rows]

    def jobs_since(self, since_iso: str, kind: str | None = None) -> list[dict]:
        """Every job created at or after `since_iso` (for spend and daily-limit checks)."""
        sql, args = "SELECT * FROM runs WHERE created_at >= ?", [since_iso]
        if kind:
            sql += " AND kind = ?"
            args.append(kind)
        with self._conn() as c:
            return [_decode_run(dict(r)) for r in c.execute(sql, args).fetchall()]

    def children(self, parent_id: str, kind: str | None = None) -> list[dict]:
        sql, args = "SELECT * FROM runs WHERE parent_id = ?", [parent_id]
        if kind:
            sql += " AND kind = ?"
            args.append(kind)
        sql += " ORDER BY created_at"
        with self._conn() as c:
            return [_decode_run(dict(r)) for r in c.execute(sql, args).fetchall()]

    def runs_since(self, tester: str, since_iso: str) -> int:
        with self._conn() as c:
            return c.execute(
                "SELECT COUNT(*) FROM runs WHERE tester = ? AND created_at >= ?", (tester, since_iso)
            ).fetchone()[0]

    def usage(self, window_minutes: int = 10) -> list[dict]:
        """Per tester: total jobs, jobs in the last 24 h, busiest window, compute seconds, estimated cost."""
        with self._conn() as c:
            rows = c.execute(
                "SELECT tester, created_at, execution_ms, cost_usd FROM runs ORDER BY tester, created_at"
            ).fetchall()
        day_ago = (dt.datetime.now(dt.timezone.utc) - dt.timedelta(days=1)).isoformat(timespec="seconds")
        out: dict[str, dict] = {}
        times: dict[str, list[dt.datetime]] = {}
        for tester, created, exec_ms, cost in rows:
            u = out.setdefault(
                tester,
                {
                    "tester": tester,
                    "runs": 0,
                    "last_24h": 0,
                    "gpu_seconds": 0.0,
                    "cost_usd": 0.0,
                    "first": created,
                    "last": created,
                },
            )
            u["runs"] += 1
            u["last"] = created
            u["last_24h"] += created >= day_ago
            u["gpu_seconds"] += (exec_ms or 0) / 1000
            u["cost_usd"] += cost or 0.0
            times.setdefault(tester, []).append(dt.datetime.fromisoformat(created))
        window = dt.timedelta(minutes=window_minutes)
        for tester, ts in times.items():
            best, j = 0, 0
            for i, t in enumerate(ts):
                while t - ts[j] > window:
                    j += 1
                best = max(best, i - j + 1)
            out[tester][f"max_in_{window_minutes}min"] = best
            out[tester]["gpu_seconds"] = round(out[tester]["gpu_seconds"], 1)
            out[tester]["cost_usd"] = round(out[tester]["cost_usd"], 3)
        return sorted(out.values(), key=lambda u: -u["runs"])

    @staticmethod
    def _decode(row: sqlite3.Row) -> dict:
        d = dict(row)
        for f in JSON_FIELDS:
            if d.get(f):
                d[f] = json.loads(d[f])
        if d.get("valid") is not None:
            d["valid"] = bool(d["valid"])
        return d