"""
Storage layer for PromptVault.
- JSON file  → prompt version history
- SQLite DB  → run logs and analytics
"""

import json
import sqlite3
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional


class PromptStorage:
    """Handles all persistence for PromptVault — no external dependencies."""

    def __init__(self, vault_dir: Path):
        self.vault_dir = Path(vault_dir)
        self.vault_dir.mkdir(parents=True, exist_ok=True)
        self._history_path = self.vault_dir / "history.json"
        self._db_path = self.vault_dir / "runs.db"
        self._init_history()
        self._init_db()

    # ------------------------------------------------------------------
    # Version History (JSON)
    # ------------------------------------------------------------------

    def _init_history(self):
        if not self._history_path.exists():
            self._history_path.write_text(json.dumps({}))

    def _load_history(self) -> dict:
        try:
            return json.loads(self._history_path.read_text())
        except (json.JSONDecodeError, FileNotFoundError):
            return {}

    def _save_history(self, data: dict):
        self._history_path.write_text(json.dumps(data, indent=2))

    def save(self, prompt_dict: Dict[str, Any]):
        """Save a prompt version to history."""
        data = self._load_history()
        name = prompt_dict["name"]
        version = prompt_dict["version"]

        if name not in data:
            data[name] = []

        # Don't duplicate same version
        existing = {v["version"] for v in data[name]}
        if version not in existing:
            data[name].append(prompt_dict)

        self._save_history(data)

    def get_versions(self, name: str) -> List[Dict[str, Any]]:
        """Return all stored versions for a prompt name."""
        data = self._load_history()
        return data.get(name, [])

    def list_prompts(self) -> List[str]:
        """Return all prompt names stored in the vault."""
        return list(self._load_history().keys())

    # ------------------------------------------------------------------
    # Run Logs (SQLite)
    # ------------------------------------------------------------------

    def _connect(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self._db_path)
        conn.row_factory = sqlite3.Row
        return conn

    def _init_db(self):
        with self._connect() as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS runs (
                    run_id          TEXT PRIMARY KEY,
                    prompt_name     TEXT NOT NULL,
                    prompt_version  TEXT NOT NULL,
                    rendered_prompt TEXT,
                    response        TEXT,
                    model           TEXT,
                    latency_ms      REAL,
                    tokens          INTEGER,
                    timestamp       TEXT,
                    extra           TEXT DEFAULT '{}'
                )
            """)
            conn.commit()

    def log_run(
        self,
        prompt_name: str,
        prompt_version: str,
        rendered_prompt: str,
        response: str,
        model: Optional[str] = None,
        latency_ms: Optional[float] = None,
        tokens: Optional[int] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> str:
        """Log a single prompt/response run. Returns the run_id."""
        run_id = str(uuid.uuid4())[:12]
        with self._connect() as conn:
            conn.execute("""
                INSERT INTO runs (
                    run_id, prompt_name, prompt_version,
                    rendered_prompt, response, model,
                    latency_ms, tokens, timestamp, extra
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                run_id,
                prompt_name,
                prompt_version,
                rendered_prompt,
                response,
                model,
                latency_ms,
                tokens,
                datetime.utcnow().isoformat(),
                json.dumps(extra or {}),
            ))
            conn.commit()
        return run_id

    def get_runs(
        self,
        prompt_name: str,
        prompt_version: Optional[str] = None,
        last_n: int = 50,
    ) -> List[Dict[str, Any]]:
        """Retrieve recent runs for a prompt."""
        query = "SELECT * FROM runs WHERE prompt_name = ?"
        params: list = [prompt_name]
        if prompt_version:
            query += " AND prompt_version = ?"
            params.append(prompt_version)
        query += " ORDER BY timestamp DESC LIMIT ?"
        params.append(last_n)

        with self._connect() as conn:
            rows = conn.execute(query, params).fetchall()

        results = []
        for row in rows:
            d = dict(row)
            d["extra"] = json.loads(d.get("extra", "{}"))
            results.append(d)
        return results

    def get_stats(
        self,
        prompt_name: str,
        prompt_version: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Compute aggregate stats for a prompt's runs."""
        runs = self.get_runs(prompt_name, prompt_version, last_n=10000)

        if not runs:
            return {
                "prompt_name": prompt_name,
                "prompt_version": prompt_version or "all",
                "run_count": 0,
                "message": "No runs logged yet.",
            }

        latencies = [r["latency_ms"] for r in runs if r["latency_ms"] is not None]
        tokens    = [r["tokens"]     for r in runs if r["tokens"]     is not None]
        models    = list({r["model"] for r in runs if r["model"]})

        def avg(lst):
            return round(sum(lst) / len(lst), 2) if lst else None

        return {
            "prompt_name":      prompt_name,
            "prompt_version":   prompt_version or "all",
            "run_count":        len(runs),
            "avg_latency_ms":   avg(latencies),
            "min_latency_ms":   round(min(latencies), 2) if latencies else None,
            "max_latency_ms":   round(max(latencies), 2) if latencies else None,
            "avg_tokens":       avg(tokens),
            "total_tokens":     sum(tokens) if tokens else None,
            "models_used":      models,
            "first_run":        runs[-1]["timestamp"],
            "last_run":         runs[0]["timestamp"],
        }
