"""SQLite-backed experiment store for tracking training runs."""

from __future__ import annotations

import hashlib
import json
import sqlite3
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path

from pydantic import BaseModel, Field

from bilgeai.core.constants import PROJECT_DIR


class RunRecord(BaseModel):
    """A single training run record."""

    run_id: str
    timestamp: str
    config_hash: str
    model_name: str
    status: str = "created"  # created, running, completed, failed

    # Training config snapshot
    config_json: str = "{}"

    # Hardware snapshot
    hardware_json: str = "{}"

    # Environment snapshot
    env_json: str = "{}"

    # Results (filled after training)
    final_loss: float | None = None
    best_loss: float | None = None
    total_steps: int | None = None
    training_time_seconds: float | None = None
    output_dir: str | None = None

    # Metadata
    notes: str = ""
    tags: list[str] = Field(default_factory=list)


class ExperimentStore:
    """SQLite store for experiment tracking."""

    def __init__(self, project_dir: Path | None = None) -> None:
        self.project_dir = project_dir or (Path.cwd() / PROJECT_DIR)
        self.db_path = self.project_dir / "experiments.db"
        self._ensure_db()

    def _ensure_db(self) -> None:
        """Create DB and tables if they don't exist."""
        self.project_dir.mkdir(parents=True, exist_ok=True)
        conn = sqlite3.connect(str(self.db_path))
        try:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS runs (
                    run_id TEXT PRIMARY KEY,
                    timestamp TEXT NOT NULL,
                    config_hash TEXT NOT NULL,
                    model_name TEXT NOT NULL,
                    status TEXT NOT NULL DEFAULT 'created',
                    config_json TEXT NOT NULL DEFAULT '{}',
                    hardware_json TEXT NOT NULL DEFAULT '{}',
                    env_json TEXT NOT NULL DEFAULT '{}',
                    final_loss REAL,
                    best_loss REAL,
                    total_steps INTEGER,
                    training_time_seconds REAL,
                    output_dir TEXT,
                    notes TEXT DEFAULT '',
                    tags TEXT DEFAULT '[]'
                )
            """)
            conn.commit()
        finally:
            conn.close()

    @staticmethod
    def _generate_run_id() -> str:
        """Generate a short, unique run ID."""
        now = datetime.now(timezone.utc)
        ts = now.strftime("%Y%m%d_%H%M%S")
        # Add a short hash for uniqueness
        h = hashlib.sha256(now.isoformat().encode()).hexdigest()[:6]
        return f"run_{ts}_{h}"

    @staticmethod
    def compute_config_hash(config_dict: dict) -> str:
        """Compute a deterministic hash for a config dict."""
        # Sort keys for determinism, exclude non-essential fields
        clean = {
            k: v
            for k, v in sorted(config_dict.items())
            if k not in ("explanations", "schema_version")
        }
        raw = json.dumps(clean, sort_keys=True, default=str)
        return hashlib.sha256(raw.encode()).hexdigest()[:12]

    @staticmethod
    def capture_environment() -> dict:
        """Capture current environment for reproducibility."""
        env: dict = {
            "python_version": sys.version,
            "platform": sys.platform,
        }

        # pip freeze
        try:
            result = subprocess.run(
                [sys.executable, "-m", "pip", "freeze"],
                capture_output=True,
                text=True,
                timeout=10,
            )
            if result.returncode == 0:
                env["pip_freeze"] = result.stdout.strip().split("\n")
        except Exception:
            pass

        return env

    def create_run(
        self,
        config_dict: dict,
        hardware_dict: dict | None = None,
        model_name: str = "",
        notes: str = "",
        tags: list[str] | None = None,
    ) -> RunRecord:
        """Create a new run record and save to DB."""
        run_id = self._generate_run_id()
        config_hash = self.compute_config_hash(config_dict)
        env_dict = self.capture_environment()

        record = RunRecord(
            run_id=run_id,
            timestamp=datetime.now(timezone.utc).isoformat(),
            config_hash=config_hash,
            model_name=model_name or config_dict.get("model_name", "unknown"),
            status="created",
            config_json=json.dumps(config_dict, default=str),
            hardware_json=json.dumps(hardware_dict or {}, default=str),
            env_json=json.dumps(env_dict, default=str),
            notes=notes,
            tags=tags or [],
        )

        self._insert_run(record)
        return record

    def _insert_run(self, record: RunRecord) -> None:
        """Insert a run record into the DB."""
        conn = sqlite3.connect(str(self.db_path))
        try:
            conn.execute(
                """INSERT INTO runs
                   (run_id, timestamp, config_hash, model_name, status,
                    config_json, hardware_json, env_json,
                    final_loss, best_loss, total_steps, training_time_seconds,
                    output_dir, notes, tags)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    record.run_id,
                    record.timestamp,
                    record.config_hash,
                    record.model_name,
                    record.status,
                    record.config_json,
                    record.hardware_json,
                    record.env_json,
                    record.final_loss,
                    record.best_loss,
                    record.total_steps,
                    record.training_time_seconds,
                    record.output_dir,
                    record.notes,
                    json.dumps(record.tags),
                ),
            )
            conn.commit()
        finally:
            conn.close()

    def update_run(
        self,
        run_id: str,
        status: str | None = None,
        final_loss: float | None = None,
        best_loss: float | None = None,
        total_steps: int | None = None,
        training_time_seconds: float | None = None,
        output_dir: str | None = None,
    ) -> None:
        """Update fields of an existing run."""
        updates = []
        values = []
        if status is not None:
            updates.append("status = ?")
            values.append(status)
        if final_loss is not None:
            updates.append("final_loss = ?")
            values.append(final_loss)
        if best_loss is not None:
            updates.append("best_loss = ?")
            values.append(best_loss)
        if total_steps is not None:
            updates.append("total_steps = ?")
            values.append(total_steps)
        if training_time_seconds is not None:
            updates.append("training_time_seconds = ?")
            values.append(training_time_seconds)
        if output_dir is not None:
            updates.append("output_dir = ?")
            values.append(output_dir)

        if not updates:
            return

        values.append(run_id)
        conn = sqlite3.connect(str(self.db_path))
        try:
            conn.execute(
                f"UPDATE runs SET {', '.join(updates)} WHERE run_id = ?",
                values,
            )
            conn.commit()
        finally:
            conn.close()

    def get_run(self, run_id: str) -> RunRecord | None:
        """Get a single run by ID."""
        conn = sqlite3.connect(str(self.db_path))
        conn.row_factory = sqlite3.Row
        try:
            row = conn.execute("SELECT * FROM runs WHERE run_id = ?", (run_id,)).fetchone()
            if row is None:
                return None
            return self._row_to_record(row)
        finally:
            conn.close()

    def list_runs(
        self,
        limit: int = 20,
        status: str | None = None,
        model_name: str | None = None,
    ) -> list[RunRecord]:
        """List runs, most recent first."""
        conn = sqlite3.connect(str(self.db_path))
        conn.row_factory = sqlite3.Row
        try:
            query = "SELECT * FROM runs"
            params: list = []
            conditions = []

            if status:
                conditions.append("status = ?")
                params.append(status)
            if model_name:
                conditions.append("model_name LIKE ?")
                params.append(f"%{model_name}%")

            if conditions:
                query += " WHERE " + " AND ".join(conditions)

            query += " ORDER BY timestamp DESC LIMIT ?"
            params.append(limit)

            rows = conn.execute(query, params).fetchall()
            return [self._row_to_record(r) for r in rows]
        finally:
            conn.close()

    def get_best_run(self, metric: str = "best_loss") -> RunRecord | None:
        """Get the run with the best metric value."""
        if metric not in ("best_loss", "final_loss"):
            return None

        conn = sqlite3.connect(str(self.db_path))
        conn.row_factory = sqlite3.Row
        try:
            row = conn.execute(
                f"SELECT * FROM runs WHERE {metric} IS NOT NULL ORDER BY {metric} ASC LIMIT 1"
            ).fetchone()
            if row is None:
                return None
            return self._row_to_record(row)
        finally:
            conn.close()

    def diff_runs(self, run_id_a: str, run_id_b: str) -> dict:
        """Compare two runs and return differences."""
        a = self.get_run(run_id_a)
        b = self.get_run(run_id_b)

        if a is None or b is None:
            missing = run_id_a if a is None else run_id_b
            return {"error": f"Run not found: {missing}"}

        config_a = json.loads(a.config_json)
        config_b = json.loads(b.config_json)

        # Find config differences
        all_keys = set(config_a.keys()) | set(config_b.keys())
        config_diff = {}
        for key in sorted(all_keys):
            val_a = config_a.get(key)
            val_b = config_b.get(key)
            if val_a != val_b:
                config_diff[key] = {"run_a": val_a, "run_b": val_b}

        # Results comparison
        results_diff = {}
        for field in ("final_loss", "best_loss", "total_steps", "training_time_seconds"):
            val_a = getattr(a, field)
            val_b = getattr(b, field)
            if val_a != val_b:
                results_diff[field] = {"run_a": val_a, "run_b": val_b}

        return {
            "run_a": {"id": a.run_id, "timestamp": a.timestamp, "status": a.status},
            "run_b": {"id": b.run_id, "timestamp": b.timestamp, "status": b.status},
            "same_config": a.config_hash == b.config_hash,
            "config_diff": config_diff,
            "results_diff": results_diff,
        }

    @staticmethod
    def _row_to_record(row: sqlite3.Row) -> RunRecord:
        """Convert a DB row to RunRecord."""
        tags_raw = row["tags"]
        try:
            tags = json.loads(tags_raw) if tags_raw else []
        except (json.JSONDecodeError, TypeError):
            tags = []

        return RunRecord(
            run_id=row["run_id"],
            timestamp=row["timestamp"],
            config_hash=row["config_hash"],
            model_name=row["model_name"],
            status=row["status"],
            config_json=row["config_json"],
            hardware_json=row["hardware_json"],
            env_json=row["env_json"],
            final_loss=row["final_loss"],
            best_loss=row["best_loss"],
            total_steps=row["total_steps"],
            training_time_seconds=row["training_time_seconds"],
            output_dir=row["output_dir"],
            notes=row["notes"] or "",
            tags=tags,
        )
