"""SQLite result storage for eval runs and cases."""

from __future__ import annotations

import json
import sqlite3
from datetime import datetime
from pathlib import Path

from hatchdx.agent.eval.models import (
    AssertionResult,
    EvalCaseResult,
    EvalRunResult,
)


# ---------------------------------------------------------------------------
# Schema
# ---------------------------------------------------------------------------

_SCHEMA = """\
CREATE TABLE IF NOT EXISTS eval_runs (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    agent_name TEXT NOT NULL,
    suite_name TEXT NOT NULL,
    model TEXT NOT NULL,
    temperature REAL NOT NULL,
    started_at TEXT NOT NULL,
    completed_at TEXT,
    total_cases INTEGER NOT NULL DEFAULT 0,
    passed INTEGER NOT NULL DEFAULT 0,
    failed INTEGER NOT NULL DEFAULT 0,
    total_cost_usd REAL NOT NULL DEFAULT 0.0,
    total_tokens INTEGER NOT NULL DEFAULT 0
);

CREATE TABLE IF NOT EXISTS eval_cases (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    run_id INTEGER NOT NULL REFERENCES eval_runs(id),
    case_name TEXT NOT NULL,
    input TEXT NOT NULL,
    output TEXT,
    status TEXT NOT NULL,
    tool_calls_json TEXT NOT NULL DEFAULT '[]',
    assertion_results_json TEXT NOT NULL DEFAULT '[]',
    latency_ms REAL NOT NULL DEFAULT 0.0,
    tokens_used INTEGER NOT NULL DEFAULT 0,
    cost_usd REAL NOT NULL DEFAULT 0.0,
    error_message TEXT
);
"""


# ---------------------------------------------------------------------------
# Storage class
# ---------------------------------------------------------------------------


class EvalStorage:
    """SQLite-backed storage for eval results."""

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

    def _connect(self) -> sqlite3.Connection:
        if self._conn is None:
            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.executescript(_SCHEMA)
        return self._conn

    def close(self) -> None:
        if self._conn is not None:
            self._conn.close()
            self._conn = None

    def save_run(self, result: EvalRunResult) -> int:
        """Save a complete eval run. Returns the run ID."""
        conn = self._connect()
        cursor = conn.execute(
            """INSERT INTO eval_runs
               (agent_name, suite_name, model, temperature,
                started_at, completed_at,
                total_cases, passed, failed,
                total_cost_usd, total_tokens)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                result.agent_name,
                result.suite_name,
                result.model,
                result.temperature,
                result.started_at.isoformat(),
                result.completed_at.isoformat() if result.completed_at else None,
                result.total_cases,
                result.passed,
                result.failed,
                result.total_cost,
                result.total_tokens,
            ),
        )
        run_id = cursor.lastrowid

        for case_result in result.case_results:
            assertion_dicts = [
                {
                    "type": ar.assertion_type,
                    "passed": ar.passed,
                    "message": ar.message,
                }
                for ar in case_result.assertion_results
            ]
            conn.execute(
                """INSERT INTO eval_cases
                   (run_id, case_name, input, output, status,
                    tool_calls_json, assertion_results_json,
                    latency_ms, tokens_used, cost_usd, error_message)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    run_id,
                    case_result.case_name,
                    case_result.input,
                    case_result.output,
                    "passed" if case_result.passed else "failed",
                    json.dumps(case_result.tool_calls_json),
                    json.dumps(assertion_dicts),
                    case_result.latency_ms,
                    case_result.tokens_used,
                    case_result.cost_usd,
                    case_result.error_message,
                ),
            )

        conn.commit()
        return run_id

    def get_latest_run(
        self, agent_name: str, suite_name: str
    ) -> dict | None:
        """Get the most recent run for a given agent+suite (for comparison)."""
        conn = self._connect()
        row = conn.execute(
            """SELECT * FROM eval_runs
               WHERE agent_name = ? AND suite_name = ?
               ORDER BY id DESC LIMIT 1""",
            (agent_name, suite_name),
        ).fetchone()
        if row is None:
            return None
        return dict(row)

    def get_run_cases(self, run_id: int) -> list[dict]:
        """Get all case results for a run."""
        conn = self._connect()
        rows = conn.execute(
            "SELECT * FROM eval_cases WHERE run_id = ? ORDER BY id",
            (run_id,),
        ).fetchall()
        return [dict(r) for r in rows]

    def get_previous_run(
        self, agent_name: str, suite_name: str, before_id: int
    ) -> dict | None:
        """Get the run immediately before a given run ID (for regression detection)."""
        conn = self._connect()
        row = conn.execute(
            """SELECT * FROM eval_runs
               WHERE agent_name = ? AND suite_name = ? AND id < ?
               ORDER BY id DESC LIMIT 1""",
            (agent_name, suite_name, before_id),
        ).fetchone()
        if row is None:
            return None
        return dict(row)
