"""RepairBandit — UCB1 multi-armed bandit for SelfHealer strategy selection.

Reads/writes strategy outcomes to SQLite for persistence across restarts.
Purpose:    Selects the best healing strategy using UCB1 exploration-exploitation. Persists strategy outcomes to SQLite across restarts.
Exports:    RepairBandit
Added:      v0.55.0
Status:     Stable
"""
from __future__ import annotations
import math
import sqlite3
from pathlib import Path
from typing import Optional


class RepairBandit:
    """UCB1 bandit that selects the best SelfHealer repair strategy.

    Args:
        strategies: list of strategy names.
        db_path: SQLite path for persistence. None = in-memory only.
        exploration_c: UCB1 exploration constant.
    """

    def __init__(
        self,
        strategies: list[str],
        db_path: Optional[str] = None,
        exploration_c: float = 1.0,
    ) -> None:
        self.strategies = strategies
        self.exploration_c = exploration_c
        self._counts: dict[str, int] = {s: 0 for s in strategies}
        self._rewards: dict[str, float] = {s: 0.0 for s in strategies}
        self._db_path = db_path
        if db_path:
            self._init_db()
            self._load_from_db()

    def _init_db(self) -> None:
        Path(self._db_path).parent.mkdir(parents=True, exist_ok=True)
        with sqlite3.connect(self._db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS bandit_stats (
                    strategy TEXT PRIMARY KEY,
                    n_pulls  INTEGER DEFAULT 0,
                    total_reward REAL DEFAULT 0.0
                )
            """)

    def _load_from_db(self) -> None:
        with sqlite3.connect(self._db_path) as conn:
            rows = conn.execute(
                "SELECT strategy, n_pulls, total_reward FROM bandit_stats"
            ).fetchall()
        for strategy, n, r in rows:
            if strategy in self._counts:
                self._counts[strategy] = n
                self._rewards[strategy] = r

    def _save_to_db(self, strategy: str) -> None:
        with sqlite3.connect(self._db_path) as conn:
            conn.execute("""
                INSERT INTO bandit_stats (strategy, n_pulls, total_reward)
                VALUES (?, ?, ?)
                ON CONFLICT(strategy) DO UPDATE SET
                    n_pulls = excluded.n_pulls,
                    total_reward = excluded.total_reward
            """, (strategy, self._counts[strategy], self._rewards[strategy]))

    def select(self) -> str:
        """Select a strategy using UCB1. Pulls each arm once before UCB."""
        total = sum(self._counts.values())
        for s in self.strategies:
            if self._counts[s] == 0:
                return s
        best, best_score = None, -1.0
        for s in self.strategies:
            n = self._counts[s]
            mean = self._rewards[s] / n
            ucb = mean + self.exploration_c * math.sqrt(math.log(total) / n)
            if ucb > best_score:
                best_score = ucb
                best = s
        return best

    def update(self, strategy: str, reward: float) -> None:
        """Record outcome. reward=1.0 = success, 0.0 = failure."""
        if strategy not in self._counts:
            return
        self._counts[strategy] += 1
        self._rewards[strategy] += reward
        if self._db_path:
            self._save_to_db(strategy)
