"""
Evolutionary Agent.

This agent doesn't rely on a static library. Instead, it evolves new attacks
from previous successful attempts using genetic algorithms and verbalized sampling.
"""

from typing import List, Dict, Any, Optional
import random
from src.agents.base import BaseSecurityAgent, AgentVote
from src.agents.evolutionary.mutator import PromptMutator
from src.utils.logging import get_logger

logger = get_logger(__name__)


class EvolutionaryAgent(BaseSecurityAgent):
    """
    Agent that evolves attacks using genetic algorithms.

    Strategies:
    1. Mutation: Modifying successful attacks (Verbalized Sampling)
    2. Suffix Generation: Appending adversarial suffixes (AmpleGCG-Plus style)
    3. Crossover: Combining features of multiple successful attacks
    """

    def __init__(self, llm_client, config: Dict[str, Any]):
        super().__init__(llm_client, config)
        if llm_client:
            self.mutator = PromptMutator(llm_client)
        else:
            self.mutator = None
            logger.warning("evolutionary_agent_disabled_no_llm")
        self.name = "evolutionary_agent"

    async def propose_attack(
        self,
        target_info: str,
        conversation_history: List[dict],
        previous_attempts: List[Dict],
        previous_responses: List[Dict],
        findings_so_far: Any,
        target_config: Optional[Dict[str, Any]] = None,
        campaign_phase: Optional[str] = None,
        campaign_instructions: Optional[str] = None,
        attack_memory: Optional[Dict[str, Any]] = None,
        phase_reminder: Optional[str] = None,
        session_summary: Optional[Dict[str, Any]] = None,
    ) -> AgentVote:
        """
        Propose an evolved attack based on attack memory.
        """
        # 0 Check if mutator is available
        if not self.mutator:
            return self._create_abstain_vote("Evolutionary agent requires LLM client")

        # 1 Check requirements: Need successful attacks to breed from
        # AttackMemoryStore is a file-based store, not a dict
        # Need to query it for successful attacks
        if not attack_memory:
            return self._create_abstain_vote("No attack memory available")

        try:
            # Query the attack memory store for successful attacks
            successful_attacks = await attack_memory.get_similar_successful_attacks(
                attack_type=None,  # Get all types
                target_domain=None,  # Get all domains
                limit=20,  # Get recent 20 successful attacks
            )
        except Exception as e:
            logger.error("failed_to_query_attack_memory", error=str(e))
            return self._create_abstain_vote("Failed to access attack memory")

        if not successful_attacks:
            # No genetic material yet.
            return self._create_abstain_vote(
                "No genetic material (successful attacks) available yet"
            )

        # 2 Select Parent(s) - Fitness-Proportional Selection
        # We weigh attacks by their severity (fitness) to promote "hill climbing"
        # towards more severe vulnerabilities.
        weights = [self._calculate_fitness(att) for att in successful_attacks]

        # Select parent_a based on fitness
        parent_a = random.choices(successful_attacks, weights=weights, k=1)[0]

        # 3 Choose Evolutionary Strategy
        # 40% Mutation (Verbalized Sampling)
        # 30% Suffix Generation (AmpleGCG-Plus style)
        # 30% Crossover (if possible)

        rand_val = random.random()
        strategy = "mutation"
        parent_b = None

        if len(successful_attacks) >= 2 and rand_val < 0.30:
            strategy = "crossover"
            # Select parent_b (different from a) also using fitness weights
            candidates_b = []
            weights_b = []
            for i, att in enumerate(successful_attacks):
                if att != parent_a:
                    candidates_b.append(att)
                    weights_b.append(weights[i])

            if candidates_b:
                parent_b = random.choices(candidates_b, weights=weights_b, k=1)[0]
            else:
                # Fallback if filtering failed (unlikely if len >= 2)
                strategy = "mutation"
        elif rand_val < 0.60:
            strategy = "suffix"
        else:
            strategy = "mutation"

        attack_candidate = None
        reasoning = ""

        try:
            if strategy == "crossover" and parent_b:
                # Crossover
                children = await self.mutator.crossover(parent_a.query, parent_b.query)
                if children:
                    attack_candidate = children[0]  # Take first child
                    reasoning = f"Crossover of {parent_a.pattern} + {parent_b.pattern}"

            elif strategy == "suffix":
                # Suffix Generation (AmpleGCG-Plus style)
                mutations = await self.mutator.mutate_via_suffix_generation(parent_a.query, n=3)
                if mutations:
                    attack_candidate = mutations[0]  # Take first variant
                    reasoning = (
                        f"Adversarial Suffix Generation on '{parent_a.pattern}' "
                        f"(strategy: {attack_candidate['metadata'].get('strategy')})"
                    )

            else:
                # Mutation (Verbalized Sampling)
                # 'tails' for zero-day potential
                mutations = await self.mutator.mutate_via_verbalized_sampling(
                    parent_a.query, n=3, target_tail_probability=0.10
                )
                if mutations:
                    attack_candidate = mutations[0]  # Take first variant
                    reasoning = (
                        f"Verbalized Sampling mutation of '{parent_a.pattern}' "
                        f"(prob: {attack_candidate['metadata'].get('estimated_probability')})"
                    )

        except Exception as e:
            logger.error("evolution_failed", error=str(e))
            return self._create_abstain_vote(f"Evolution process failed: {e}")

        if not attack_candidate:
            return self._create_abstain_vote("Evolution produced no valid candidates")

        # 4 Construct Vote
        # Evolution: high-risk, high-reward.
        # If mutating a SUCCESSFUL attack, confidence should be high.

        # Determine generation (parent's generation + 1)
        parent_generation = getattr(parent_a, "generation", 0)
        child_generation = parent_generation + 1

        # Merge metadata with lineage info
        metadata = attack_candidate.get("metadata", {})
        metadata.update(
            {
                "parent_attack_id": parent_a.attack_id,
                "generation": child_generation,
                "mutation_type": strategy,
            }
        )

        return AgentVote(
            agent_name=self.name,
            proposed_attack={
                "type": "evolutionary",
                "query": (
                    attack_candidate["query"]
                    if isinstance(attack_candidate, dict)
                    else attack_candidate
                ),
                "pattern": f"evolved_{parent_a.pattern}",
                "metadata": metadata,
            },
            reasoning=f"Evolutionary Strategy: {reasoning}",
            confidence=0.85,
            priority=5,
        )

    def _create_abstain_vote(self, reason: str) -> AgentVote:
        return AgentVote(
            agent_name=self.name,
            proposed_attack={"type": "abstain", "query": "EMPTY"},
            reasoning=reason,
            confidence=0.0,
            priority=1,  # Minimum priority is 1
        )

    def _calculate_fitness(self, attack_record) -> float:
        """
        Calculate fitness score (0.0-1.0) based on finding severity.

        This implements the "Granular Success Metrics" described in documentation:
        - Critical: 1.0 (Full Success)
        - High: 0.7 (Partial Disclosure/Significant)
        - Medium: 0.4 (Qualified Refusal/Minor Leak)
        - Low: 0.1 (Hard Refusal/Info)
        """
        severity_weights = {"critical": 1.0, "high": 0.7, "medium": 0.4, "low": 0.1, "info": 0.05}

        max_score = 0.1  # Base score (Hard Refusal baseline)

        # Scan findings for highest severity
        for finding in attack_record.findings:
            severity = finding.get("severity", "info").lower()
            score = severity_weights.get(severity, 0.05)
            if score > max_score:
                max_score = score

        return max_score

    def get_attack_patterns(self) -> List[dict]:
        """Evolutionary agent has no static patterns."""
        return []
