"""Bandit (op × field) arm 採樣 — UCB / Thompson 雙模式。

State 持久化於 ~/.wq_gui_bandit.json（chmod 600）。GUI 端讀寫；
arm key = 'op_category::field_category'，例如 'ts_rank::anl4'。
"""

from __future__ import annotations

import json
import math
import os
import random as _r
import re
import threading
from typing import Dict

BANDIT_PATH = os.path.expanduser("~/.wq_gui_bandit.json")
_lock = threading.Lock()


def expr_op_category(expr: str) -> str:
    """從 alpha 表達式抽出最外層 op 名稱（'(' 之前）。"""
    s = expr.strip()
    if "(" not in s:
        return "plain"
    return s.split("(", 1)[0]


def expr_field_category(expr: str) -> str:
    """抽第一個 field 的 prefix（'_' 之前），如 anl4_qf_eps → 'anl4'。"""
    m = re.search(r"ts_backfill\(([a-zA-Z0-9_]+)", expr)
    if not m:
        m = re.search(r"\b([a-z][a-z0-9_]+)", expr)
    if not m:
        return "unknown"
    field = m.group(1)
    if "_" in field:
        return field.split("_", 1)[0]
    return field[:5]


def arm_key(expr: str) -> str:
    return f"{expr_op_category(expr)}::{expr_field_category(expr)}"


def load_weights() -> Dict[str, dict]:
    if not os.path.exists(BANDIT_PATH):
        return {}
    try:
        with open(BANDIT_PATH, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return {}


def save_weights(weights: Dict[str, dict]) -> None:
    try:
        with open(BANDIT_PATH, "w", encoding="utf-8") as f:
            json.dump(weights, f, indent=2)
        os.chmod(BANDIT_PATH, 0o600)
    except Exception:
        pass


def _get_total_score(arm: dict) -> float:
    """Read total_score, falling back to legacy wins key."""
    return float(arm.get("total_score", arm.get("wins", 0)))


def score_ucb(arm: dict, total_trials: int, explore_scale: float = 1.0) -> float:
    t = arm.get("trials", 0)
    if t == 0:
        return float("inf")
    win_rate = _get_total_score(arm) / t
    explore = math.sqrt(2.0 * math.log(max(2, total_trials)) / t) * explore_scale
    return win_rate + explore


def score_thompson(arm: dict) -> float:
    total_score = _get_total_score(arm)
    trials = arm.get("trials", 0)
    alpha_p = total_score + 1.0
    beta_p = max(0.1, trials - total_score) + 1.0
    return _r.betavariate(alpha_p, beta_p)


def record(expr: str, score: float, *, is_pass: bool | None = None) -> None:
    """單筆觀測。thread-safe atomic load+modify+save。

    score: continuous [0,1] composite quality score.
    is_pass: deprecated kwarg — if provided, overrides score with 1.0 or 0.0.
    """
    if is_pass is not None:
        score = 1.0 if is_pass else 0.0
    with _lock:
        weights = load_weights()
        key = arm_key(expr)
        arm = weights.setdefault(key, {"trials": 0, "total_score": 0.0})
        arm["trials"] = arm.get("trials", 0) + 1
        arm["total_score"] = _get_total_score(arm) + score
        save_weights(weights)
