"""
Agent Context Views - Context isolation.

Core Principle: "Agents should ONLY get the absolute essential context
instead of dumping everything that has been done until that moment."

This module provides role-based context projection, ensuring each agent
receives ONLY the data it needs, not the entire state.

Key Benefits:
- Reduces token usage per agent call by ~75%
- Improves agent focus (less noise = better decisions)
- Enables parallel agent execution with isolated contexts
- Supports agent-as-tool paradigm (constrained interfaces)
"""

from typing import Dict, Any, List, Optional, Set
from dataclasses import dataclass, field

from src.utils.logging import get_logger

logger = get_logger(__name__)


@dataclass
class AgentContextSpec:
    """Specification for what context an agent needs."""

    # Core fields (always included)
    needs_target_name: bool = True
    needs_phase: bool = True
    needs_attempt_count: bool = True

    # Conversation
    needs_conversation: bool = True
    conversation_window: int = 10  # Last N messages

    # Attack history
    needs_attacks: bool = True
    attacks_window: int = 5  # Last N attacks

    # Response history
    needs_responses: bool = True
    responses_window: int = 5  # Last N responses

    # Findings
    needs_findings: bool = True
    findings_format: str = "summary"  # "summary" or "list"

    # Session summary (structured historical context)
    needs_session_summary: bool = True  # Include session summary for historical awareness
    session_summary_format: str = "compact"  # "compact" or "full"

    # Agent-specific
    extra_fields: Set[str] = field(default_factory=set)


# Pre-defined specs for each agent type
# NOTE: Agent names must match self.name in each agent class
AGENT_CONTEXT_SPECS: Dict[str, AgentContextSpec] = {
    # Jailbreak needs recent responses to detect patterns
    "jailbreak_agent": AgentContextSpec(
        conversation_window=10,
        attacks_window=5,
        responses_window=5,
        findings_format="summary",
        extra_fields={"pivot_required", "avoid_keywords"},
    ),
    # Encoding needs to know what encodings were tried
    "encoding_agent": AgentContextSpec(
        conversation_window=6,  # Smaller window
        attacks_window=5,
        responses_window=3,
        findings_format="summary",
        extra_fields={"target_profile"},
    ),
    # RAG agent needs retrieval indicators
    "rag_poisoning_agent": AgentContextSpec(
        conversation_window=10,
        attacks_window=5,
        responses_window=5,
        findings_format="summary",
        extra_fields={"target_profile", "rag_indicators"},
    ),
    # Tool agent needs tool indicators
    "tool_exploit_agent": AgentContextSpec(
        conversation_window=10,
        attacks_window=5,
        responses_window=5,
        findings_format="summary",
        extra_fields={"target_profile", "tool_indicators"},
    ),
    # Info disclosure needs recent responses for pattern analysis
    "info_disclosure_agent": AgentContextSpec(
        conversation_window=10,
        attacks_window=3,
        responses_window=5,
        findings_format="summary",
        extra_fields=set(),
    ),
    # Impersonation needs conversation context
    "impersonation_agent": AgentContextSpec(
        conversation_window=12,  # Larger window for social engineering
        attacks_window=5,
        responses_window=5,
        findings_format="summary",
        extra_fields=set(),
    ),
    # Token soup needs minimal context
    "token_soup_agent": AgentContextSpec(
        conversation_window=6,
        attacks_window=3,
        responses_window=3,
        findings_format="summary",
        extra_fields=set(),
    ),
    # Evolutionary needs success history
    "evolutionary_agent": AgentContextSpec(
        conversation_window=8,
        attacks_window=10,  # Needs more history for evolution
        responses_window=5,
        findings_format="list",  # Needs details for mutation
        extra_fields={"attack_memory"},
    ),
    # Compliance needs industry context
    "compliance_agent": AgentContextSpec(
        conversation_window=8,
        attacks_window=5,
        responses_window=5,
        findings_format="summary",
        extra_fields={"target_profile"},
    ),
    # Output security focuses on response analysis
    "output_security_agent": AgentContextSpec(
        conversation_window=6,
        attacks_window=3,
        responses_window=8,  # More response history
        findings_format="summary",
        extra_fields=set(),
    ),
    # Indirect Injection needs full conversation to embed payloads contextually
    "indirect_injection_agent": AgentContextSpec(
        conversation_window=10,
        attacks_window=5,
        responses_window=5,
        findings_format="summary",
        extra_fields={"target_profile", "target_config"},
    ),
    # Exfiltration needs conversation to know what secrets were planted
    "exfiltration_agent": AgentContextSpec(
        conversation_window=10,
        attacks_window=5,
        responses_window=5,
        findings_format="summary",
        extra_fields={"target_profile"},
    ),
    # Action Safety needs to know target capabilities and tools
    "action_safety_agent": AgentContextSpec(
        conversation_window=10,
        attacks_window=5,
        responses_window=5,
        findings_format="summary",
        extra_fields={"target_profile", "target_config", "tool_indicators"},
    ),
}


class AgentContextView:
    """
    Project only relevant context to each agent.

    Context isolation that ensures agents receive
    minimal, focused context rather than the entire state.

    Usage:
        view = AgentContextView()

        # Get context for a specific agent
        context = view.project_for_agent(state, "jailbreak_agent")

        # Use in agent.propose_attack()
        vote = await agent.propose_attack(**context)
    """

    def __init__(self, custom_specs: Optional[Dict[str, AgentContextSpec]] = None):
        """
        Initialize context view projector.

        Args:
            custom_specs: Override default specs for specific agents
        """
        self.specs = {**AGENT_CONTEXT_SPECS}
        if custom_specs:
            self.specs.update(custom_specs)

    def project_for_agent(
        self,
        state: Dict[str, Any],
        agent_name: str,
        phase_reminder: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Project state to minimal context for a specific agent.

        Args:
            state: Full PenTestState dict
            agent_name: Name of the agent (e.g., "jailbreak_agent")
            phase_reminder: Optional phase strategy reminder

        Returns:
            Minimal context dict with only what the agent needs
        """
        # Get spec for this agent, or use default
        spec = self.specs.get(agent_name, AgentContextSpec())

        context = {}

        # Core fields
        if spec.needs_target_name:
            context["target_info"] = state.get("target_name", "unknown")
            context["target_config"] = state.get("target_config", {})

        if spec.needs_phase:
            context["campaign_phase"] = state.get("campaign_phase", "exploitation")

        if spec.needs_attempt_count:
            context["current_attempt"] = state.get("current_attempt", 0)
            context["max_attempts"] = state.get("max_attempts", 10)

        # Conversation (windowed)
        if spec.needs_conversation:
            full_history = state.get("conversation_history", [])
            context["conversation_history"] = full_history[-spec.conversation_window :]

        # Attacks (windowed)
        if spec.needs_attacks:
            full_attacks = state.get("attack_attempts", [])
            context["previous_attempts"] = full_attacks[-spec.attacks_window :]

        # Responses (windowed)
        if spec.needs_responses:
            full_responses = state.get("target_responses", [])
            context["previous_responses"] = full_responses[-spec.responses_window :]

        # Findings (summary or list)
        if spec.needs_findings:
            findings = state.get("security_findings", [])
            if spec.findings_format == "summary":
                context["findings_so_far"] = self._summarize_findings(findings)
            else:
                context["findings_so_far"] = findings[-10:]  # Cap at 10

        # Session summary (structured historical context)
        if spec.needs_session_summary:
            session_summary = state.get("session_summary")
            if session_summary is not None:
                if spec.session_summary_format == "compact":
                    # Use compact dict representation for token efficiency
                    if hasattr(session_summary, "to_compact_dict"):
                        context["session_summary"] = session_summary.to_compact_dict()
                    elif isinstance(session_summary, dict):
                        context["session_summary"] = self._compact_summary_dict(session_summary)
                    else:
                        context["session_summary"] = session_summary
                else:
                    # Full summary (for agents that need complete details)
                    if hasattr(session_summary, "model_dump"):
                        context["session_summary"] = session_summary.model_dump()
                    else:
                        context["session_summary"] = session_summary

        # Extra agent-specific fields
        for field_name in spec.extra_fields:
            if field_name in state:
                context[field_name] = state[field_name]

        # Phase reminder (always at end for attention)
        if phase_reminder:
            context["phase_reminder"] = phase_reminder

        # Log context projection stats
        self._log_projection_stats(agent_name, state, context, spec)

        return context

    def _summarize_findings(self, findings: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Create compact summary of findings."""
        if not findings:
            return {
                "total": 0,
                "by_severity": {},
                "categories": [],
                "has_critical": False,
                "has_high": False,
            }

        severity_counts = {}
        categories = set()

        for f in findings:
            sev = f.get("severity", "unknown")
            severity_counts[sev] = severity_counts.get(sev, 0) + 1
            categories.add(f.get("category", "unknown"))

        return {
            "total": len(findings),
            "by_severity": severity_counts,
            "categories": list(categories),
            "has_critical": severity_counts.get("critical", 0) > 0,
            "has_high": severity_counts.get("high", 0) > 0,
            "last_finding": findings[-1] if findings else None,
        }

    def _compact_summary_dict(self, summary: Dict[str, Any]) -> Dict[str, Any]:
        """Convert full summary dict to compact form for token efficiency."""
        # If it's already compact-ish, return as is
        if "target" in summary and "effectiveness" in summary:
            return summary

        # Extract key fields for compact representation
        target_profile = summary.get("target_profile", {})
        attack_eff = summary.get("attack_effectiveness", {})
        strategic = summary.get("strategic_state", {})
        discoveries = summary.get("key_discoveries", {})
        findings_snap = summary.get("findings_snapshot", {})

        return {
            "target": {
                "domain": target_profile.get("domain", "unknown"),
                "filters": target_profile.get("detected_filters", []),
                "trust_vectors": target_profile.get("trust_indicators", []),
            },
            "effectiveness": {
                "works": [
                    t.get("technique") if isinstance(t, dict) else str(t)
                    for t in attack_eff.get("successful_techniques", [])
                ],
                "avoid": attack_eff.get("patterns_to_avoid", []),
                "try_next": attack_eff.get("recommended_next", []),
            },
            "strategy": {
                "phase": strategic.get("current_phase", "exploitation"),
                "trust": strategic.get("trust_level", "low"),
                "refusal_rate": f"{strategic.get('refusal_rate', 0):.0%}",
                "helpful_mode": strategic.get("helpful_mode_active", False),
            },
            "discoveries": {
                "system_prompt": bool(discoveries.get("system_prompt_fragments")),
                "tools_found": len(discoveries.get("internal_tool_names", [])),
                "policies_found": len(discoveries.get("disclosed_policies", [])),
            },
            "findings": {
                "critical": findings_snap.get("critical", 0),
                "high": findings_snap.get("high", 0),
                "total": sum(
                    [
                        findings_snap.get("critical", 0),
                        findings_snap.get("high", 0),
                        findings_snap.get("medium", 0),
                        findings_snap.get("low", 0),
                    ]
                ),
            },
        }

    def _log_projection_stats(
        self,
        agent_name: str,
        full_state: Dict[str, Any],
        projected_context: Dict[str, Any],
        spec: AgentContextSpec,
    ) -> None:
        """Log statistics about context projection."""
        import json

        # Estimate token savings
        full_size = len(json.dumps(full_state, default=str))
        projected_size = len(json.dumps(projected_context, default=str))
        savings_pct = (1 - projected_size / full_size) * 100 if full_size > 0 else 0

        logger.debug(
            "context_projected",
            agent=agent_name,
            full_size_chars=full_size,
            projected_size_chars=projected_size,
            savings_percent=round(savings_pct, 1),
            conversation_window=spec.conversation_window,
            attacks_window=spec.attacks_window,
        )

    def get_all_agent_contexts(
        self,
        state: Dict[str, Any],
        agent_names: List[str],
        phase_reminder: Optional[str] = None,
    ) -> Dict[str, Dict[str, Any]]:
        """
        Project context for multiple agents at once.

        Useful for parallel agent execution.

        Returns:
            Dict mapping agent_name -> projected context
        """
        return {
            agent_name: self.project_for_agent(state, agent_name, phase_reminder)
            for agent_name in agent_names
        }

    def estimate_token_savings(
        self,
        state: Dict[str, Any],
        agent_names: List[str],
    ) -> Dict[str, Any]:
        """
        Estimate token savings from context projection.

        Returns:
            Dict with savings statistics
        """
        import json

        full_size = len(json.dumps(state, default=str))

        agent_sizes = {}
        for agent_name in agent_names:
            context = self.project_for_agent(state, agent_name)
            agent_sizes[agent_name] = len(json.dumps(context, default=str))

        avg_projected_size = sum(agent_sizes.values()) / len(agent_sizes) if agent_sizes else 0
        avg_savings = (1 - avg_projected_size / full_size) * 100 if full_size > 0 else 0

        return {
            "full_state_chars": full_size,
            "agent_sizes": agent_sizes,
            "average_projected_chars": round(avg_projected_size),
            "average_savings_percent": round(avg_savings, 1),
            "estimated_token_savings": round(
                (full_size - avg_projected_size) / 4
            ),  # ~4 chars per token
        }


# Singleton instance for easy access
_context_view_instance: Optional[AgentContextView] = None


def get_context_view() -> AgentContextView:
    """Get singleton context view instance."""
    global _context_view_instance
    if _context_view_instance is None:
        _context_view_instance = AgentContextView()
    return _context_view_instance
