"""Exploit pattern database for adversarial verifier testing results.

Store, query, and analyze results from running adversarial tests
against verifiers at scale. This is the data moat.
"""

import json
import sqlite3
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from typing import Literal

from pydantic import BaseModel, Field

DB_PATH = Path.home() / '.deepgym' / 'data' / 'exploit_results.db'


class ExploitRecord(BaseModel):
    """Record of one adversarial test against a verifier."""

    verifier_id: str
    benchmark: str
    exploitable: bool
    patterns: list[str] = Field(default_factory=list)
    num_attacks: int = 0
    num_exploits: int = 0
    max_exploit_score: float = 0.0
    attacks: list[dict] = Field(default_factory=list)
    verifier_hash: str = ''
    attack_type: Literal['white-box', 'black-box'] = 'white-box'
    tested_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat())


class ExploitDB:
    """SQLite database for exploit test results."""

    def __init__(self, db_path: Path | None = None) -> None:
        self._path = db_path or 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._create_tables()

    def _create_tables(self) -> None:
        """Create the results table if it does not exist, and migrate old schemas."""
        self._conn.execute("""CREATE TABLE IF NOT EXISTS results (
            verifier_id TEXT NOT NULL,
            benchmark TEXT NOT NULL,
            exploitable BOOLEAN NOT NULL,
            patterns TEXT NOT NULL DEFAULT '[]',
            num_attacks INTEGER NOT NULL DEFAULT 0,
            num_exploits INTEGER NOT NULL DEFAULT 0,
            max_exploit_score REAL NOT NULL DEFAULT 0.0,
            attacks TEXT NOT NULL DEFAULT '[]',
            verifier_hash TEXT NOT NULL DEFAULT '',
            attack_type TEXT NOT NULL DEFAULT 'white-box',
            tested_at TEXT NOT NULL,
            PRIMARY KEY (verifier_id, attack_type)
        )""")
        self._conn.commit()
        self._migrate()

    def _migrate(self) -> None:
        """Add the attack_type column to old databases that lack it."""
        columns = {row[1] for row in self._conn.execute('PRAGMA table_info(results)').fetchall()}
        if 'attack_type' not in columns:
            self._conn.execute(
                "ALTER TABLE results ADD COLUMN attack_type TEXT NOT NULL DEFAULT 'white-box'"
            )
            self._conn.commit()

    def save(self, record: ExploitRecord) -> None:
        """Save or update an exploit test result."""
        self._conn.execute(
            """INSERT OR REPLACE INTO results
               (verifier_id, benchmark, exploitable, patterns, num_attacks,
                num_exploits, max_exploit_score, attacks, verifier_hash,
                attack_type, tested_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                record.verifier_id,
                record.benchmark,
                record.exploitable,
                json.dumps(record.patterns),
                record.num_attacks,
                record.num_exploits,
                record.max_exploit_score,
                json.dumps(record.attacks),
                record.verifier_hash,
                record.attack_type,
                record.tested_at,
            ),
        )
        self._conn.commit()

    def get(
        self,
        verifier_id: str,
        attack_type: str = 'white-box',
    ) -> ExploitRecord | None:
        """Get a single result by verifier ID and attack type.

        Args:
            verifier_id: The unique identifier for the verifier.
            attack_type: 'white-box' or 'black-box'.

        Returns:
            The matching ExploitRecord, or None if not found.
        """
        row = self._conn.execute(
            'SELECT * FROM results WHERE verifier_id = ? AND attack_type = ?',
            (verifier_id, attack_type),
        ).fetchone()
        if row is None:
            return None
        return self._row_to_record(row)

    def get_all(
        self,
        benchmark: str | None = None,
        attack_type: str | None = None,
    ) -> list[ExploitRecord]:
        """Get all results, optionally filtered by benchmark and attack type.

        Args:
            benchmark: If provided, filter to this benchmark only.
            attack_type: If provided, filter to 'white-box' or 'black-box'.

        Returns:
            List of matching ExploitRecord objects.
        """
        clauses: list[str] = []
        params: list[str] = []
        if benchmark is not None:
            clauses.append('benchmark = ?')
            params.append(benchmark)
        if attack_type is not None:
            clauses.append('attack_type = ?')
            params.append(attack_type)

        query = 'SELECT * FROM results'
        if clauses:
            query += ' WHERE ' + ' AND '.join(clauses)
        query += ' ORDER BY verifier_id'

        rows = self._conn.execute(query, params).fetchall()
        return [self._row_to_record(r) for r in rows]

    def stats(self) -> dict:
        """Get aggregate statistics across all tested verifiers.

        Returns:
            Dict with total_tested, total_exploitable, exploit_rate,
            pattern_counts, and by_benchmark breakdown.
        """
        rows = self.get_all()
        total = len(rows)
        exploitable = sum(1 for r in rows if r.exploitable)

        # Pattern frequency
        pattern_counts = dict(self.pattern_frequency())

        # Per-benchmark breakdown
        by_benchmark: dict[str, dict] = {}
        for r in rows:
            bm = by_benchmark.setdefault(r.benchmark, {'tested': 0, 'exploitable': 0})
            bm['tested'] += 1
            if r.exploitable:
                bm['exploitable'] += 1

        for bm_stats in by_benchmark.values():
            tested = bm_stats['tested']
            bm_stats['exploit_rate'] = bm_stats['exploitable'] / tested if tested > 0 else 0.0

        return {
            'total_tested': total,
            'total_exploitable': exploitable,
            'exploit_rate': exploitable / total if total > 0 else 0.0,
            'pattern_counts': pattern_counts,
            'by_benchmark': by_benchmark,
        }

    def pattern_frequency(self) -> dict[str, int]:
        """Get frequency of each exploit pattern across all results.

        Returns:
            Dict mapping pattern name to occurrence count, sorted descending.
        """
        rows = self._conn.execute('SELECT patterns FROM results').fetchall()
        counter: Counter[str] = Counter()
        for row in rows:
            patterns = json.loads(row['patterns'])
            counter.update(patterns)
        return dict(counter.most_common())

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

    @staticmethod
    def _row_to_record(row: sqlite3.Row) -> ExploitRecord:
        """Convert a database row to an ExploitRecord."""
        return ExploitRecord(
            verifier_id=row['verifier_id'],
            benchmark=row['benchmark'],
            exploitable=bool(row['exploitable']),
            patterns=json.loads(row['patterns']),
            num_attacks=row['num_attacks'],
            num_exploits=row['num_exploits'],
            max_exploit_score=row['max_exploit_score'],
            attacks=json.loads(row['attacks']),
            verifier_hash=row['verifier_hash'],
            attack_type=row['attack_type'],
            tested_at=row['tested_at'],
        )
