import sqlite3
import json
from datetime import datetime, timezone
from pathlib import Path


class RunLogger:
    """Logs LLM runs to a local SQLite file for comparison and debugging."""

    def __init__(self, path: str = "promptlock_runs.db"):
        self.path = Path(path)
        self._conn = sqlite3.connect(self.path)
        self._create_table()

    def _create_table(self):
        self._conn.execute("""
            CREATE TABLE IF NOT EXISTS runs (
                id          INTEGER PRIMARY KEY AUTOINCREMENT,
                prompt_name TEXT NOT NULL,
                version     TEXT,
                model       TEXT,
                input       TEXT,
                output      TEXT,
                validated   TEXT NOT NULL DEFAULT 'not_checked',
                error       TEXT,
                timestamp   TEXT NOT NULL
            )
        """)
        self._conn.commit()

    def log(
        self,
        prompt_name: str,
        output: str | dict,
        version: str = None,
        model: str = None,
        input: str = None,
        validated: bool = None,
        error: str = None,
    ):
        """Log a single LLM run.

        Args:
            prompt_name: name of the prompt used
            output: the LLM output (string or dict)
            version: prompt version tag
            model: model name (e.g. 'gpt-4o', 'claude-3-5-sonnet')
            input: the rendered prompt sent to the LLM
            validated: True if contract passed, False if failed, None if not checked
            error: the ContractViolation message if validation failed
        """
        if isinstance(output, dict):
            output = json.dumps(output)

        if validated is None:
            validated_str = "not_checked"
        elif validated:
            validated_str = "passed"
        else:
            validated_str = "failed"

        timestamp = datetime.now(timezone.utc).isoformat()

        self._conn.execute(
            """
            INSERT INTO runs (prompt_name, version, model, input, output, validated, error, timestamp)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (prompt_name, version, model, input, output, validated_str, error, timestamp),
        )
        self._conn.commit()

    def get_runs(
        self,
        prompt_name: str = None,
        version: str = None,
        validated: str = None,
        limit: int = 50,
    ) -> list[dict]:
        """Retrieve logged runs with optional filters.

        Args:
            prompt_name: filter by prompt name
            version: filter by version
            validated: filter by 'passed', 'failed', or 'not_checked'
            limit: max number of rows to return

        Returns:
            list of run dicts ordered by most recent first
        """
        query = "SELECT * FROM runs WHERE 1=1"
        params = []

        if prompt_name:
            query += " AND prompt_name = ?"
            params.append(prompt_name)
        if version:
            query += " AND version = ?"
            params.append(version)
        if validated:
            query += " AND validated = ?"
            params.append(validated)

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

        cursor = self._conn.execute(query, params)
        columns = [col[0] for col in cursor.description]
        return [dict(zip(columns, row)) for row in cursor.fetchall()]

    def summary(self, prompt_name: str = None) -> dict:
        """Return a quick summary of run counts grouped by validation result.

        Args:
            prompt_name: optional filter by prompt name
        """
        query = "SELECT validated, COUNT(*) as count FROM runs"
        params = []

        if prompt_name:
            query += " WHERE prompt_name = ?"
            params.append(prompt_name)

        query += " GROUP BY validated"
        cursor = self._conn.execute(query, params)
        return {row[0]: row[1] for row in cursor.fetchall()}

    def close(self):
        """Close the SQLite connection."""
        self._conn.close()

    def __repr__(self) -> str:
        return f"RunLogger(path='{self.path}')"