"""Protocol runs — the log of a protocol being executed.

A protocol is not just a document: starting one commits the agent to a fixed
order of actions, and the run is what records that the commitment was taken and
how it ended. The row snapshots the steps and the catalog version at start, so
editing the protocol mid-run cannot quietly change what the running agent owes
the user, and the log keeps saying which version was actually executed.

One run per session may be active. The steps themselves live in the session's
todo plan (``navi.tools.todo.insert_steps``), tagged with ``protocol_run_id``;
this table only holds the run's state and its closing report.
"""

from __future__ import annotations

import asyncio
import json
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any

import asyncpg
import structlog

log = structlog.get_logger()

STATUSES = ("running", "done", "failed", "cancelled")

_DDL = """
CREATE TABLE IF NOT EXISTS protocol_runs (
    id               TEXT PRIMARY KEY,
    protocol_id      TEXT NOT NULL DEFAULT '',
    protocol_name    TEXT NOT NULL DEFAULT '',
    protocol_version INTEGER NOT NULL DEFAULT 0,
    session_id       TEXT NOT NULL DEFAULT '',
    user_id          TEXT NOT NULL DEFAULT '',
    status           TEXT NOT NULL DEFAULT 'running'
                     CHECK (status IN ('running','done','failed','cancelled')),
    steps            JSONB NOT NULL DEFAULT '[]'::jsonb,
    post_checks      JSONB NOT NULL DEFAULT '[]'::jsonb,
    report           TEXT NOT NULL DEFAULT '',
    reason           TEXT NOT NULL DEFAULT '',
    started_at       TIMESTAMPTZ NOT NULL DEFAULT now(),
    finished_at      TIMESTAMPTZ
);
CREATE INDEX IF NOT EXISTS idx_protocol_runs_session
    ON protocol_runs (session_id, started_at DESC);
CREATE UNIQUE INDEX IF NOT EXISTS idx_protocol_runs_one_active
    ON protocol_runs (session_id) WHERE status = 'running';
"""

_SELECT = (
    "id, protocol_id, protocol_name, protocol_version, session_id, user_id, "
    "status, steps, post_checks, report, reason, started_at, finished_at"
)


def _json_list(value: Any) -> list[str]:
    if isinstance(value, str):
        # asyncpg hands JSONB back as text unless a codec is registered.
        try:
            value = json.loads(value)
        except ValueError:
            return []
    return [str(v) for v in value] if isinstance(value, list) else []


@dataclass
class ProtocolRun:
    id: str
    protocol_id: str
    protocol_name: str
    protocol_version: int
    session_id: str
    user_id: str
    status: str
    steps: list[str] = field(default_factory=list)
    post_checks: list[str] = field(default_factory=list)
    report: str = ""
    reason: str = ""
    started_at: datetime | None = None
    finished_at: datetime | None = None

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "protocol_id": self.protocol_id,
            "protocol_name": self.protocol_name,
            "protocol_version": self.protocol_version,
            "session_id": self.session_id,
            "user_id": self.user_id,
            "status": self.status,
            "steps": list(self.steps),
            "post_checks": list(self.post_checks),
            "report": self.report,
            "reason": self.reason,
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "finished_at": self.finished_at.isoformat() if self.finished_at else None,
        }


def _row_to_run(row) -> ProtocolRun:
    return ProtocolRun(
        id=row["id"],
        protocol_id=row["protocol_id"],
        protocol_name=row["protocol_name"],
        protocol_version=row["protocol_version"],
        session_id=row["session_id"],
        user_id=row["user_id"],
        status=row["status"],
        steps=_json_list(row["steps"]),
        post_checks=_json_list(row["post_checks"]),
        report=row["report"],
        reason=row["reason"],
        started_at=row["started_at"],
        finished_at=row["finished_at"],
    )


class ProtocolRunStore:
    """Runs in their own table; DDL runs once, lazily."""

    def __init__(self, pool: asyncpg.Pool) -> None:
        self._pool = pool
        self._ready = False
        self._lock = asyncio.Lock()

    async def _ensure_ready(self) -> None:
        if self._ready:
            return
        async with self._lock:
            if self._ready:
                return
            async with self._pool.acquire() as conn:
                await conn.execute(_DDL)
            self._ready = True

    async def start(
        self,
        *,
        protocol_id: str,
        protocol_name: str,
        protocol_version: int,
        session_id: str,
        user_id: str,
        steps: list[str],
        post_checks: list[str] | None = None,
    ) -> ProtocolRun:
        await self._ensure_ready()
        run_id = uuid.uuid4().hex
        async with self._pool.acquire() as conn:
            await conn.execute(
                """
                INSERT INTO protocol_runs
                    (id, protocol_id, protocol_name, protocol_version, session_id,
                     user_id, status, steps, post_checks)
                VALUES ($1, $2, $3, $4, $5, $6, 'running', $7::jsonb, $8::jsonb)
                """,
                run_id, protocol_id, protocol_name, int(protocol_version), session_id,
                user_id, json.dumps(steps, ensure_ascii=False),
                json.dumps(post_checks or [], ensure_ascii=False),
            )
        run = await self.get(run_id)
        assert run is not None  # the INSERT just succeeded
        return run

    async def get(self, run_id: str) -> ProtocolRun | None:
        await self._ensure_ready()
        async with self._pool.acquire() as conn:
            row = await conn.fetchrow(
                f"SELECT {_SELECT} FROM protocol_runs WHERE id = $1", run_id,
            )
        return _row_to_run(row) if row else None

    async def active(self, session_id: str) -> ProtocolRun | None:
        await self._ensure_ready()
        async with self._pool.acquire() as conn:
            row = await conn.fetchrow(
                f"SELECT {_SELECT} FROM protocol_runs "
                "WHERE session_id = $1 AND status = 'running' "
                "ORDER BY started_at DESC LIMIT 1",
                session_id,
            )
        return _row_to_run(row) if row else None

    async def close(
        self, run_id: str, *, status: str, report: str = "", reason: str = ""
    ) -> ProtocolRun | None:
        """Close a run. Only a ``running`` row is touched; an already closed one
        is returned as it is, so a repeated call cannot rewrite history."""
        if status not in ("done", "failed", "cancelled"):
            raise ValueError(f"cannot close a run as {status!r}")
        await self._ensure_ready()
        async with self._pool.acquire() as conn:
            await conn.execute(
                """
                UPDATE protocol_runs
                   SET status = $2, report = $3, reason = $4, finished_at = now()
                 WHERE id = $1 AND status = 'running'
                """,
                run_id, status, report, reason,
            )
        return await self.get(run_id)

    async def list(self, session_id: str, limit: int = 20) -> list[ProtocolRun]:
        await self._ensure_ready()
        async with self._pool.acquire() as conn:
            rows = await conn.fetch(
                f"SELECT {_SELECT} FROM protocol_runs WHERE session_id = $1 "
                "ORDER BY started_at DESC LIMIT $2",
                session_id, max(1, min(int(limit), 200)),
            )
        return [_row_to_run(r) for r in rows]