"""SQLite state database helpers for RCA and snapshot metadata."""

from __future__ import annotations

import json
import sqlite3
import time
from collections.abc import Iterator
from contextlib import contextmanager
from pathlib import Path
from threading import Lock
from typing import Any, cast

STATE_DB_FILENAME = "llmdebug_state.sqlite3"
_BUSY_TIMEOUT_MS = 30_000
_WAL_RETRY_ATTEMPTS = 8
_WAL_BASE_SLEEP_SEC = 0.01
_WAL_INIT_LOCK = Lock()
_WAL_INITIALIZED_PATHS: set[str] = set()


def state_db_path(out_dir: str | Path) -> Path:
    """Return the SQLite database path for a snapshot output directory."""
    return Path(out_dir) / STATE_DB_FILENAME


def _enable_wal_mode(
    conn: sqlite3.Connection,
    *,
    attempts: int = _WAL_RETRY_ATTEMPTS,
    base_sleep_sec: float = _WAL_BASE_SLEEP_SEC,
) -> None:
    """Set WAL mode with retries to tolerate brief startup lock contention."""
    total_attempts = max(1, attempts)
    for attempt in range(total_attempts):
        try:
            conn.execute("PRAGMA journal_mode=WAL")
            return
        except sqlite3.OperationalError as exc:
            if "locked" not in str(exc).lower() or attempt == total_attempts - 1:
                raise
            if base_sleep_sec > 0:
                time.sleep(base_sleep_sec * float(attempt + 1))


def _configure_wal_mode_once(conn: sqlite3.Connection, *, db_path: Path) -> None:
    key = str(db_path.resolve())
    with _WAL_INIT_LOCK:
        if key in _WAL_INITIALIZED_PATHS:
            return
        _enable_wal_mode(conn)
        _WAL_INITIALIZED_PATHS.add(key)


def _configure_connection(conn: sqlite3.Connection, *, db_path: Path) -> None:
    conn.execute(f"PRAGMA busy_timeout={_BUSY_TIMEOUT_MS}")
    _configure_wal_mode_once(conn, db_path=db_path)
    conn.execute("PRAGMA synchronous=NORMAL")
    conn.execute("PRAGMA temp_store=MEMORY")


def _ensure_schema(conn: sqlite3.Connection) -> None:
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS rca_attempts (
            attempt_id TEXT PRIMARY KEY,
            session_id TEXT NOT NULL,
            updated_at TEXT NOT NULL,
            created_at TEXT NOT NULL,
            failure_signature TEXT NOT NULL,
            state TEXT NOT NULL,
            payload_json TEXT NOT NULL
        )
        """
    )
    conn.execute(
        """
        CREATE INDEX IF NOT EXISTS idx_rca_attempts_session_updated
            ON rca_attempts(session_id, updated_at DESC)
        """
    )
    conn.execute(
        """
        CREATE INDEX IF NOT EXISTS idx_rca_attempts_updated
            ON rca_attempts(updated_at DESC)
        """
    )
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS snapshots (
            path TEXT PRIMARY KEY,
            mtime REAL NOT NULL,
            name TEXT NOT NULL,
            exc_type TEXT NOT NULL,
            exc_message TEXT NOT NULL,
            file TEXT NOT NULL,
            line TEXT NOT NULL
        )
        """
    )
    conn.execute(
        """
        CREATE INDEX IF NOT EXISTS idx_snapshots_mtime
            ON snapshots(mtime DESC)
        """
    )


@contextmanager
def open_state_db(out_dir: str | Path) -> Iterator[sqlite3.Connection]:
    """Open the state DB connection with schema initialization."""
    path = state_db_path(out_dir)
    path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(path, timeout=30.0)
    try:
        _configure_connection(conn, db_path=path)
        _ensure_schema(conn)
        yield conn
        conn.commit()
    finally:
        conn.close()


def count_rca_records(out_dir: str | Path) -> int:
    with open_state_db(out_dir) as conn:
        row = conn.execute("SELECT COUNT(*) FROM rca_attempts").fetchone()
    return int(row[0]) if row else 0


def append_rca_record(out_dir: str | Path, record: dict[str, Any], *, max_records: int) -> None:
    payload = json.dumps(record, ensure_ascii=False, sort_keys=True)
    with open_state_db(out_dir) as conn:
        conn.execute(
            """
            INSERT OR REPLACE INTO rca_attempts(
                attempt_id, session_id, updated_at, created_at, failure_signature, state, payload_json
            )
            VALUES (?, ?, ?, ?, ?, ?, ?)
            """,
            (
                str(record.get("attempt_id", "")),
                str(record.get("session_id", "")),
                str(record.get("updated_at", "")),
                str(record.get("created_at", "")),
                str(record.get("failure_signature", "")),
                str(record.get("state", "")),
                payload,
            ),
        )
        if max_records > 0:
            conn.execute(
                """
                DELETE FROM rca_attempts
                WHERE rowid IN (
                    SELECT rowid
                    FROM rca_attempts
                    ORDER BY updated_at DESC, rowid DESC
                    LIMIT -1 OFFSET ?
                )
                """,
                (max_records,),
            )


def append_rca_records_bulk(out_dir: str | Path, records: list[dict[str, Any]]) -> None:
    if not records:
        return
    rows = [
        (
            str(record.get("attempt_id", "")),
            str(record.get("session_id", "")),
            str(record.get("updated_at", "")),
            str(record.get("created_at", "")),
            str(record.get("failure_signature", "")),
            str(record.get("state", "")),
            json.dumps(record, ensure_ascii=False, sort_keys=True),
        )
        for record in records
    ]
    with open_state_db(out_dir) as conn:
        conn.executemany(
            """
            INSERT OR REPLACE INTO rca_attempts(
                attempt_id, session_id, updated_at, created_at, failure_signature, state, payload_json
            )
            VALUES (?, ?, ?, ?, ?, ?, ?)
            """,
            rows,
        )


def _decode_payload_rows(rows: list[tuple[str]]) -> list[dict[str, Any]]:
    decoded: list[dict[str, Any]] = []
    for (payload_json,) in rows:
        try:
            payload = json.loads(payload_json)
        except json.JSONDecodeError:
            continue
        if isinstance(payload, dict):
            decoded.append(cast(dict[str, Any], payload))
    return decoded


def fetch_latest_rca_record(
    out_dir: str | Path, *, session_id: str | None = None
) -> dict[str, Any] | None:
    sql = "SELECT payload_json FROM rca_attempts"
    params: tuple[Any, ...] = ()
    if session_id:
        sql += " WHERE session_id = ?"
        params = (session_id,)
    sql += " ORDER BY updated_at DESC, rowid DESC LIMIT 1"
    with open_state_db(out_dir) as conn:
        rows = conn.execute(sql, params).fetchall()
    decoded = _decode_payload_rows(rows)
    return decoded[0] if decoded else None


def fetch_rca_history_records(
    out_dir: str | Path,
    *,
    session_id: str | None = None,
    limit: int = 20,
) -> list[dict[str, Any]]:
    sql = "SELECT payload_json FROM rca_attempts"
    params: list[Any] = []
    if session_id:
        sql += " WHERE session_id = ?"
        params.append(session_id)
    sql += " ORDER BY updated_at DESC, rowid DESC"
    if limit > 0:
        sql += " LIMIT ?"
        params.append(limit)
    with open_state_db(out_dir) as conn:
        rows = conn.execute(sql, tuple(params)).fetchall()
    return _decode_payload_rows(rows)


def upsert_snapshot_record(out_dir: str | Path, record: dict[str, Any]) -> None:
    with open_state_db(out_dir) as conn:
        conn.execute(
            """
            INSERT OR REPLACE INTO snapshots(path, mtime, name, exc_type, exc_message, file, line)
            VALUES (?, ?, ?, ?, ?, ?, ?)
            """,
            (
                str(record.get("path", "")),
                float(record.get("mtime", 0.0)),
                str(record.get("name", "?")),
                str(record.get("exc_type", "?")),
                str(record.get("exc_message", "")),
                str(record.get("file", "?")),
                str(record.get("line", "?")),
            ),
        )


def remove_snapshot_records(out_dir: str | Path, paths: list[str]) -> None:
    if not paths:
        return
    with open_state_db(out_dir) as conn:
        conn.executemany("DELETE FROM snapshots WHERE path = ?", [(path,) for path in paths])


def list_snapshot_records(
    out_dir: str | Path,
    *,
    limit: int | None = None,
) -> list[dict[str, Any]]:
    sql = """
        SELECT path, mtime, name, exc_type, exc_message, file, line
        FROM snapshots
        ORDER BY mtime DESC, path DESC
    """
    params: tuple[Any, ...] = ()
    if limit is not None and limit > 0:
        sql += " LIMIT ?"
        params = (limit,)
    with open_state_db(out_dir) as conn:
        rows = conn.execute(sql, params).fetchall()
    out: list[dict[str, Any]] = []
    for path, mtime, name, exc_type, exc_message, file_name, line in rows:
        out.append(
            {
                "path": str(path),
                "mtime": float(mtime),
                "name": str(name),
                "exc_type": str(exc_type),
                "exc_message": str(exc_message),
                "file": str(file_name),
                "line": str(line),
            }
        )
    return out


def replace_snapshot_records(out_dir: str | Path, records: list[dict[str, Any]]) -> None:
    with open_state_db(out_dir) as conn:
        conn.execute("DELETE FROM snapshots")
        if records:
            rows = [
                (
                    str(record.get("path", "")),
                    float(record.get("mtime", 0.0)),
                    str(record.get("name", "?")),
                    str(record.get("exc_type", "?")),
                    str(record.get("exc_message", "")),
                    str(record.get("file", "?")),
                    str(record.get("line", "?")),
                )
                for record in records
            ]
            conn.executemany(
                """
                INSERT OR REPLACE INTO snapshots(path, mtime, name, exc_type, exc_message, file, line)
                VALUES (?, ?, ?, ?, ?, ?, ?)
                """,
                rows,
            )
