"""DAG persistence via SQLite.

Stores the phase graph, checkpoints, and journal for architect runs.
Uses WAL mode for safe concurrent reads during observation.
"""

from __future__ import annotations

import json
import sqlite3
from collections import deque
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import Any


@dataclass
class Phase:
    """A single phase in the architect DAG."""

    id: str
    label: str
    description: str = ""
    phase_type: str = "cleave_run"  # cleave_run | direct
    status: str = "pending"  # pending | running | completed | failed | paused
    wave: int | None = None
    config_json: str = ""
    workspace_path: str = ""
    cost_usd: float = 0.0
    started_at: str = ""
    completed_at: str = ""
    error_message: str = ""
    observation_json: str = ""
    created_at: str = ""
    updated_at: str = ""

    def __post_init__(self) -> None:
        now = datetime.now(UTC).isoformat()
        if not self.created_at:
            self.created_at = now
        if not self.updated_at:
            self.updated_at = now

    def to_row(self) -> tuple:
        """Return a tuple for SQLite INSERT."""
        return (
            self.id, self.label, self.description, self.phase_type,
            self.status, self.wave, self.config_json, self.workspace_path,
            self.cost_usd, self.started_at, self.completed_at,
            self.error_message, self.observation_json,
            self.created_at, self.updated_at,
        )

    @classmethod
    def from_row(cls, row: sqlite3.Row) -> Phase:
        """Construct from a sqlite3.Row."""
        return cls(
            id=row["id"],
            label=row["label"],
            description=row["description"],
            phase_type=row["phase_type"],
            status=row["status"],
            wave=row["wave"],
            config_json=row["config_json"],
            workspace_path=row["workspace_path"],
            cost_usd=row["cost_usd"],
            started_at=row["started_at"],
            completed_at=row["completed_at"],
            error_message=row["error_message"],
            observation_json=row["observation_json"],
            created_at=row["created_at"],
            updated_at=row["updated_at"],
        )


@dataclass
class Checkpoint:
    """A gate evaluation checkpoint."""

    id: int | None = None
    phase_id: str | None = None
    gate_type: str = ""  # structural | budget | failure
    confidence_score: float | None = None
    decision: str = ""  # continue | pause | fail
    reason: str = ""
    created_at: str = ""

    def __post_init__(self) -> None:
        if not self.created_at:
            self.created_at = datetime.now(UTC).isoformat()

    @classmethod
    def from_row(cls, row: sqlite3.Row) -> Checkpoint:
        return cls(
            id=row["id"],
            phase_id=row["phase_id"],
            gate_type=row["gate_type"],
            confidence_score=row["confidence_score"],
            decision=row["decision"],
            reason=row["reason"],
            created_at=row["created_at"],
        )


_SCHEMA = """\
CREATE TABLE IF NOT EXISTS metadata (
    key TEXT PRIMARY KEY,
    value_json TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS phases (
    id TEXT PRIMARY KEY,
    label TEXT NOT NULL,
    description TEXT DEFAULT '',
    phase_type TEXT NOT NULL DEFAULT 'cleave_run',
    status TEXT NOT NULL DEFAULT 'pending',
    wave INTEGER,
    config_json TEXT DEFAULT '',
    workspace_path TEXT DEFAULT '',
    cost_usd REAL DEFAULT 0.0,
    started_at TEXT DEFAULT '',
    completed_at TEXT DEFAULT '',
    error_message TEXT DEFAULT '',
    observation_json TEXT DEFAULT '',
    created_at TEXT NOT NULL,
    updated_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS dependencies (
    phase_id TEXT NOT NULL REFERENCES phases(id),
    depends_on TEXT NOT NULL REFERENCES phases(id),
    PRIMARY KEY (phase_id, depends_on)
);

CREATE TABLE IF NOT EXISTS checkpoints (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    phase_id TEXT REFERENCES phases(id),
    gate_type TEXT NOT NULL,
    confidence_score REAL,
    decision TEXT NOT NULL,
    reason TEXT DEFAULT '',
    created_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS journal (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    timestamp TEXT NOT NULL,
    event_type TEXT NOT NULL,
    phase_id TEXT REFERENCES phases(id),
    message TEXT NOT NULL
);
"""


class DAGStore:
    """SQLite-backed DAG storage for architect phases."""

    def __init__(self, db_path: Path) -> None:
        self.db_path = db_path
        self._conn: sqlite3.Connection | None = None

    def open(self) -> None:
        """Connect and initialize the schema."""
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._conn = sqlite3.connect(str(self.db_path))
        self._conn.row_factory = sqlite3.Row
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._conn.execute("PRAGMA foreign_keys=ON")
        self._conn.executescript(_SCHEMA)
        self._conn.commit()

    def close(self) -> None:
        """Close the database connection."""
        if self._conn:
            self._conn.close()
            self._conn = None

    def __enter__(self) -> DAGStore:
        self.open()
        return self

    def __exit__(self, *exc: object) -> None:
        self.close()

    @property
    def _db(self) -> sqlite3.Connection:
        if self._conn is None:
            raise RuntimeError("DAGStore is not open — call open() or use as context manager")
        return self._conn

    # ── Phase CRUD ────────────────────────────────────────────────

    def add_phase(self, phase: Phase) -> None:
        self._db.execute(
            "INSERT INTO phases VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
            phase.to_row(),
        )
        self._db.commit()

    def get_phase(self, phase_id: str) -> Phase | None:
        row = self._db.execute("SELECT * FROM phases WHERE id = ?", (phase_id,)).fetchone()
        return Phase.from_row(row) if row else None

    def update_phase(self, phase_id: str, **fields: Any) -> None:
        if not fields:
            return
        fields["updated_at"] = datetime.now(UTC).isoformat()
        sets = ", ".join(f"{k} = ?" for k in fields)
        vals = list(fields.values()) + [phase_id]
        self._db.execute(f"UPDATE phases SET {sets} WHERE id = ?", vals)
        self._db.commit()

    def list_phases(self, status: str | None = None) -> list[Phase]:
        if status:
            rows = self._db.execute(
                "SELECT * FROM phases WHERE status = ? ORDER BY wave, id", (status,)
            ).fetchall()
        else:
            rows = self._db.execute("SELECT * FROM phases ORDER BY wave, id").fetchall()
        return [Phase.from_row(r) for r in rows]

    def remove_phase(self, phase_id: str) -> None:
        self._db.execute("DELETE FROM dependencies WHERE phase_id = ? OR depends_on = ?",
                         (phase_id, phase_id))
        self._db.execute("DELETE FROM phases WHERE id = ?", (phase_id,))
        self._db.commit()

    # ── Dependencies ──────────────────────────────────────────────

    def add_dependency(self, phase_id: str, depends_on: str) -> None:
        self._db.execute(
            "INSERT OR IGNORE INTO dependencies VALUES (?, ?)",
            (phase_id, depends_on),
        )
        self._db.commit()

    def get_dependencies(self, phase_id: str) -> list[str]:
        rows = self._db.execute(
            "SELECT depends_on FROM dependencies WHERE phase_id = ?", (phase_id,)
        ).fetchall()
        return [r["depends_on"] for r in rows]

    def get_dependents(self, phase_id: str) -> list[str]:
        rows = self._db.execute(
            "SELECT phase_id FROM dependencies WHERE depends_on = ?", (phase_id,)
        ).fetchall()
        return [r["phase_id"] for r in rows]

    # ── Graph algorithms ──────────────────────────────────────────

    def ready_set(self) -> list[Phase]:
        """Return pending phases whose dependencies are all completed."""
        pending = self.list_phases(status="pending")
        ready = []
        for phase in pending:
            deps = self.get_dependencies(phase.id)
            if not deps:
                ready.append(phase)
                continue
            dep_statuses = []
            for dep_id in deps:
                dep = self.get_phase(dep_id)
                if dep:
                    dep_statuses.append(dep.status)
            if all(s == "completed" for s in dep_statuses):
                ready.append(phase)
        return ready

    def topological_sort(self) -> list[str]:
        """Kahn's algorithm for topological ordering."""
        phases = self.list_phases()
        in_degree: dict[str, int] = {p.id: 0 for p in phases}

        for phase in phases:
            deps = self.get_dependencies(phase.id)
            in_degree[phase.id] = len(deps)

        queue: deque[str] = deque(pid for pid, deg in in_degree.items() if deg == 0)
        result: list[str] = []

        while queue:
            node = queue.popleft()
            result.append(node)
            for dependent in self.get_dependents(node):
                in_degree[dependent] -= 1
                if in_degree[dependent] == 0:
                    queue.append(dependent)

        return result

    def detect_cycles(self) -> list[list[str]]:
        """DFS-based cycle detection. Returns list of cycles found."""
        phases = self.list_phases()
        phase_ids = {p.id for p in phases}
        visited: set[str] = set()
        rec_stack: set[str] = set()
        cycles: list[list[str]] = []

        def _dfs(node: str, path: list[str]) -> None:
            visited.add(node)
            rec_stack.add(node)
            path.append(node)

            for dependent in self.get_dependents(node):
                if dependent not in phase_ids:
                    continue
                if dependent not in visited:
                    _dfs(dependent, path)
                elif dependent in rec_stack:
                    # Found cycle: extract it from path
                    cycle_start = path.index(dependent)
                    cycles.append(path[cycle_start:] + [dependent])

            path.pop()
            rec_stack.discard(node)

        for pid in phase_ids:
            if pid not in visited:
                _dfs(pid, [])

        return cycles

    def compute_waves(self) -> list[list[str]]:
        """Level-order grouping based on dependency depth."""
        phases = self.list_phases()
        if not phases:
            return []

        # Build adjacency: phase_id -> list of dependencies
        depths: dict[str, int] = {}
        phase_ids = {p.id for p in phases}

        def _depth(pid: str, seen: set[str] | None = None) -> int:
            if pid in depths:
                return depths[pid]
            if seen is None:
                seen = set()
            if pid in seen:
                return 0  # cycle guard
            seen.add(pid)
            deps = [d for d in self.get_dependencies(pid) if d in phase_ids]
            if not deps:
                depths[pid] = 0
                return 0
            d = max(_depth(dep, seen) for dep in deps) + 1
            depths[pid] = d
            return d

        for p in phases:
            _depth(p.id)

        max_depth = max(depths.values()) if depths else 0
        waves: list[list[str]] = []
        for w in range(max_depth + 1):
            wave = sorted(pid for pid, d in depths.items() if d == w)
            if wave:
                waves.append(wave)

        return waves

    # ── Aggregate queries ─────────────────────────────────────────

    def total_cost(self) -> float:
        row = self._db.execute("SELECT COALESCE(SUM(cost_usd), 0.0) AS total FROM phases").fetchone()
        return float(row["total"])

    def all_completed(self) -> bool:
        row = self._db.execute(
            "SELECT COUNT(*) AS cnt FROM phases WHERE status != 'completed'"
        ).fetchone()
        return row["cnt"] == 0

    def has_failed(self) -> bool:
        row = self._db.execute(
            "SELECT COUNT(*) AS cnt FROM phases WHERE status = 'failed'"
        ).fetchone()
        return row["cnt"] > 0

    def consecutive_failures(self) -> int:
        """Count consecutive failed phases from the most recently completed/failed."""
        rows = self._db.execute(
            "SELECT status FROM phases WHERE status IN ('completed', 'failed') "
            "ORDER BY completed_at DESC"
        ).fetchall()
        count = 0
        for row in rows:
            if row["status"] == "failed":
                count += 1
            else:
                break
        return count

    # ── Checkpoints ───────────────────────────────────────────────

    def add_checkpoint(
        self,
        gate_type: str,
        decision: str,
        reason: str = "",
        phase_id: str | None = None,
        confidence_score: float | None = None,
    ) -> int:
        now = datetime.now(UTC).isoformat()
        cur = self._db.execute(
            "INSERT INTO checkpoints (phase_id, gate_type, confidence_score, decision, reason, created_at) "
            "VALUES (?, ?, ?, ?, ?, ?)",
            (phase_id, gate_type, confidence_score, decision, reason, now),
        )
        self._db.commit()
        return cur.lastrowid  # type: ignore[return-value]

    def get_checkpoints(self, phase_id: str | None = None) -> list[Checkpoint]:
        if phase_id:
            rows = self._db.execute(
                "SELECT * FROM checkpoints WHERE phase_id = ? ORDER BY created_at", (phase_id,)
            ).fetchall()
        else:
            rows = self._db.execute("SELECT * FROM checkpoints ORDER BY created_at").fetchall()
        return [Checkpoint.from_row(r) for r in rows]

    # ── Journal ───────────────────────────────────────────────────

    def log_event(self, event_type: str, message: str, phase_id: str | None = None) -> None:
        now = datetime.now(UTC).isoformat()
        self._db.execute(
            "INSERT INTO journal (timestamp, event_type, phase_id, message) VALUES (?, ?, ?, ?)",
            (now, event_type, phase_id, message),
        )
        self._db.commit()

    def get_journal(self, limit: int = 100) -> list[dict]:
        rows = self._db.execute(
            "SELECT * FROM journal ORDER BY timestamp DESC LIMIT ?", (limit,)
        ).fetchall()
        return [
            {
                "id": r["id"],
                "timestamp": r["timestamp"],
                "event_type": r["event_type"],
                "phase_id": r["phase_id"],
                "message": r["message"],
            }
            for r in rows
        ]

    # ── Metadata ──────────────────────────────────────────────────

    def set_metadata(self, key: str, value: Any) -> None:
        val_json = json.dumps(value)
        self._db.execute(
            "INSERT INTO metadata (key, value_json) VALUES (?, ?) "
            "ON CONFLICT(key) DO UPDATE SET value_json = excluded.value_json",
            (key, val_json),
        )
        self._db.commit()

    def get_metadata(self, key: str) -> Any | None:
        row = self._db.execute("SELECT value_json FROM metadata WHERE key = ?", (key,)).fetchone()
        if row is None:
            return None
        return json.loads(row["value_json"])

    # ── Export ─────────────────────────────────────────────────────

    def snapshot(self) -> dict:
        """Full JSON-serializable export of the DAG state."""
        phases = self.list_phases()
        return {
            "phases": [
                {
                    "id": p.id,
                    "label": p.label,
                    "description": p.description,
                    "phase_type": p.phase_type,
                    "status": p.status,
                    "wave": p.wave,
                    "cost_usd": p.cost_usd,
                    "started_at": p.started_at,
                    "completed_at": p.completed_at,
                    "error_message": p.error_message,
                    "dependencies": self.get_dependencies(p.id),
                }
                for p in phases
            ],
            "checkpoints": [
                {
                    "id": c.id,
                    "phase_id": c.phase_id,
                    "gate_type": c.gate_type,
                    "decision": c.decision,
                    "reason": c.reason,
                    "created_at": c.created_at,
                }
                for c in self.get_checkpoints()
            ],
            "journal": self.get_journal(limit=500),
            "total_cost_usd": self.total_cost(),
        }
