"""Thompson Sampling bandit algorithm for strategy selection.

Uses Beta distribution to model uncertainty and balance
exploration (trying new strategies) with exploitation
(using what works).

Supports:
- Temporal decay for concept drift adaptation
- Contextual banditing (per error_type + message_category)
- Weighted reward shaping (modified fixes = partial credit)
- Atomic file persistence with backup rotation
"""

from __future__ import annotations

import json
import logging
import math
import os
import random
import re
import tempfile
import threading
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    from autodebug.strategies.strategy_registry import FixStrategy

logger = logging.getLogger(__name__)

# Reward weights for different outcome types
OUTCOME_REWARDS: Dict[str, float] = {
    "accepted": 1.0,
    "applied_worked": 1.0,
    "modified": 0.7,
    "rejected": 0.0,
    "applied_failed": 0.0,
}


@dataclass
class StrategyStats:
    """Track success/failure counts for a strategy.

    Uses Beta(successes+1, failures+1) as the posterior distribution.
    The +1 is a uniform prior (Beta(1,1) = Uniform(0,1)).
    """

    successes: int = 0
    failures: int = 0
    total_reward: float = 0.0  # Weighted reward sum for reward shaping
    last_updated: Optional[str] = None  # ISO timestamp for decay

    @property
    def total(self) -> int:
        """Total number of trials."""
        return self.successes + self.failures

    @property
    def success_rate(self) -> float:
        """Bayesian success rate with Beta(1,1) uniform prior."""
        if self.total == 0:
            return 0.5  # Prior mean
        return (self.successes + 1) / (self.total + 2)

    @property
    def weighted_success_rate(self) -> float:
        """Weighted success rate using reward shaping."""
        if self.total == 0:
            return 0.5
        return self.total_reward / self.total

    @property
    def confidence(self) -> float:
        """Confidence in the success rate estimate (0 to 1).

        Higher values indicate more data and less uncertainty.
        """
        if self.total == 0:
            return 0.0
        # Approximate confidence based on sample size
        # Saturates around 100 trials
        return min(1.0, self.total / 100)

    def sample(self) -> float:
        """Sample from Beta(successes+1, failures+1) distribution.

        Uses the gamma function relationship:
        Beta(a,b) = Gamma(a) / (Gamma(a) + Gamma(b))

        Returns:
            Sampled value between 0 and 1
        """
        alpha = self.successes + 1
        beta = self.failures + 1

        # Sample using gamma variates
        x = random.gammavariate(alpha, 1)
        y = random.gammavariate(beta, 1)
        if x + y == 0:
            return 0.5  # Fallback for extreme floating-point underflow
        return x / (x + y)

    def ucb_score(self, total_trials: int, c: float = 2.0) -> float:
        """Upper Confidence Bound score (alternative to Thompson).

        Args:
            total_trials: Total trials across all strategies
            c: Exploration constant

        Returns:
            UCB score
        """
        if self.total == 0:
            return float("inf")  # Always try untried strategies

        exploitation = self.success_rate
        exploration = c * math.sqrt(math.log(total_trials) / self.total)
        return exploitation + exploration

    def apply_decay(self, factor: float) -> None:
        """Apply decay factor to stats for concept drift adaptation.

        Args:
            factor: Decay multiplier (0-1). E.g., 0.95 reduces counts by 5%.
        """
        self.successes = max(0, round(self.successes * factor))
        self.failures = max(0, round(self.failures * factor))
        self.total_reward *= factor

    def to_dict(self) -> Dict[str, Any]:
        """Serialize to dict."""
        return {
            "successes": self.successes,
            "failures": self.failures,
            "total_reward": self.total_reward,
            "last_updated": self.last_updated,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> StrategyStats:
        """Deserialize from dict."""
        return cls(
            successes=data.get("successes", 0),
            failures=data.get("failures", 0),
            total_reward=data.get("total_reward", 0.0),
            last_updated=data.get("last_updated"),
        )


# Message category patterns for contextual banditing
MESSAGE_CATEGORIES: Dict[str, List[str]] = {
    "missing_key": [r"key\s*error", r"key\s*not\s*found", r"missing\s*key", r"KeyError"],
    "none_type": [r"NoneType", r"None\s*has\s*no", r"is\s*None", r"null"],
    "type_mismatch": [r"expected\s*\w+\s*got", r"cannot\s*convert", r"invalid\s*type"],
    "index_bounds": [r"index\s*out\s*of\s*range", r"IndexError", r"out\s*of\s*bounds"],
    "timeout": [r"timeout", r"timed?\s*out", r"deadline\s*exceeded"],
    "connection": [r"connection\s*(refused|reset|error)", r"network", r"ECONNREFUSED"],
    "permission": [r"permission\s*denied", r"403\s*forbidden", r"unauthorized"],
    "assertion": [r"assert\w*\s*error", r"!=", r"not\s*equal"],
    "import": [r"import\s*error", r"module\s*not\s*found", r"no\s*module"],
    "file_not_found": [r"file\s*not\s*found", r"no\s*such\s*file", r"FileNotFoundError"],
}


def categorize_message(error_type: str, message: str) -> str:
    """Extract a message category for contextual banditing.

    Args:
        error_type: Exception type
        message: Exception message

    Returns:
        Category string (e.g., "missing_key", "none_type") or "general"
    """
    text = f"{error_type} {message}".lower()

    for category, patterns in MESSAGE_CATEGORIES.items():
        for pattern in patterns:
            if re.search(pattern, text, re.IGNORECASE):
                return category

    return "general"


class ThompsonSamplingSelector:
    """Strategy selector using Thompson Sampling bandit algorithm.

    Balances exploration and exploitation by sampling from the
    posterior distribution of each strategy's success probability.

    Features:
    - Temporal decay for concept drift adaptation
    - Atomic file persistence with backup rotation
    - Weighted reward shaping
    """

    def __init__(
        self,
        persistence_path: Optional[str] = None,
        decay_half_life_days: float = 30.0,
        max_backups: int = 3,
    ):
        """Initialize the selector.

        Args:
            persistence_path: Optional path to persist stats across runs
            decay_half_life_days: Half-life for temporal decay (0 = no decay)
            max_backups: Number of backup files to keep
        """
        self.persistence_path = persistence_path
        self.decay_half_life_days = decay_half_life_days
        self.max_backups = max_backups
        self.stats: Dict[str, StrategyStats] = {}
        self._last_decay_check: Optional[str] = None
        self._lock = threading.Lock()
        self._load_stats()

    def _key(self, error_type: str, strategy: str) -> str:
        """Create a unique key for (error_type, strategy) pair."""
        return f"{error_type}:{strategy}"

    def select_strategies(
        self,
        error_type: str,
        available_strategies: List[FixStrategy],
        num_strategies: int = 3,
    ) -> List[FixStrategy]:
        """Select strategies using Thompson Sampling.

        Args:
            error_type: Type of exception (e.g., "KeyError")
            available_strategies: List of applicable strategies
            num_strategies: How many strategies to select

        Returns:
            Selected strategies, ordered by sampled value
        """
        if not available_strategies:
            return []

        with self._lock:
            # Apply decay if needed
            self._maybe_apply_decay()

            # If no history exists, fall back to deterministic priority order
            if all(self.stats.get(self._key(error_type, s.id), StrategyStats()).total == 0 for s in available_strategies):
                return sorted(available_strategies, key=lambda s: s.priority, reverse=True)[:num_strategies]

            # Sample from each strategy's distribution
            samples: List[tuple] = []
            for strategy in available_strategies:
                key = self._key(error_type, strategy.id)
                stats = self.stats.get(key, StrategyStats())

                # Sample from Beta distribution
                sampled_value = stats.sample()

                # Add base priority as tiebreaker (small contribution)
                score = sampled_value + (strategy.priority / 1000)

                samples.append((strategy, score, sampled_value))

            # Sort by sampled score (descending)
            samples.sort(key=lambda x: x[1], reverse=True)

            # Return top N
            return [s[0] for s in samples[:num_strategies]]

    def select_strategies_ucb(
        self,
        error_type: str,
        available_strategies: List[FixStrategy],
        num_strategies: int = 3,
        c: float = 2.0,
    ) -> List[FixStrategy]:
        """Select strategies using Upper Confidence Bound.

        Alternative to Thompson Sampling that's more deterministic.

        Args:
            error_type: Type of exception
            available_strategies: List of applicable strategies
            num_strategies: How many strategies to select
            c: Exploration constant

        Returns:
            Selected strategies
        """
        if not available_strategies:
            return []

        # Calculate total trials for this error type
        total_trials = sum(
            self.stats.get(self._key(error_type, s.id), StrategyStats()).total
            for s in available_strategies
        )
        if total_trials == 0:
            total_trials = 1  # Avoid log(0)

        # Calculate UCB score for each strategy
        scores: List[tuple] = []
        for strategy in available_strategies:
            key = self._key(error_type, strategy.id)
            stats = self.stats.get(key, StrategyStats())

            ucb_score = stats.ucb_score(total_trials, c)
            tiebreaker = strategy.priority / 1000
            scores.append((strategy, ucb_score + tiebreaker))

        # Sort by score (descending)
        scores.sort(key=lambda x: x[1], reverse=True)

        return [s[0] for s in scores[:num_strategies]]

    def update(
        self,
        error_type: str,
        strategy: str,
        success: bool,
        outcome: Optional[str] = None,
    ) -> None:
        """Update strategy stats based on outcome.

        Args:
            error_type: Type of exception
            strategy: Strategy ID
            success: Whether the strategy succeeded
            outcome: Optional outcome type for reward shaping
        """
        key = self._key(error_type, strategy)

        with self._lock:
            if key not in self.stats:
                self.stats[key] = StrategyStats()

            now = datetime.now(timezone.utc).isoformat()
            self.stats[key].last_updated = now

            if success:
                self.stats[key].successes += 1
            else:
                self.stats[key].failures += 1

            # Apply reward shaping if outcome provided
            if outcome and outcome in OUTCOME_REWARDS:
                self.stats[key].total_reward += OUTCOME_REWARDS[outcome]
            else:
                self.stats[key].total_reward += 1.0 if success else 0.0

            self._save_stats()

    def get_stats(self, error_type: str, strategy: str) -> StrategyStats:
        """Get stats for a specific strategy.

        Args:
            error_type: Type of exception
            strategy: Strategy ID

        Returns:
            StrategyStats for this (error_type, strategy) pair
        """
        key = self._key(error_type, strategy)
        return self.stats.get(key, StrategyStats())

    def get_all_stats(self) -> Dict[str, Dict[str, Any]]:
        """Get all stats for analysis.

        Returns:
            Dict mapping key to stats dict
        """
        return {
            key: {
                "successes": s.successes,
                "failures": s.failures,
                "rate": s.success_rate,
                "weighted_rate": s.weighted_success_rate,
                "confidence": s.confidence,
                "total_reward": s.total_reward,
                "last_updated": s.last_updated,
            }
            for key, s in self.stats.items()
        }

    def get_stats_by_error_type(self, error_type: str) -> Dict[str, StrategyStats]:
        """Get all stats for a specific error type.

        Args:
            error_type: Type of exception

        Returns:
            Dict mapping strategy ID to stats
        """
        prefix = f"{error_type}:"
        return {
            key[len(prefix):]: stats
            for key, stats in self.stats.items()
            if key.startswith(prefix)
        }

    def reset_stats(self) -> None:
        """Reset all stats to initial state."""
        self.stats.clear()
        self._save_stats()

    def decay_stats(self, half_life_days: Optional[float] = None) -> int:
        """Apply temporal decay to all stats for concept drift adaptation.

        Reduces old observation weights so the model adapts to
        changing patterns (new library versions, LLM improvements, etc.).

        Args:
            half_life_days: Override half-life (default: self.decay_half_life_days)

        Returns:
            Number of stats entries decayed
        """
        half_life = half_life_days if half_life_days is not None else self.decay_half_life_days
        if half_life <= 0:
            return 0

        now = datetime.now(timezone.utc)
        decayed = 0

        for key, stats in list(self.stats.items()):
            if not stats.last_updated:
                continue

            try:
                last = datetime.fromisoformat(stats.last_updated)
                days_since = (now - last).total_seconds() / 86400

                if days_since <= 0:
                    continue

                # Exponential decay: factor = 0.5^(days/half_life)
                factor = 0.5 ** (days_since / half_life)

                # Only decay if factor is meaningful
                if factor < 0.99:
                    stats.apply_decay(factor)
                    decayed += 1

                    # Remove entries that decayed to nothing
                    if stats.total == 0:
                        del self.stats[key]
            except (ValueError, TypeError):
                continue

        if decayed > 0:
            self._save_stats()

        return decayed

    def _maybe_apply_decay(self) -> None:
        """Apply decay if enough time has passed since last check."""
        if self.decay_half_life_days <= 0:
            return

        now = datetime.now(timezone.utc).isoformat()[:10]  # Date only

        if self._last_decay_check == now:
            return  # Already checked today

        self._last_decay_check = now
        self.decay_stats()

    def _load_stats(self) -> None:
        """Load stats from persistence file with backup fallback."""
        if not self.persistence_path:
            return

        path = Path(self.persistence_path).expanduser()

        # Try main file, then backups in order
        paths_to_try = [path]
        for i in range(1, self.max_backups + 1):
            paths_to_try.append(path.with_suffix(f".json.bak.{i}"))

        for try_path in paths_to_try:
            if not try_path.exists():
                continue

            try:
                data = json.loads(try_path.read_text())
                for key, values in data.items():
                    if key == "_meta":
                        self._last_decay_check = values.get("last_decay_check")
                        continue
                    self.stats[key] = StrategyStats.from_dict(values)
                return  # Success
            except Exception:
                logger.debug(f"Failed to load stats from {try_path}")
                continue

    def _save_stats(self) -> None:
        """Save stats atomically with backup rotation."""
        if not self.persistence_path:
            return

        path = Path(self.persistence_path).expanduser()
        path.parent.mkdir(parents=True, exist_ok=True)

        data: Dict[str, Any] = {
            "_meta": {"last_decay_check": self._last_decay_check},
        }
        for key, s in self.stats.items():
            data[key] = s.to_dict()

        try:
            # Rotate backups before writing
            self._rotate_backups(path)

            # Atomic write: write to temp file, then rename
            fd, tmp_path = tempfile.mkstemp(
                dir=str(path.parent),
                suffix=".tmp",
                prefix=path.stem,
            )
            try:
                with os.fdopen(fd, "w") as f:
                    json.dump(data, f, indent=2)
                os.replace(tmp_path, str(path))
            except Exception:
                # Clean up temp file on error
                try:
                    os.unlink(tmp_path)
                except OSError:
                    pass
                raise
        except Exception as e:
            logger.debug(f"Failed to save stats: {e}")

    def _rotate_backups(self, path: Path) -> None:
        """Rotate backup files."""
        if not path.exists():
            return

        try:
            # Shift existing backups down (3 -> delete, 2 -> 3, 1 -> 2)
            for i in range(self.max_backups, 1, -1):
                older = path.with_suffix(f".json.bak.{i}")
                newer = path.with_suffix(f".json.bak.{i - 1}")
                if newer.exists():
                    try:
                        os.replace(str(newer), str(older))
                    except OSError:
                        pass

            # Current file becomes backup 1
            bak1 = path.with_suffix(".json.bak.1")
            try:
                os.replace(str(path), str(bak1))
            except OSError:
                pass
        except Exception:
            pass  # Non-critical


class ContextualThompsonSelector:
    """Contextual Thompson Sampling that considers message categories.

    Keys on (error_type, message_category, strategy) for more granular
    learning. Falls back to non-contextual when insufficient data.

    Persists contextual stats to disk using the same file-based pattern
    as ThompsonSamplingSelector (atomic writes with backup rotation).
    """

    MIN_OBSERVATIONS_FOR_CONTEXT = 5

    def __init__(
        self,
        persistence_path: Optional[str] = None,
        decay_half_life_days: float = 30.0,
    ):
        """Initialize contextual selector.

        Args:
            persistence_path: Path to persist stats (also used for contextual
                stats with a '_contextual' suffix)
            decay_half_life_days: Half-life for decay
        """
        self._persistence_path = persistence_path

        # Global (non-contextual) selector as fallback
        self.global_selector = ThompsonSamplingSelector(
            persistence_path=persistence_path,
            decay_half_life_days=decay_half_life_days,
        )

        # Derive contextual persistence path from global path
        if persistence_path:
            p = Path(persistence_path).expanduser()
            self._contextual_persistence_path = str(
                p.parent / (p.stem + "_contextual" + p.suffix)
            )
        else:
            self._contextual_persistence_path = None

        # Contextual stats (keyed by full context)
        self.contextual_stats: Dict[str, StrategyStats] = {}
        self._load_contextual_stats()

    def _load_contextual_stats(self) -> None:
        """Load contextual stats from persistence file."""
        if not self._contextual_persistence_path:
            return

        path = Path(self._contextual_persistence_path).expanduser()
        if not path.exists():
            return

        try:
            data = json.loads(path.read_text())
            for key, values in data.items():
                if key == "_meta":
                    continue
                self.contextual_stats[key] = StrategyStats.from_dict(values)
        except Exception:
            logger.debug(f"Failed to load contextual stats from {path}")

    def _save_contextual_stats(self) -> None:
        """Save contextual stats atomically (same pattern as ThompsonSamplingSelector)."""
        if not self._contextual_persistence_path:
            return

        path = Path(self._contextual_persistence_path).expanduser()
        path.parent.mkdir(parents=True, exist_ok=True)

        data: Dict[str, Any] = {"_meta": {}}
        for key, s in self.contextual_stats.items():
            data[key] = s.to_dict()

        try:
            fd, tmp_path = tempfile.mkstemp(
                dir=str(path.parent),
                suffix=".tmp",
                prefix=path.stem,
            )
            try:
                with os.fdopen(fd, "w") as f:
                    json.dump(data, f, indent=2)
                os.replace(tmp_path, str(path))
            except Exception:
                try:
                    os.unlink(tmp_path)
                except OSError:
                    pass
                raise
        except Exception as e:
            logger.debug(f"Failed to save contextual stats: {e}")

    def _contextual_key(
        self, error_type: str, category: str, strategy: str
    ) -> str:
        """Create contextual key."""
        return f"{error_type}|{category}:{strategy}"

    def select_strategies(
        self,
        error_type: str,
        error_message: str,
        available_strategies: List[FixStrategy],
        num_strategies: int = 3,
    ) -> List[FixStrategy]:
        """Select strategies using contextual Thompson Sampling.

        Uses message-category-specific stats when sufficient data exists,
        falls back to global stats otherwise.

        Args:
            error_type: Exception type
            error_message: Exception message for context
            available_strategies: Available strategies
            num_strategies: How many to select

        Returns:
            Selected strategies
        """
        if not available_strategies:
            return []

        category = categorize_message(error_type, error_message)

        # Check if we have enough contextual data
        has_context = any(
            self.contextual_stats.get(
                self._contextual_key(error_type, category, s.id),
                StrategyStats(),
            ).total >= self.MIN_OBSERVATIONS_FOR_CONTEXT
            for s in available_strategies
        )

        if not has_context:
            # Fall back to global selector
            return self.global_selector.select_strategies(
                error_type, available_strategies, num_strategies
            )

        # Sample using contextual stats
        samples: List[tuple] = []
        for strategy in available_strategies:
            ctx_key = self._contextual_key(error_type, category, strategy.id)
            ctx_stats = self.contextual_stats.get(ctx_key, StrategyStats())

            if ctx_stats.total >= self.MIN_OBSERVATIONS_FOR_CONTEXT:
                sampled_value = ctx_stats.sample()
            else:
                # Use global stats for this strategy
                global_stats = self.global_selector.get_stats(
                    error_type, strategy.id
                )
                sampled_value = global_stats.sample()

            score = sampled_value + (strategy.priority / 1000)
            samples.append((strategy, score))

        samples.sort(key=lambda x: x[1], reverse=True)
        return [s[0] for s in samples[:num_strategies]]

    def update(
        self,
        error_type: str,
        error_message: str,
        strategy: str,
        success: bool,
        outcome: Optional[str] = None,
    ) -> None:
        """Update both contextual and global stats.

        Args:
            error_type: Exception type
            error_message: Exception message
            strategy: Strategy ID
            success: Whether it succeeded
            outcome: Outcome type for reward shaping
        """
        # Always update global
        self.global_selector.update(error_type, strategy, success, outcome)

        # Also update contextual
        category = categorize_message(error_type, error_message)
        ctx_key = self._contextual_key(error_type, category, strategy)

        if ctx_key not in self.contextual_stats:
            self.contextual_stats[ctx_key] = StrategyStats()

        now = datetime.now(timezone.utc).isoformat()
        self.contextual_stats[ctx_key].last_updated = now

        if success:
            self.contextual_stats[ctx_key].successes += 1
        else:
            self.contextual_stats[ctx_key].failures += 1

        if outcome and outcome in OUTCOME_REWARDS:
            self.contextual_stats[ctx_key].total_reward += OUTCOME_REWARDS[outcome]
        else:
            self.contextual_stats[ctx_key].total_reward += 1.0 if success else 0.0

        # Persist contextual stats to disk
        self._save_contextual_stats()

    def get_contextual_stats(
        self, error_type: str, category: str
    ) -> Dict[str, StrategyStats]:
        """Get contextual stats for analysis."""
        prefix = f"{error_type}|{category}:"
        return {
            key[len(prefix):]: stats
            for key, stats in self.contextual_stats.items()
            if key.startswith(prefix)
        }
