captcha_solver/learning/db.py

"""SQLite learning database for tracking solver accuracy and optimization history."""
from __future__ import annotations

import json
import sqlite3
import time
from pathlib import Path
from typing import Any

DATA_DIR = Path(__file__).resolve().parent.parent.parent / "data"
DATA_DIR.mkdir(parents=True, exist_ok=True)
DB_PATH = DATA_DIR / "learning.db"

SCHEMA = """
CREATE TABLE IF NOT EXISTS attempts (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    captcha_type TEXT NOT NULL,
    solver_used TEXT NOT NULL,
    hint TEXT,
    image_hash TEXT,
    answer TEXT,
    expected TEXT,
    correct INTEGER,
    confidence REAL,
    latency_ms INTEGER,
    preprocess_steps TEXT,
    run_source TEXT DEFAULT 'api',
    timestamp TEXT DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS solver_stats (
    solver_name TEXT NOT NULL,
    captcha_type TEXT NOT NULL,
    total INTEGER DEFAULT 0,
    correct INTEGER DEFAULT 0,
    avg_latency_ms REAL DEFAULT 0,
    last_run TEXT DEFAULT (datetime('now')),
    PRIMARY KEY (solver_name, captcha_type)
);

CREATE TABLE IF NOT EXISTS optimization_log (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    cycle INTEGER,
    action TEXT,
    before_val TEXT,
    after_val TEXT,
    before_acc REAL,
    after_acc REAL,
    notes TEXT,
    timestamp TEXT DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS cycles (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    status TEXT DEFAULT 'running',
    total_solves INTEGER DEFAULT 0,
    correct_solves INTEGER DEFAULT 0,
    avg_latency_ms REAL DEFAULT 0,
    started_at TEXT DEFAULT (datetime('now')),
    finished_at TEXT
);
"""


class LearningDB:
    def __init__(self, db_path: str | Path = DB_PATH) -> None:
        self.db_path = Path(db_path)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._conn: sqlite3.Connection | None = None

    @property
    def conn(self) -> sqlite3.Connection:
        if self._conn is None:
            self._conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
            self._conn.row_factory = sqlite3.Row
            self._conn.executescript(SCHEMA)
            self._conn.commit()
        return self._conn

    def close(self) -> None:
        if self._conn:
            self._conn.close()
            self._conn = None

    # --- Attempts ---
    def record_attempt(
        self,
        captcha_type: str,
        solver_used: str,
        answer: str | None,
        expected: str | None = None,
        correct: bool | None = None,
        confidence: float | None = None,
        latency_ms: int | None = None,
        hint: str | None = None,
        image_hash: str | None = None,
        preprocess_steps: str | None = None,
        run_source: str = "api",
    ) -> int:
        if correct is None and expected is not None and answer is not None:
            correct = answer.strip().lower() == expected.strip().lower()
        cur = self.conn.execute(
            """INSERT INTO attempts
               (captcha_type, solver_used, hint, image_hash, answer, expected, correct, confidence, latency_ms, preprocess_steps, run_source)
               VALUES (?,?,?,?,?,?,?,?,?,?,?)""",
            (captcha_type, solver_used, hint, image_hash, answer, expected,
             int(correct) if correct is not None else None,
             confidence, latency_ms, preprocess_steps, run_source),
        )
        self.conn.commit()
        self._update_stats(solver_used, captcha_type, correct, latency_ms)
        return cur.lastrowid

    def _update_stats(self, solver: str, ctype: str, correct: bool | None, latency_ms: int | None) -> None:
        row = self.conn.execute(
            "SELECT total, correct, avg_latency_ms FROM solver_stats WHERE solver_name=? AND captcha_type=?",
            (solver, ctype)
        ).fetchone()
        if row:
            t = row["total"] + 1
            c = row["correct"] + (1 if correct else 0) if correct is not None else row["correct"]
            avg = row["avg_latency_ms"]
            if latency_ms:
                avg = ((avg * row["total"]) + latency_ms) / t
            self.conn.execute(
                "UPDATE solver_stats SET total=?, correct=?, avg_latency_ms=?, last_run=datetime('now') WHERE solver_name=? AND captcha_type=?",
                (t, c, avg, solver, ctype)
            )
        else:
            self.conn.execute(
                "INSERT INTO solver_stats (solver_name, captcha_type, total, correct, avg_latency_ms) VALUES (?,?,1,?,?)",
                (solver, ctype, (1 if correct else 0) if correct is not None else 0, latency_ms or 0)
            )
        self.conn.commit()

    # --- Stats ---
    def get_best_solver(self, captcha_type: str, min_samples: int = 5) -> dict | None:
        rows = self.conn.execute(
            """SELECT solver_name, total, correct,
                      CAST(correct AS REAL) / MAX(total, 1) AS accuracy,
                      avg_latency_ms
               FROM solver_stats
               WHERE captcha_type=? AND total>=?
               ORDER BY accuracy DESC, total DESC
               LIMIT 1""",
            (captcha_type, min_samples),
        ).fetchall()
        if rows:
            d = dict(rows[0])
            d["accuracy"] = round(d["accuracy"], 3) if d.get("accuracy") is not None else 0.0
            return d
        return None

    def get_solver_ranking(self, captcha_type: str | None = None) -> list[dict]:
        if captcha_type:
            rows = self.conn.execute(
                "SELECT * FROM solver_stats WHERE captcha_type=? ORDER BY CAST(correct AS REAL)/MAX(total,1) DESC",
                (captcha_type,)
            ).fetchall()
        else:
            rows = self.conn.execute(
                "SELECT * FROM solver_stats ORDER BY CAST(correct AS REAL)/MAX(total,1) DESC"
            ).fetchall()
        return [dict(r) for r in rows]

    def get_recent_failures(self, limit: int = 20) -> list[dict]:
        rows = self.conn.execute(
            """SELECT * FROM attempts WHERE correct=0 AND answer IS NOT NULL
               ORDER BY timestamp DESC LIMIT ?""",
            (limit,)
        ).fetchall()
        return [dict(r) for r in rows]

    # --- Cycles ---
    def start_cycle(self) -> int:
        cur = self.conn.execute("INSERT INTO cycles (status) VALUES ('running')")
        self.conn.commit()
        return cur.lastrowid

    def finish_cycle(self, cycle_id: int, total: int, correct: int, avg_latency: float) -> None:
        self.conn.execute(
            "UPDATE cycles SET status='completed', total_solves=?, correct_solves=?, avg_latency_ms=?, finished_at=datetime('now') WHERE id=?",
            (total, correct, avg_latency, cycle_id)
        )
        self.conn.commit()

    def log_optimization(
        self, cycle: int, action: str,
        before_val: str, after_val: str,
        before_acc: float, after_acc: float,
        notes: str = ""
    ) -> int:
        cur = self.conn.execute(
            "INSERT INTO optimization_log (cycle, action, before_val, after_val, before_acc, after_acc, notes) VALUES (?,?,?,?,?,?,?)",
            (cycle, action, before_val, after_val, before_acc, after_acc, notes)
        )
        self.conn.commit()
        return cur.lastrowid

    # --- Export ---
    def summary(self) -> dict:
        totals = self.conn.execute(
            "SELECT COUNT(*) as total, SUM(CASE WHEN correct=1 THEN 1 ELSE 0 END) as correct, AVG(latency_ms) as avg_ms FROM attempts"
        ).fetchone()
        acc = round((totals["correct"] / max(totals["total"], 1)) * 100, 1) if totals["total"] else 0
        return {
            "total_attempts": totals["total"],
            "correct": totals["correct"],
            "accuracy_pct": acc,
            "avg_latency_ms": round(totals["avg_ms"] or 0),
            "solver_count": self.conn.execute("SELECT COUNT(*) FROM solver_stats").fetchone()[0],
            "optimization_cycles": self.conn.execute("SELECT COUNT(*) FROM optimization_log").fetchone()[0],
        }

    def get_recent_attempts(self, limit: int = 50) -> list[dict]:
        rows = self.conn.execute("SELECT * FROM attempts ORDER BY timestamp DESC LIMIT ?", (limit,)).fetchall()
        return [dict(r) for r in rows]