from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Dict, Any, List
from pathlib import Path

import duckdb


@dataclass
class ArmStats:
    name: str
    wins: int
    trials: int
    ucb: float


class BanditRepo:
    """Simple DuckDB-backed repo for UCB1 arm statistics.

    This lives alongside DecisionLogger at data/learning.duckdb by default.
    """

    def __init__(self, duckdb_path: str = "data/learning.duckdb"):
        self.duckdb_path = Path(duckdb_path)
        self.duckdb_path.parent.mkdir(parents=True, exist_ok=True)
        self._init_tables()

    def _init_tables(self):
        with duckdb.connect(str(self.duckdb_path)) as conn:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS bandit_arms (
                    account_id TEXT,
                    segment TEXT,
                    arm TEXT,
                    wins INTEGER,
                    trials INTEGER,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    PRIMARY KEY (account_id, segment, arm)
                )
                """
            )

    def get_ucb_snapshot(self, account_id: str, segment: str) -> Dict[str, Any]:
        with duckdb.connect(str(self.duckdb_path)) as conn:
            rows = conn.execute(
                """
                SELECT arm, wins, trials
                FROM bandit_arms
                WHERE account_id = ? AND segment = ?
                """,
                [str(account_id), str(segment)],
            ).fetchall()
            total_trials = sum(r[2] for r in rows) or 1
            arms: List[ArmStats] = []
            for arm, wins, trials in rows:
                if trials <= 0:
                    ucb = float("inf")
                else:
                    rate = (wins or 0) / max(1, trials)
                    ucb = rate + math.sqrt(2 * math.log(total_trials) / trials)
                arms.append(
                    ArmStats(
                        name=arm, wins=int(wins or 0), trials=int(trials or 0), ucb=ucb
                    )
                )
            arms.sort(key=lambda a: a.ucb, reverse=True)
            return {
                "arms": [a.__dict__ for a in arms],
                "updated_at": None,
            }

    def upsert_arm_stats(
        self,
        account_id: str,
        segment: str,
        arm: str,
        *,
        wins_delta: int,
        trials_delta: int,
    ) -> None:
        with duckdb.connect(str(self.duckdb_path)) as conn:
            # Read current
            row = conn.execute(
                """
                SELECT wins, trials FROM bandit_arms
                WHERE account_id = ? AND segment = ? AND arm = ?
                """,
                [str(account_id), str(segment), str(arm)],
            ).fetchone()
            if row:
                new_wins = int(row[0] or 0) + int(wins_delta or 0)
                new_trials = int(row[1] or 0) + int(trials_delta or 0)
                conn.execute(
                    """
                    UPDATE bandit_arms
                    SET wins = ?, trials = ?, updated_at = CURRENT_TIMESTAMP
                    WHERE account_id = ? AND segment = ? AND arm = ?
                    """,
                    [new_wins, new_trials, str(account_id), str(segment), str(arm)],
                )
            else:
                conn.execute(
                    """
                    INSERT INTO bandit_arms (account_id, segment, arm, wins, trials)
                    VALUES (?, ?, ?, ?, ?)
                    """,
                    [
                        str(account_id),
                        str(segment),
                        str(arm),
                        int(wins_delta or 0),
                        int(trials_delta or 0),
                    ],
                )
