import csv
import json
import sqlite3
from datetime import datetime, timezone
from pathlib import Path
from typing import Literal


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 export(
        self,
        path: str,
        format: Literal["csv", "json", "excel"] = "csv",
        prompt_name: str = None,
        validated: str = None,
        limit: int = 1000,
    ):
        """Export logged runs to a CSV, JSON, or Excel file.

        Args:
            path: output file path (e.g. 'runs.csv', 'runs.json', 'runs.xlsx')
            format: one of 'csv', 'json', or 'excel'
            prompt_name: optional filter by prompt name
            validated: optional filter by 'passed', 'failed', or 'not_checked'
            limit: max number of rows to export
        """
        runs = self.get_runs(prompt_name=prompt_name, validated=validated, limit=limit)

        if not runs:
            print("No runs found to export.")
            return

        output_path = Path(path)

        if format == "csv":
            with open(output_path, "w", newline="", encoding="utf-8") as f:
                writer = csv.DictWriter(f, fieldnames=runs[0].keys())
                writer.writeheader()
                writer.writerows(runs)

        elif format == "json":
            with open(output_path, "w", encoding="utf-8") as f:
                json.dump(runs, f, indent=2, ensure_ascii=False)

        elif format == "excel":
            from openpyxl import Workbook
            from openpyxl.styles import Font, PatternFill

            wb = Workbook()
            ws = wb.active
            ws.title = "Runs"

            headers = list(runs[0].keys())

            # style the header row
            for col_idx, header in enumerate(headers, start=1):
                cell = ws.cell(row=1, column=col_idx, value=header)
                cell.font = Font(bold=True, color="FFFFFF")
                cell.fill = PatternFill(fill_type="solid", fgColor="2563EB")

            # write data rows
            for row_idx, run in enumerate(runs, start=2):
                for col_idx, key in enumerate(headers, start=1):
                    ws.cell(row=row_idx, column=col_idx, value=run[key])

            # auto size columns
            for col in ws.columns:
                max_length = max(len(str(cell.value or "")) for cell in col)
                ws.column_dimensions[col[0].column_letter].width = min(max_length + 4, 60)

            wb.save(output_path)

        else:
            raise ValueError(f"Unsupported format '{format}'. Use 'csv', 'json', or 'excel'.")

        print(f"Exported {len(runs)} runs to {output_path}")

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

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