"""
SQLite-backed index of all benchmark runs.

The JSON files remain the authoritative source of truth.
This module maintains a fast read-side cache in `runs_cache.db`
so the History page never has to scan thousands of files.

Schema
──────
runs                  – one row per run JSON file
stats_by_llm          – precomputed per-LLM aggregate (rebuilt on every write)
stats_by_llm_model    – precomputed per-(LLM, model_group) aggregate
stats_by_milestone    – precomputed per-milestone aggregate
"""

import json
import os
import sqlite3
import threading
from typing import Optional

import pandas as pd

_DB_NAME = "runs_cache.db"

# ── Base schema ───────────────────────────────────────────────────────────────
_CREATE_BASE_SQL = """
CREATE TABLE IF NOT EXISTS runs (
    file                    TEXT PRIMARY KEY,
    timestamp               TEXT,
    llm_used                TEXT,
    input_group             TEXT,
    model_group             TEXT,
    milestone               INTEGER,
    model_attractors_count  REAL,
    total_phenotypes        INTEGER,
    verified_count          INTEGER,
    success_rate            REAL,
    coverage                REAL,
    request_tokens          INTEGER,
    response_tokens         INTEGER,
    total_tokens            INTEGER,
    prompt                  TEXT
);
CREATE INDEX IF NOT EXISTS idx_runs_llm       ON runs(llm_used);
CREATE INDEX IF NOT EXISTS idx_runs_milestone ON runs(milestone);
CREATE INDEX IF NOT EXISTS idx_runs_timestamp ON runs(timestamp);
"""

# ── Precomputed summary tables ────────────────────────────────────────────────
_CREATE_SUMMARY_SQL = """
CREATE TABLE IF NOT EXISTS stats_by_llm (
    llm_used         TEXT PRIMARY KEY,
    total_runs       INTEGER NOT NULL DEFAULT 0,
    correct_runs     INTEGER NOT NULL DEFAULT 0,
    avg_success_rate REAL,
    avg_coverage     REAL,
    avg_tokens       REAL
);

CREATE TABLE IF NOT EXISTS stats_by_llm_model (
    llm_used         TEXT NOT NULL,
    model_group      TEXT NOT NULL,
    total_runs       INTEGER NOT NULL DEFAULT 0,
    correct_runs     INTEGER NOT NULL DEFAULT 0,
    avg_success_rate REAL,
    avg_coverage     REAL,
    avg_tokens       REAL,
    PRIMARY KEY (llm_used, model_group)
);

CREATE TABLE IF NOT EXISTS stats_by_milestone (
    milestone         INTEGER PRIMARY KEY,
    total_runs        INTEGER NOT NULL DEFAULT 0,
    total_phenotypes  INTEGER NOT NULL DEFAULT 0,
    verified_count    INTEGER NOT NULL DEFAULT 0,
    avg_success_rate  REAL,
    avg_coverage      REAL,
    total_tokens      INTEGER
);
"""

_CREATE_SQL = _CREATE_BASE_SQL + _CREATE_SUMMARY_SQL

# ── Summary rebuild SQL ───────────────────────────────────────────────────────
# coverage uses plain AVG (NULLs ignored) so LLMs with no coverage data show NULL.
# correct_runs uses COALESCE(coverage,0) to treat missing coverage as 0 (not perfect).
_REFRESH_SUMMARIES_SQL = """
DELETE FROM stats_by_llm;
INSERT INTO stats_by_llm
    SELECT llm_used,
           COUNT(*)                                                                            AS total_runs,
           SUM(CASE WHEN success_rate >= 100 AND COALESCE(coverage,0) >= 100 THEN 1 ELSE 0 END) AS correct_runs,
           ROUND(AVG(success_rate), 1)                                                         AS avg_success_rate,
           ROUND(AVG(coverage), 1)                                                             AS avg_coverage,
           ROUND(AVG(total_tokens), 0)                                                         AS avg_tokens
    FROM runs GROUP BY llm_used;

DELETE FROM stats_by_llm_model;
INSERT INTO stats_by_llm_model
    SELECT llm_used, model_group,
           COUNT(*)                                                                            AS total_runs,
           SUM(CASE WHEN success_rate >= 100 AND COALESCE(coverage,0) >= 100 THEN 1 ELSE 0 END) AS correct_runs,
           ROUND(AVG(success_rate), 1)                                                         AS avg_success_rate,
           ROUND(AVG(coverage), 1)                                                             AS avg_coverage,
           ROUND(AVG(total_tokens), 0)                                                         AS avg_tokens
    FROM runs GROUP BY llm_used, model_group;

DELETE FROM stats_by_milestone;
INSERT INTO stats_by_milestone
    SELECT milestone,
           COUNT(*)                              AS total_runs,
           SUM(total_phenotypes)                 AS total_phenotypes,
           SUM(verified_count)                   AS verified_count,
           ROUND(AVG(success_rate), 1)           AS avg_success_rate,
           ROUND(AVG(coverage), 1)               AS avg_coverage,
           SUM(COALESCE(total_tokens, 0))        AS total_tokens
    FROM runs GROUP BY milestone;
"""

# ── Persistent per-path connection pool ──────────────────────────────────────
_conn_cache: dict[str, sqlite3.Connection] = {}
_conn_lock = threading.Lock()


def _db_path(runs_dir: str) -> str:
    return os.path.join(runs_dir, _DB_NAME)


def _ensure_summary_tables(conn: sqlite3.Connection) -> None:
    """Create summary tables if missing — handles DBs created before this update."""
    existing = {r[0] for r in conn.execute(
        "SELECT name FROM sqlite_master WHERE type='table'"
    ).fetchall()}
    if "stats_by_llm" not in existing:
        conn.executescript(_CREATE_SUMMARY_SQL)
        conn.commit()


def _get_connection(runs_dir: str) -> sqlite3.Connection:
    """Return a long-lived, thread-safe connection (WAL mode).
    Schema is initialised only on first connect per path."""
    with _conn_lock:
        conn = _conn_cache.get(runs_dir)
        if conn is None:
            conn = sqlite3.connect(
                _db_path(runs_dir),
                check_same_thread=False,
            )
            conn.execute("PRAGMA journal_mode=WAL;")
            conn.execute("PRAGMA synchronous=NORMAL;")
            conn.execute("PRAGMA cache_size=-8000;")   # 8 MB page cache
            conn.executescript(_CREATE_SQL)
            conn.commit()
            _conn_cache[runs_dir] = conn
        else:
            # Ensure summary tables exist for connections predating this schema update
            _ensure_summary_tables(conn)
        return conn


def _connect(runs_dir: str) -> sqlite3.Connection:
    """Backwards-compatible alias — returns the shared connection."""
    return _get_connection(runs_dir)


def _invalidate_connection(runs_dir: str) -> None:
    """Close and remove the cached connection (used after rebuild_index)."""
    with _conn_lock:
        conn = _conn_cache.pop(runs_dir, None)
    if conn:
        try:
            conn.close()
        except Exception:
            pass


def _refresh_summaries(conn: sqlite3.Connection) -> None:
    """Rebuild all three precomputed summary tables from the runs table (pure SQL)."""
    conn.executescript(_REFRESH_SUMMARIES_SQL)


# ── Row parser ────────────────────────────────────────────────────────────────

def _parse_run_row(filename: str, run_json: dict):
    """Parse a run JSON dict into a tuple for DB insertion. Returns None on failure."""
    raw_results = run_json.get("results", [])

    if isinstance(raw_results, str):
        results = []
        model_attractors = None
    elif isinstance(raw_results, dict):
        results = raw_results.get("results", [])
        if not isinstance(results, list):
            results = []
        raw_att = raw_results.get("model_attractors_count")
        model_attractors = float(raw_att) if isinstance(raw_att, (int, float)) else None
    elif isinstance(raw_results, list):
        results = raw_results
        model_attractors = None
    else:
        results = []
        model_attractors = None

    total_phenotypes = len(results)
    verified_count   = sum(1 for item in results
                           if isinstance(item, dict) and item.get("verified") is True)
    success_rate     = (verified_count / total_phenotypes * 100) if total_phenotypes > 0 else 0.0
    coverage: Optional[float] = None
    if model_attractors and model_attractors > 0:
        coverage = round(total_phenotypes / model_attractors * 100, 2)

    tokens_info = run_json.get("tokens_used", {})
    req_tok = res_tok = tot_tok = None
    if isinstance(tokens_info, dict):
        req_tok = tokens_info.get("request_tokens")
        res_tok = tokens_info.get("response_tokens")
        tot_tok = tokens_info.get("total_tokens")

    return (
        filename,
        run_json.get("timestamp", ""),
        run_json.get("llm_used", ""),
        run_json.get("input_group", ""),
        run_json.get("model_group", ""),
        int(run_json.get("milestone", 1)),
        model_attractors,
        total_phenotypes,
        verified_count,
        round(success_rate, 2),
        coverage,
        req_tok,
        res_tok,
        tot_tok,
        run_json.get("prompt", ""),
    )


_INSERT_SQL = """
    INSERT OR REPLACE INTO runs
      (file, timestamp, llm_used, input_group, model_group,
       milestone, model_attractors_count,
       total_phenotypes, verified_count, success_rate, coverage,
       request_tokens, response_tokens, total_tokens, prompt)
    VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
"""


# ── Public API ────────────────────────────────────────────────────────────────

def upsert_run(runs_dir: str, filename: str, run_json: dict) -> None:
    """Insert or replace a single run record and refresh precomputed summaries."""
    row = _parse_run_row(filename, run_json)
    conn = _get_connection(runs_dir)
    conn.execute(_INSERT_SQL, row)
    conn.commit()
    _refresh_summaries(conn)


def load_runs_df(runs_dir: str) -> pd.DataFrame:
    """
    Return a DataFrame of all indexed runs.
    Automatically indexes any new JSON files not yet in the DB.
    """
    _ensure_synced(runs_dir)
    conn = _get_connection(runs_dir)
    return pd.read_sql("SELECT * FROM runs ORDER BY timestamp DESC", conn)


def load_llm_summary_df(runs_dir: str) -> pd.DataFrame:
    """Return precomputed per-LLM aggregate statistics (no Python groupby needed)."""
    _ensure_synced(runs_dir)
    conn = _get_connection(runs_dir)
    return pd.read_sql("SELECT * FROM stats_by_llm ORDER BY total_runs DESC", conn)


def load_llm_model_summary_df(runs_dir: str) -> pd.DataFrame:
    """Return precomputed per-(LLM, model_group) aggregate statistics."""
    _ensure_synced(runs_dir)
    conn = _get_connection(runs_dir)
    return pd.read_sql(
        "SELECT * FROM stats_by_llm_model ORDER BY llm_used, model_group", conn
    )


def load_milestone_summary_df(runs_dir: str) -> pd.DataFrame:
    """Return precomputed per-milestone aggregate statistics."""
    _ensure_synced(runs_dir)
    conn = _get_connection(runs_dir)
    return pd.read_sql("SELECT * FROM stats_by_milestone ORDER BY milestone", conn)


def _count_json_runs(runs_dir: str) -> int:
    """Count run JSON files without loading them."""
    try:
        return sum(
            1 for f in os.listdir(runs_dir)
            if f.startswith("run_") and f.endswith(".json")
        )
    except OSError:
        return 0


def _count_indexed(runs_dir: str) -> int:
    """Count rows already in the DB (fast single-query check)."""
    conn = _get_connection(runs_dir)
    return conn.execute("SELECT COUNT(*) FROM runs").fetchone()[0]


def _ensure_synced(runs_dir: str) -> None:
    """Index any JSON run files that are not yet in the DB, then refresh summaries.

    Fast path: uses a single COUNT comparison; only scans the directory when counts differ.
    Also repopulates summary tables if they are empty but run rows exist (e.g. after
    a schema migration on an existing DB).
    """
    conn = _get_connection(runs_dir)

    json_count    = _count_json_runs(runs_dir)
    db_count      = conn.execute("SELECT COUNT(*) FROM runs").fetchone()[0]
    summary_count = conn.execute("SELECT COUNT(*) FROM stats_by_llm").fetchone()[0]

    if json_count == db_count:
        if summary_count == 0 and db_count > 0:
            # Summary tables are empty (e.g. first start after schema migration) — just rebuild them
            _refresh_summaries(conn)
        return   # nothing new — skip directory scan

    # Find and index only the new files in one batched transaction.
    indexed = {
        row[0]
        for row in conn.execute("SELECT file FROM runs").fetchall()
    }
    rows = []
    for fname in os.listdir(runs_dir):
        if not (fname.startswith("run_") and fname.endswith(".json")):
            continue
        if fname in indexed:
            continue
        fpath = os.path.join(runs_dir, fname)
        try:
            with open(fpath, "r", encoding="utf-8-sig") as f:
                run_json = json.load(f)
            rows.append(_parse_run_row(fname, run_json))
        except Exception:
            pass   # corrupt / partial file — skip silently

    if rows:
        conn.executemany(_INSERT_SQL, rows)
        conn.commit()
        _refresh_summaries(conn)


def rebuild_index(runs_dir: str) -> int:
    """
    Drop and fully rebuild the index from all JSON files.
    Uses a single connection and batched transaction for speed.
    Returns the number of runs indexed.
    """
    _invalidate_connection(runs_dir)

    db = _db_path(runs_dir)
    if os.path.exists(db):
        os.remove(db)

    conn = _get_connection(runs_dir)   # creates fresh DB + schema
    rows = []
    for fname in os.listdir(runs_dir):
        if not (fname.startswith("run_") and fname.endswith(".json")):
            continue
        fpath = os.path.join(runs_dir, fname)
        try:
            with open(fpath, "r", encoding="utf-8-sig") as f:
                run_json = json.load(f)
            rows.append(_parse_run_row(fname, run_json))
        except Exception:
            pass
    if rows:
        conn.executemany(_INSERT_SQL, rows)
        conn.commit()
    _refresh_summaries(conn)
    return len(rows)
