#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""Contextual Bandit for Source Selection (Thompson Sampling).

Implements Gap 2 from QR8: Contextual Bandit for Source Selection.
Instead of fixed weights or regex heuristics, the brain learns which of the
18 search sources are historically useful for different task domains using
online reinforcement learning (Thompson Sampling / Beta Distributions).
"""
import logging
import random
from typing import List
import sqlite3

from pulse.core.brain_db import get_conn

log = logging.getLogger("pulse.contextual_bandit")

DEFAULT_ALPHA = 1.0
DEFAULT_BETA = 1.0

def _ensure_table():
    conn = get_conn()
    conn.execute("""
        CREATE TABLE IF NOT EXISTS contextual_bandit (
            domain TEXT NOT NULL,
            source TEXT NOT NULL,
            alpha REAL NOT NULL DEFAULT 1.0,
            beta REAL NOT NULL DEFAULT 1.0,
            PRIMARY KEY (domain, source)
        )
    """)

def update_source_reward(domain: str, retrieved_sources: List[str], success: bool):
    """Reward sources that were retrieved during a successful task."""
    try:
        _ensure_table()
        conn = get_conn()
        domain = domain.lower().strip() if domain else "general"
        
        if not retrieved_sources:
            return
            
        for source in retrieved_sources:
            if not source:
                continue
                
            row = conn.execute("SELECT alpha, beta FROM contextual_bandit WHERE domain=? AND source=?", (domain, source)).fetchone()
            alpha = row[0] if row else DEFAULT_ALPHA
            beta = row[1] if row else DEFAULT_BETA
            
            if success:
                alpha += 1.0
            else:
                beta += 1.0
                
            conn.execute(
                "INSERT OR REPLACE INTO contextual_bandit (domain, source, alpha, beta) VALUES (?, ?, ?, ?)",
                (domain, source, alpha, beta)
            )
            
        log.info(f"Updated Thompson Sampling distributions for domain '{domain}' (Success={success})")
    except Exception as e:
        log.error(f"Failed to update source reward: {e}")

def select_top_sources(domain: str, available_sources: List[str], top_k: int) -> List[str]:
    """Use Thompson Sampling to select the top_k sources for a query in the given domain."""
    try:
        _ensure_table()
        conn = get_conn()
        domain = domain.lower().strip() if domain else "general"
        
        if not available_sources:
            return []
            
        if len(available_sources) <= top_k:
            return available_sources
            
        samples = []
        for source in available_sources:
            row = conn.execute("SELECT alpha, beta FROM contextual_bandit WHERE domain=? AND source=?", (domain, source)).fetchone()
            alpha = row[0] if row else DEFAULT_ALPHA
            beta = row[1] if row else DEFAULT_BETA
            
            try:
                sample = random.betavariate(alpha, beta)
            except ValueError:
                sample = 0.5
                
            samples.append((source, sample))
            
        samples.sort(key=lambda x: x[1], reverse=True)
        return [s[0] for s in samples[:top_k]]
    except Exception as e:
        log.error(f"Failed to select top sources: {e}")
        return available_sources[:top_k]
