"""SQLite persistence layer for learning engine data.

Stores metric runs and experiments in a local SQLite database.
Stdlib-only (sqlite3 is part of Python stdlib).
"""

from __future__ import annotations

import json
import sqlite3
from pathlib import Path

from orxaq.intelligence.learning.experiment import Experiment
from orxaq.intelligence.learning.metrics import DiscoveryMetrics

_DEFAULT_DB_DIR = Path.home() / ".orxaq"
_DEFAULT_DB_PATH = _DEFAULT_DB_DIR / "learning.db"

_CREATE_METRICS = """
CREATE TABLE IF NOT EXISTS metrics_runs (
    run_id          TEXT PRIMARY KEY,
    timestamp       REAL NOT NULL,
    algorithm       TEXT NOT NULL DEFAULT '',
    dataset_size    INTEGER NOT NULL DEFAULT 0,
    num_variables   INTEGER NOT NULL DEFAULT 0,
    shd             INTEGER NOT NULL DEFAULT 0,
    edge_precision  REAL NOT NULL DEFAULT 0.0,
    edge_recall     REAL NOT NULL DEFAULT 0.0,
    edge_f1         REAL NOT NULL DEFAULT 0.0,
    orientation_accuracy REAL NOT NULL DEFAULT 0.0,
    execution_time_ms REAL NOT NULL DEFAULT 0.0,
    config          TEXT NOT NULL DEFAULT '{}'
)
"""

_CREATE_EXPERIMENTS = """
CREATE TABLE IF NOT EXISTS experiments (
    experiment_id   TEXT PRIMARY KEY,
    name            TEXT NOT NULL,
    algorithm       TEXT NOT NULL,
    algorithm_config TEXT NOT NULL DEFAULT '{}',
    data_config     TEXT NOT NULL DEFAULT '{}',
    results         TEXT NOT NULL DEFAULT '{}',
    start_time      REAL NOT NULL DEFAULT 0.0,
    end_time        REAL NOT NULL DEFAULT 0.0,
    status          TEXT NOT NULL DEFAULT 'pending',
    tags            TEXT NOT NULL DEFAULT '[]'
)
"""


class LearningDB:
    """SQLite storage for learning engine data."""

    def __init__(self, db_path: str | Path | None = None) -> None:
        if db_path is None:
            _DEFAULT_DB_DIR.mkdir(parents=True, exist_ok=True)
            self._path = _DEFAULT_DB_PATH
        else:
            self._path = Path(db_path)
            self._path.parent.mkdir(parents=True, exist_ok=True)

        self._conn = sqlite3.connect(str(self._path))
        self._conn.row_factory = sqlite3.Row
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._create_tables()

    def _create_tables(self) -> None:
        """Create tables if they do not exist."""
        self._conn.execute(_CREATE_METRICS)
        self._conn.execute(_CREATE_EXPERIMENTS)
        self._conn.commit()

    # -- metrics ----------------------------------------------------------

    def save_metric(self, metrics: DiscoveryMetrics) -> None:
        """Insert a metric run (or replace if run_id exists)."""
        self._conn.execute(
            """INSERT OR REPLACE INTO metrics_runs
               (run_id, timestamp, algorithm, dataset_size, num_variables,
                shd, edge_precision, edge_recall, edge_f1,
                orientation_accuracy, execution_time_ms, config)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                metrics.run_id,
                metrics.timestamp,
                metrics.algorithm,
                metrics.dataset_size,
                metrics.num_variables,
                metrics.shd,
                metrics.edge_precision,
                metrics.edge_recall,
                metrics.edge_f1,
                metrics.orientation_accuracy,
                metrics.execution_time_ms,
                json.dumps(metrics.config),
            ),
        )
        self._conn.commit()

    def load_metrics(self, algorithm: str | None = None, limit: int = 100) -> list[dict]:
        """Query metric runs, optionally filtered by algorithm."""
        if algorithm:
            cursor = self._conn.execute(
                "SELECT * FROM metrics_runs WHERE algorithm = ? ORDER BY timestamp DESC LIMIT ?",
                (algorithm, limit),
            )
        else:
            cursor = self._conn.execute(
                "SELECT * FROM metrics_runs ORDER BY timestamp DESC LIMIT ?",
                (limit,),
            )
        rows = cursor.fetchall()
        return [self._row_to_metric_dict(r) for r in rows]

    def load_metric_by_id(self, run_id: str) -> dict | None:
        """Load a single metric run by its run_id."""
        cursor = self._conn.execute("SELECT * FROM metrics_runs WHERE run_id = ?", (run_id,))
        row = cursor.fetchone()
        if row is None:
            return None
        return self._row_to_metric_dict(row)

    @staticmethod
    def _row_to_metric_dict(row: sqlite3.Row) -> dict:
        """Convert a sqlite3.Row from metrics_runs to a plain dict."""
        d = dict(row)
        d["config"] = json.loads(d.get("config", "{}"))
        return d

    # -- experiments ------------------------------------------------------

    def save_experiment(self, experiment: Experiment) -> None:
        """Upsert an experiment."""
        self._conn.execute(
            """INSERT OR REPLACE INTO experiments
               (experiment_id, name, algorithm, algorithm_config,
                data_config, results, start_time, end_time, status, tags)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                experiment.experiment_id,
                experiment.name,
                experiment.algorithm,
                json.dumps(experiment.algorithm_config),
                json.dumps(experiment.data_config),
                json.dumps(experiment.results),
                experiment.start_time,
                experiment.end_time,
                experiment.status,
                json.dumps(experiment.tags),
            ),
        )
        self._conn.commit()

    def load_experiments(self, status: str | None = None, limit: int = 100) -> list[dict]:
        """Query experiments, optionally filtered by status."""
        if status:
            cursor = self._conn.execute(
                "SELECT * FROM experiments WHERE status = ? ORDER BY start_time DESC LIMIT ?",
                (status, limit),
            )
        else:
            cursor = self._conn.execute(
                "SELECT * FROM experiments ORDER BY start_time DESC LIMIT ?",
                (limit,),
            )
        rows = cursor.fetchall()
        return [self._row_to_experiment_dict(r) for r in rows]

    @staticmethod
    def _row_to_experiment_dict(row: sqlite3.Row) -> dict:
        """Convert a sqlite3.Row from experiments to a plain dict."""
        d = dict(row)
        for json_field in ("algorithm_config", "data_config", "results", "tags"):
            d[json_field] = json.loads(d.get(json_field, "{}"))
        return d

    # -- summary ----------------------------------------------------------

    def summary(self) -> dict:
        """Return aggregate stats across metrics and experiments."""
        # Metrics summary
        m_row = self._conn.execute(
            "SELECT COUNT(*) as cnt, "
            "MIN(shd) as best_shd, "
            "MAX(edge_f1) as best_f1, "
            "AVG(execution_time_ms) as avg_time "
            "FROM metrics_runs"
        ).fetchone()

        metric_count = m_row["cnt"] if m_row else 0
        algo_rows = self._conn.execute(
            "SELECT DISTINCT algorithm FROM metrics_runs ORDER BY algorithm"
        ).fetchall()
        algorithms = [r["algorithm"] for r in algo_rows]

        # Experiments summary
        e_row = self._conn.execute("SELECT COUNT(*) as cnt FROM experiments").fetchone()
        experiment_count = e_row["cnt"] if e_row else 0

        status_rows = self._conn.execute(
            "SELECT status, COUNT(*) as cnt FROM experiments GROUP BY status"
        ).fetchall()
        by_status = {r["status"]: r["cnt"] for r in status_rows}

        return {
            "metric_runs": metric_count,
            "algorithms": algorithms,
            "best_shd": m_row["best_shd"] if metric_count > 0 else None,
            "best_f1": (round(m_row["best_f1"], 4) if metric_count > 0 else None),
            "avg_execution_ms": (round(m_row["avg_time"], 2) if metric_count > 0 else None),
            "experiments": experiment_count,
            "experiments_by_status": by_status,
        }

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