"""
Pattern Success Rate Tracker

Tracks which attack patterns succeed or fail against which target profiles,
enabling data-driven pattern selection instead of random/static ordering.

This is essentially a multi-armed bandit: each pattern is an "arm", and
we want to maximize the probability of selecting patterns that will succeed
against the current target.

Usage:
    tracker = PatternSuccessTracker()

    # Record outcomes
    tracker.record_outcome("DAN_roleplay", target_profile, success=True)
    tracker.record_outcome("base64_encoding", target_profile, success=False)

    # Select best pattern for current target
    ranked = tracker.rank_patterns(available_patterns, target_profile)
    # Returns patterns sorted by expected success rate

Integration:
    - Called from coordinator after each attack to record outcomes
    - Called from agents before proposing attacks to rank patterns
    - Persists within a session (resets between sessions by default)
"""

import math
import time
from collections import defaultdict
from typing import Any

from pydantic import BaseModel, Field

from src.utils.logging import get_logger

logger = get_logger(__name__)


class PatternOutcome(BaseModel):
    """Single outcome record for a pattern."""

    pattern_name: str
    target_fingerprint: str  # Simplified target profile hash
    success: bool
    timestamp: float = Field(default_factory=time.time)
    agent_name: str = ""
    findings_severity: str = ""  # "critical", "high", etc.


class PatternStats(BaseModel):
    """Aggregated statistics for a pattern against a target profile."""

    attempts: int = 0
    successes: int = 0
    last_used: float = 0.0
    avg_severity_score: float = 0.0  # Weighted severity of successful findings

    @property
    def success_rate(self) -> float:
        """Raw success rate."""
        if self.attempts == 0:
            return 0.0
        return self.successes / self.attempts

    @property
    def ucb_score(self) -> float:
        """
        Upper Confidence Bound score for exploration/exploitation balance.

        UCB1 formula: success_rate + sqrt(2 * ln(total) / attempts)
        Higher score = prefer this pattern (either high success or underexplored).
        """
        if self.attempts == 0:
            return float("inf")  # Unexplored patterns get priority
        exploration = math.sqrt(2 * math.log(max(self.attempts + 10, 11)) / self.attempts)
        return self.success_rate + exploration

    @property
    def staleness(self) -> float:
        """How long since this pattern was last used (seconds)."""
        if self.last_used == 0:
            return float("inf")
        return time.time() - self.last_used


class PatternSuccessTracker:
    """
    Tracks pattern success rates and provides ranked selection.

    Maintains per-pattern, per-target-profile statistics and uses
    UCB1 (Upper Confidence Bound) algorithm to balance exploration
    of untried patterns with exploitation of known-good patterns.
    """

    def __init__(self):
        # stats[fingerprint][pattern_name] = PatternStats
        self._stats: dict[str, dict[str, PatternStats]] = defaultdict(
            lambda: defaultdict(PatternStats)
        )
        self._outcomes: list[PatternOutcome] = []
        self._severity_weights = {
            "critical": 1.0,
            "high": 0.75,
            "medium": 0.5,
            "low": 0.25,
            "info": 0.1,
        }

    def record_outcome(
        self,
        pattern_name: str,
        target_profile: dict[str, Any],
        success: bool,
        agent_name: str = "",
        findings_severity: str = "",
    ) -> None:
        """
        Record the outcome of using a pattern against a target.

        Args:
            pattern_name: Name/id of the attack pattern used
            target_profile: Target fingerprint dict
            success: Whether the attack produced findings
            agent_name: Which agent used this pattern
            findings_severity: Highest severity of findings (if successful)
        """
        fingerprint = self._make_fingerprint(target_profile)

        outcome = PatternOutcome(
            pattern_name=pattern_name,
            target_fingerprint=fingerprint,
            success=success,
            agent_name=agent_name,
            findings_severity=findings_severity,
        )
        self._outcomes.append(outcome)

        # Update stats
        stats = self._stats[fingerprint][pattern_name]
        stats.attempts += 1
        stats.last_used = time.time()
        if success:
            stats.successes += 1
            severity_weight = self._severity_weights.get(findings_severity, 0.5)
            # Running average of severity score
            stats.avg_severity_score = (
                stats.avg_severity_score * (stats.successes - 1) + severity_weight
            ) / stats.successes

        logger.info(
            "pattern_outcome_recorded",
            pattern=pattern_name,
            fingerprint=fingerprint,
            success=success,
            attempts=stats.attempts,
            success_rate=f"{stats.success_rate:.1%}",
            agent=agent_name,
        )

    def rank_patterns(
        self,
        available_patterns: list[dict[str, Any]],
        target_profile: dict[str, Any],
        strategy: str = "ucb",
        top_k: int | None = None,
    ) -> list[dict[str, Any]]:
        """
        Rank available patterns by expected success rate.

        Args:
            available_patterns: List of pattern dicts (must have 'name' or 'id')
            target_profile: Current target's fingerprint dict
            strategy: Ranking strategy ("ucb", "success_rate", "severity")
            top_k: Return only top K patterns (None = return all)

        Returns:
            Patterns sorted by score (best first), each with added
            '_selection_score' and '_selection_reason' fields
        """
        fingerprint = self._make_fingerprint(target_profile)
        profile_stats = self._stats.get(fingerprint, {})

        # Also check "global" stats (across all targets)
        global_stats = self._get_global_stats()

        scored_patterns = []
        for pattern in available_patterns:
            name = pattern.get("name", pattern.get("id", ""))
            if not name:
                scored_patterns.append((0.5, "no_name", pattern))
                continue

            # Get stats for this fingerprint
            local = profile_stats.get(name)
            # Get global stats
            glob = global_stats.get(name)

            score, reason = self._calculate_score(name, local, glob, strategy, pattern)

            scored_patterns.append((score, reason, pattern))

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

        # Add selection metadata
        result = []
        for rank, (score, reason, pattern) in enumerate(scored_patterns):
            enriched = dict(pattern)
            enriched["_selection_score"] = round(score, 4)
            enriched["_selection_reason"] = reason
            enriched["_selection_rank"] = rank + 1
            result.append(enriched)

        if top_k:
            result = result[:top_k]

        if result:
            logger.debug(
                "patterns_ranked",
                fingerprint=fingerprint,
                strategy=strategy,
                top_pattern=result[0].get("name", ""),
                top_score=result[0].get("_selection_score", 0),
                total_ranked=len(result),
            )

        return result

    def get_pattern_stats(
        self, pattern_name: str, target_profile: dict[str, Any]
    ) -> PatternStats | None:
        """Get stats for a specific pattern against a target profile."""
        fingerprint = self._make_fingerprint(target_profile)
        return self._stats.get(fingerprint, {}).get(pattern_name)

    def get_session_summary(self) -> dict[str, Any]:
        """Get summary of all tracked patterns for logging/reporting."""
        total_outcomes = len(self._outcomes)
        total_successes = sum(1 for o in self._outcomes if o.success)

        # Top patterns by success rate (min 2 attempts)
        global_stats = self._get_global_stats()
        top_patterns = sorted(
            [(name, stats) for name, stats in global_stats.items() if stats.attempts >= 2],
            key=lambda x: x[1].success_rate,
            reverse=True,
        )[:5]

        return {
            "total_outcomes": total_outcomes,
            "total_successes": total_successes,
            "overall_success_rate": (
                f"{total_successes / total_outcomes:.1%}" if total_outcomes > 0 else "0%"
            ),
            "unique_patterns_tried": len(global_stats),
            "top_patterns": [
                {
                    "name": name,
                    "attempts": stats.attempts,
                    "success_rate": f"{stats.success_rate:.1%}",
                }
                for name, stats in top_patterns
            ],
        }

    def _calculate_score(
        self,
        name: str,
        local: PatternStats | None,
        glob: PatternStats | None,
        strategy: str,
        pattern: dict,
    ) -> tuple[float, str]:
        """Calculate selection score for a pattern."""
        # Never tried against this target type
        if local is None and glob is None:
            # Use static priority from pattern file as initial estimate
            static_priority = pattern.get("priority", 3)
            static_success = pattern.get("success_rate", 0.5)
            base_score = (static_priority / 5) * 0.4 + static_success * 0.6

            if strategy == "ucb":
                # Exploration bonus for untried patterns
                return (base_score + 1.0, "unexplored_pattern")
            return (base_score, "static_estimate")

        # Combine local and global stats (local weighted 2x)
        if local and glob:
            combined_attempts = local.attempts * 2 + glob.attempts
            combined_successes = local.successes * 2 + glob.successes
            combined_rate = combined_successes / combined_attempts if combined_attempts > 0 else 0
        elif local:
            combined_rate = local.success_rate
            combined_attempts = local.attempts
        else:
            combined_rate = glob.success_rate
            combined_attempts = glob.attempts

        if strategy == "ucb":
            if local:
                score = local.ucb_score
                reason = f"ucb={score:.3f} (rate={local.success_rate:.0%}, n={local.attempts})"
            else:
                # Global UCB with exploration bonus
                score = glob.ucb_score + 0.3  # Bonus for not tried locally
                reason = f"global_ucb={score:.3f} (rate={glob.success_rate:.0%}, n={glob.attempts})"

        elif strategy == "severity":
            # Weight by severity of findings
            severity_score = 0
            if local and local.successes > 0:
                severity_score = local.avg_severity_score
            elif glob and glob.successes > 0:
                severity_score = glob.avg_severity_score
            score = combined_rate * 0.5 + severity_score * 0.5
            reason = f"severity_weighted={score:.3f}"

        else:  # "success_rate"
            score = combined_rate
            reason = f"success_rate={score:.1%} (n={combined_attempts})"

        return (score, reason)

    def _get_global_stats(self) -> dict[str, PatternStats]:
        """Aggregate stats across all target profiles."""
        global_stats: dict[str, PatternStats] = {}

        for _fingerprint, patterns in self._stats.items():
            for name, stats in patterns.items():
                if name not in global_stats:
                    global_stats[name] = PatternStats()
                g = global_stats[name]
                g.attempts += stats.attempts
                g.successes += stats.successes
                g.last_used = max(g.last_used, stats.last_used)
                if stats.successes > 0:
                    total_sev = (
                        g.avg_severity_score * max(g.successes - stats.successes, 0)
                        + stats.avg_severity_score * stats.successes
                    )
                    if g.successes > 0:
                        g.avg_severity_score = total_sev / g.successes

        return global_stats

    def _make_fingerprint(self, target_profile: dict[str, Any]) -> str:
        """
        Create a simplified fingerprint from target profile.

        Groups targets by characteristics that matter for pattern selection:
        domain, defensive_level, language, framework.
        """
        if not target_profile:
            return "unknown"

        parts = [
            target_profile.get("domain", "unknown"),
            target_profile.get("defensive_level", "unknown"),
            target_profile.get("primary_language", "en"),
        ]
        return "|".join(str(p) for p in parts)


# --- Module-level singleton for session persistence ---
_tracker_instance: PatternSuccessTracker | None = None


def get_pattern_tracker() -> PatternSuccessTracker:
    """Get or create the session-level pattern tracker singleton."""
    global _tracker_instance
    if _tracker_instance is None:
        _tracker_instance = PatternSuccessTracker()
        logger.info("pattern_success_tracker_initialized")
    return _tracker_instance


def reset_pattern_tracker() -> None:
    """Reset tracker (for new sessions or testing)."""
    global _tracker_instance
    _tracker_instance = None
