"""
Multi-Armed Bandits for Adaptive Configuration Selection
========================================================
Uses UCB1 algorithm for exploration/exploitation balance.
Tracks performance of threshold policies and blocking strategies.
"""

import json
import math
import logging
import time
import os
import asyncio
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, asdict
import redis.asyncio as async_redis

log = logging.getLogger(__name__)


_BANDIT_IO_MS = int(os.getenv("FM_BANDIT_IO_TIMEOUT_MS", "100"))


async def _timeout(coro, default=None):
    try:
        return await asyncio.wait_for(coro, timeout=_BANDIT_IO_MS / 1000.0)
    except Exception:
        return default


@dataclass
class ArmStats:
    """Statistics for a single bandit arm"""

    arm_id: str
    arm_type: str  # 'threshold' or 'blocking'
    trials: int = 0
    successes: float = 0.0
    reward_sum: float = 0.0
    last_selected: float = 0.0

    @property
    def mean_reward(self) -> float:
        return self.reward_sum / self.trials if self.trials > 0 else 0.0

    @property
    def success_rate(self) -> float:
        return self.successes / self.trials if self.trials > 0 else 0.0

    def ucb_score(self, total_trials: int, c: float = 1.0) -> float:
        """Calculate Upper Confidence Bound score"""
        if self.trials == 0:
            return float("inf")  # Unexplored arms have highest priority

        exploitation = self.mean_reward
        exploration = c * math.sqrt(2 * math.log(total_trials) / self.trials)
        return exploitation + exploration


class ThresholdBandit:
    """
    Bandit for selecting optimal threshold policies.
    Arms: strict (0.90), balanced (0.80), recall_oriented (0.70)
    """

    def __init__(self, redis_conn: Optional[async_redis.Redis] = None):
        self.redis_conn = redis_conn
        self.policies = {
            "strict": {
                "threshold": 0.90,
                "description": "High precision, lower recall",
            },
            "balanced": {"threshold": 0.80, "description": "Balance precision/recall"},
            "recall_oriented": {
                "threshold": 0.70,
                "description": "High recall, lower precision",
            },
        }
        self.stats_key = "bandit:threshold:stats"
        self.c = 1.0  # UCB exploration constant

    async def select(self, segment: str) -> Dict[str, Any]:
        """
        Select best threshold policy for given segment using UCB1.

        Args:
            segment: Data segment identifier (e.g., "saas_b2b_high_quality")

        Returns:
            Selected policy with name and threshold
        """
        stats = await self._load_stats(segment)

        # Calculate total trials across all arms
        total_trials = sum(s.trials for s in stats.values())

        # If no trials yet, return default
        if total_trials == 0:
            log.info(f"No trials for segment {segment}, returning balanced policy")
            return {
                "policy": "balanced",
                "threshold": self.policies["balanced"]["threshold"],
                "exploration": True,
            }

        # Calculate UCB scores for each arm
        scores = {}
        for policy_name, arm_stats in stats.items():
            scores[policy_name] = arm_stats.ucb_score(total_trials, self.c)
            log.debug(
                f"Policy {policy_name}: UCB={scores[policy_name]:.3f}, "
                f"mean={arm_stats.mean_reward:.3f}, trials={arm_stats.trials}"
            )

        # Select policy with highest UCB score
        best_policy = max(scores, key=scores.get)

        # Determine if we're exploring (high UCB due to low trials)
        is_exploration = stats[best_policy].trials < 10

        log.info(
            f"Selected {best_policy} for {segment} "
            f"(UCB={scores[best_policy]:.3f}, exploration={is_exploration})"
        )

        return {
            "policy": best_policy,
            "threshold": self.policies[best_policy]["threshold"],
            "exploration": is_exploration,
            "confidence": 1.0
            - (1.0 / (stats[best_policy].trials + 1)),  # Confidence grows with trials
        }

    async def update(self, segment: str, policy: str, reward: float):
        """
        Update statistics after observing outcome.

        Args:
            segment: Data segment identifier
            policy: Policy that was used
            reward: Observed reward (0-1, based on F-score and FP rate)
        """
        stats = await self._load_stats(segment)

        if policy not in stats:
            log.warning(f"Unknown policy {policy}, skipping update")
            return

        # Update arm statistics
        arm = stats[policy]
        arm.trials += 1
        arm.reward_sum += reward
        if reward > 0.5:  # Consider it a "success" if reward > 0.5
            arm.successes += 1
        arm.last_selected = time.time()

        # Save updated stats
        await self._save_stats(segment, stats)

        log.info(
            f"Updated {policy} for {segment}: "
            f"reward={reward:.3f}, mean={arm.mean_reward:.3f}, "
            f"trials={arm.trials}"
        )

    async def _load_stats(self, segment: str) -> Dict[str, ArmStats]:
        """Load statistics from Redis"""
        if not self.redis_conn:
            # Initialize with empty stats
            return {
                policy: ArmStats(arm_id=policy, arm_type="threshold")
                for policy in self.policies
            }

        key = f"{self.stats_key}:{segment}"
        data = await _timeout(self.redis_conn.get(key), None)

        if not data:
            # Initialize stats for new segment
            return {
                policy: ArmStats(arm_id=policy, arm_type="threshold")
                for policy in self.policies
            }

        # Deserialize stats
        raw_stats = json.loads(data)
        return {
            policy: ArmStats(
                **raw_stats.get(policy, {"arm_id": policy, "arm_type": "threshold"})
            )
            for policy in self.policies
        }

    async def _save_stats(self, segment: str, stats: Dict[str, ArmStats]):
        """Save statistics to Redis"""
        if not self.redis_conn:
            return

        key = f"{self.stats_key}:{segment}"
        data = {policy: asdict(arm) for policy, arm in stats.items()}
        await _timeout(
            self.redis_conn.set(key, json.dumps(data)), True
        )  # No TTL - keep stats forever

    async def get_performance_summary(self, segment: str) -> Dict[str, Any]:
        """Get performance summary for monitoring"""
        stats = await self._load_stats(segment)

        return {
            "segment": segment,
            "policies": {
                policy: {
                    "trials": arm.trials,
                    "success_rate": arm.success_rate,
                    "mean_reward": arm.mean_reward,
                    "last_selected": arm.last_selected,
                }
                for policy, arm in stats.items()
            },
            "total_trials": sum(s.trials for s in stats.values()),
        }


class BlockingBandit:
    """
    Bandit for selecting optimal blocking strategies.
    Arms: different blocking columns and key lengths
    """

    def __init__(self, redis_conn: Optional[async_redis.Redis] = None):
        self.redis_conn = redis_conn
        self.stats_key = "bandit:blocking:stats"
        self.c = 1.2  # Slightly higher exploration for blocking (more critical)

        # Common blocking strategies
        self.strategies = {
            "domain_3": {"column": "Domain", "key_length": 3, "variant": "alnum"},
            "domain_4": {"column": "Domain", "key_length": 4, "variant": "alnum"},
            "company_3": {"column": "Company", "key_length": 3, "variant": "alnum"},
            "name_3": {"column": "Name", "key_length": 3, "variant": "alnum"},
            "email_prefix": {
                "column": "Email",
                "key_length": 5,
                "variant": "email_prefix",
            },
            "none": {"column": None, "key_length": 0, "variant": None},  # No blocking
        }

    async def select(
        self, segment: str, available_columns: List[str]
    ) -> Dict[str, Any]:
        """
        Select best blocking strategy for given segment.

        Args:
            segment: Data segment identifier
            available_columns: List of columns available in the data

        Returns:
            Selected blocking strategy
        """
        # Filter strategies to only those with available columns
        valid_strategies = {}
        for name, strategy in self.strategies.items():
            if strategy["column"] is None or strategy["column"] in available_columns:
                valid_strategies[name] = strategy

        if not valid_strategies:
            log.warning(f"No valid blocking strategies for columns {available_columns}")
            return {"apply_blocking": False}

        stats = await self._load_stats(segment)

        # Only consider valid strategies
        valid_stats = {k: v for k, v in stats.items() if k in valid_strategies}

        # Calculate total trials
        total_trials = sum(s.trials for s in valid_stats.values())

        if total_trials == 0:
            # Default to domain blocking if available
            if "domain_3" in valid_strategies:
                return {
                    "apply_blocking": True,
                    "strategy": "domain_3",
                    **valid_strategies["domain_3"],
                    "exploration": True,
                }
            # Otherwise use first available strategy
            first_strategy = next(iter(valid_strategies))
            return {
                "apply_blocking": first_strategy != "none",
                "strategy": first_strategy,
                **valid_strategies[first_strategy],
                "exploration": True,
            }

        # Calculate UCB scores
        scores = {}
        for strategy_name in valid_strategies:
            if strategy_name not in valid_stats:
                # Initialize new arm
                valid_stats[strategy_name] = ArmStats(
                    arm_id=strategy_name, arm_type="blocking"
                )
            scores[strategy_name] = valid_stats[strategy_name].ucb_score(
                total_trials, self.c
            )

        # Select best strategy
        best_strategy = max(scores, key=scores.get)
        is_exploration = valid_stats[best_strategy].trials < 5

        result = {
            "apply_blocking": best_strategy != "none",
            "strategy": best_strategy,
            **valid_strategies[best_strategy],
            "exploration": is_exploration,
            "confidence": 1.0 - (1.0 / (valid_stats[best_strategy].trials + 1)),
        }

        # Clean up None values for "none" strategy
        if best_strategy == "none":
            result.pop("column", None)
            result.pop("key_length", None)
            result.pop("variant", None)

        log.info(
            f"Selected blocking {best_strategy} for {segment} "
            f"(UCB={scores[best_strategy]:.3f}, exploration={is_exploration})"
        )

        return result

    async def update(self, segment: str, strategy: str, reward: float):
        """Update statistics after observing outcome"""
        stats = await self._load_stats(segment)

        if strategy not in self.strategies:
            log.warning(f"Unknown strategy {strategy}, skipping update")
            return

        # Initialize if needed
        if strategy not in stats:
            stats[strategy] = ArmStats(arm_id=strategy, arm_type="blocking")

        # Update arm
        arm = stats[strategy]
        arm.trials += 1
        arm.reward_sum += reward
        if reward > 0.5:
            arm.successes += 1
        arm.last_selected = time.time()

        await self._save_stats(segment, stats)

        log.info(
            f"Updated blocking {strategy} for {segment}: "
            f"reward={reward:.3f}, mean={arm.mean_reward:.3f}, "
            f"trials={arm.trials}"
        )

    async def _load_stats(self, segment: str) -> Dict[str, ArmStats]:
        """Load statistics from Redis"""
        if not self.redis_conn:
            return {}

        key = f"{self.stats_key}:{segment}"
        data = await _timeout(self.redis_conn.get(key), None)

        if not data:
            return {}

        raw_stats = json.loads(data)
        return {name: ArmStats(**stats_dict) for name, stats_dict in raw_stats.items()}

    async def _save_stats(self, segment: str, stats: Dict[str, ArmStats]):
        """Save statistics to Redis"""
        if not self.redis_conn:
            return

        key = f"{self.stats_key}:{segment}"
        data = {name: asdict(arm) for name, arm in stats.items()}
        await _timeout(
            self.redis_conn.set(key, json.dumps(data)), True
        )  # No TTL - keep stats forever

    async def get_performance_summary(self, segment: str) -> Dict[str, Any]:
        """Get performance summary for monitoring"""
        stats = await self._load_stats(segment)

        return {
            "segment": segment,
            "strategies": {
                name: {
                    "trials": arm.trials,
                    "success_rate": arm.success_rate,
                    "mean_reward": arm.mean_reward,
                    "last_selected": arm.last_selected,
                }
                for name, arm in stats.items()
            },
            "total_trials": sum(s.trials for s in stats.values()),
        }
