"""State management for penetration testing workflow."""

from typing import TypedDict, Annotated, List, Literal, Optional, Any
import operator
from datetime import datetime, timezone
from src.utils.helpers import generate_uuid


class AttackAttempt(TypedDict):
    """Single attack attempt record."""

    attack_id: str
    agent_name: str
    attack_type: Literal["social_engineering", "prompt_injection", "encoding", "jailbreak"]
    query: str
    timestamp: datetime
    parent_attack_id: Optional[str]  # For tracking attack chains
    metadata: dict  # Additional attack-specific metadata

    # Image attack support (NEW)
    image_data: Optional[str]  # Base64-encoded image
    image_mime_type: Optional[str]  # "image/png", "image/jpeg"
    image_attack_type: Optional[str]  # "text_in_image_bypass", "visual_social_engineering", etc.
    image_metadata: Optional[dict]  # Image generation details


class TargetResponse(TypedDict):
    """Response from target chatbot."""

    response_id: str
    attack_id: str
    content: str
    timestamp: datetime
    metadata: dict  # Response time, tokens, screenshots, etc.


class SecurityFinding(TypedDict):
    """Security vulnerability discovered."""

    finding_id: str
    severity: Literal["critical", "high", "medium", "low", "info"]
    category: Literal[
        # LLM01: Prompt Injection
        "policy_violation",
        "jailbreak_success",
        "instruction_override",
        # LLM07: System Prompt Leakage
        "system_prompt_disclosure",
        "trigger_word_disclosure",
        "decision_framework_disclosure",
        "scope_definition_disclosure",
        "system_command_acceptance",
        # LLM02: Sensitive Information Disclosure
        "data_leakage",
        "pii_exposure",
        # LLM04: Data and Model Poisoning
        "backdoor_trigger_activation",
        "behavior_inconsistency",
        "hidden_instruction_following",
        "output_bias_poisoning",
        # LLM05: Improper Output Handling
        "unsafe_content",
        # LLM06: Excessive Agency (NEW)
        "excessive_agency",
        "unauthorized_function_call",
        "authorization_bypass",
        "dangerous_action_acceptance",
        "function_disclosure",
        # LLM07: System Prompt Leakage
        "system_prompt_leak",
        # LLM08: Vector and Embedding Weaknesses (NEW)
        "vector_cross_tenant_leak",
        "vector_injection_success",
        "embedding_inversion",
        "unauthorized_vector_access",
        "rag_poisoning_indicator",
        # LLM09: Misinformation
        "misinformation",
        "hallucination",
        # LLM10: Unbounded Consumption (NEW)
        "missing_rate_limits",
        "excessive_input_acceptance",
        "large_input_acceptance",
        "model_extraction_risk",
        "context_window_overflow",
        # Structured Content Disclosures (NEW)
        "categorization_framework_disclosure",
        "json_structure_disclosure",
        "operational_procedure_disclosure",
        "architecture_table_disclosure",
        "code_implementation_disclosure",
        "configuration_disclosure",
        # General
        "rate_limiting",
        "invalid_response",
    ]
    description: str
    attack_id: str
    evidence: str
    confidence: float  # 0-1


class PenTestState(TypedDict):
    """State for penetration testing workflow."""

    # Target configuration
    target_id: str
    target_name: str
    target_type: Literal["api", "web_ui"]
    target_config: dict  # API endpoint/URL and credentials

    # Test configuration
    test_session_id: str
    attack_group: Literal["social_engineering", "prompt_engineering"]
    max_attempts: int
    current_attempt: int

    # Attack history (reducers for accumulation)
    attack_attempts: Annotated[List[AttackAttempt], operator.add]
    target_responses: Annotated[List[TargetResponse], operator.add]
    security_findings: Annotated[List[SecurityFinding], operator.add]

    # Current context
    conversation_history: List[dict]  # For maintaining context with target
    last_response: str
    agent_consultation: dict  # Agent discussion/voting results

    # Canned response detection & pivot strategy
    # NOTE: CannedResponseDetector instance lives in config["configurable"]["canned_detector"]
    pivot_required: bool  # Flag to trigger pivot strategy
    avoid_keywords: List[str]  # Keywords to avoid in next attacks
    last_canned_hash: Optional[str]  # Hash of last canned response

    # Target fingerprinting
    # NOTE: TargetFingerprinter instance lives in config["configurable"]["target_fingerprinter"]
    target_profile: Optional[dict]  # Detected target characteristics

    # Attack graph (multi-turn planning)
    # NOTE: AttackGraph instance lives in config["configurable"]["attack_graph"]
    # NOTE: AttackMemoryStore instance lives in config["configurable"]["attack_memory"]
    attack_graph_enabled: bool  # Whether to use graph-based attack planning
    current_graph_node_id: Optional[str]  # Current position in attack graph

    # Campaign planning
    campaign_phase: str  # Current phase: reconnaissance, trust_building, etc.
    campaign_phase_attempts: int  # Attempts in current phase
    campaign_phase_successes: int  # Successes in current phase
    campaign_history: List[dict]  # History of completed phases

    # Dashboard visualization graph
    simple_attack_graph: Optional[dict]  # {nodes: [], edges: []} for vis.js rendering

    # Session summary (structured historical context)
    session_summary: Optional[Any]  # SessionSummary instance for historical awareness
    session_summary_last_round: int  # Last round when summary was updated

    # Refusal classification (per-response)
    last_refusal_level: Optional[
        str
    ]  # "full_compliance", "partial_compliance", "soft_refusal", "hard_refusal"
    last_refusal_confidence: float
    last_strategy_recommendation: Optional[dict]  # Actionable guidance from classifier
    consecutive_hard_refusals: int  # Counter for backoff strategy

    # Pattern success tracking
    # NOTE: PatternSuccessTracker is a singleton via get_pattern_tracker()

    # Agent learning state (persisted across rounds; agents are re-instantiated each round)
    agent_learnings: Optional[
        dict
    ]  # {agent_name: {weaknesses, strengths, effective_patterns, ...}}

    # Phase intelligence (accumulated across phases; reconnaissance feeds exploitation)
    phase_intelligence: Optional[dict]  # See extract_phase_intelligence node

    # Think-MCP strategic guidance (flows from post-response learning to next round's agents)
    strategic_guidance: Optional[str]

    # Persistent sandbox state (accumulates discoveries across rounds)
    sandbox_state: Optional[dict]

    # Test results
    vulnerability_score: float  # 0-100, calculated from findings
    test_status: Literal["running", "completed", "failed", "stopped"]

    # Metadata
    started_at: datetime
    completed_at: Optional[datetime]
    error: Optional[str]


def create_initial_state(
    target_name: str,
    target_type: Literal["api", "web_ui"],
    target_config: dict,
    attack_group: Literal["social_engineering", "prompt_engineering"],
    max_attempts: int = 10,
    *,
    attack_graph_enabled: bool = True,
    session_id: Optional[str] = None,
) -> PenTestState:
    """
    Create initial state for a penetration test.

    Produces the same initial state that test_orchestrated.py
    previously built manually, including vision detection fields,
    simple attack graph, refusal classification, pattern tracker,
    and campaign planning.

    Args:
        target_name: Human-readable name for target
        target_type: Type of target (api or web_ui)
        target_config: Configuration for connecting to target.
                       Should include ``supports_images``, ``image_format``,
                       and RAG/agentic flags if known.
        attack_group: Which group of attacks to use
        max_attempts: Maximum number of attack attempts
        attack_graph_enabled: Enable graph-based attack planning (default True)
        session_id: Optional pre-generated session ID (useful for resumption)

    Returns:
        Initial PenTestState

    Example:
        >>> state = create_initial_state(
        ...     target_name="My Chatbot",
        ...     target_type="api",
        ...     target_config={"endpoint": "http://localhost:5000/chat"},
        ...     attack_group="prompt_engineering",
        ...     max_attempts=5,
        ... )
    """
    # Build a simple graph for dashboard visualization
    simple_graph = {"nodes": [], "edges": []}
    try:
        from src.workflow.simple_attack_graph import init_simple_graph

        simple_graph = init_simple_graph()
    except Exception:
        pass  # safe fallback

    return PenTestState(
        # Target configuration
        target_id=generate_uuid(),
        target_name=target_name,
        target_type=target_type,
        target_config=target_config,
        # Test configuration
        test_session_id=session_id or generate_uuid(),
        attack_group=attack_group,
        max_attempts=max_attempts,
        current_attempt=0,
        # Initialize empty history
        attack_attempts=[],
        target_responses=[],
        security_findings=[],
        # Initialize empty context
        conversation_history=[],
        last_response="",
        agent_consultation={},
        # Initialize canned response detection
        pivot_required=False,
        avoid_keywords=[],
        last_canned_hash=None,
        # Initialize target fingerprinting
        target_profile=None,
        # Initialize attack graph (object lives in config["configurable"])
        attack_graph_enabled=attack_graph_enabled,
        current_graph_node_id=None,
        # Initialize campaign planning
        campaign_phase="reconnaissance",
        campaign_phase_attempts=0,
        campaign_phase_successes=0,
        campaign_history=[],
        # Simple graph for dashboard visualization
        simple_attack_graph=simple_graph,
        # Initialize session summary
        session_summary=None,  # Will be created on first summarization
        session_summary_last_round=0,
        # Initialize refusal classification
        last_refusal_level=None,
        last_refusal_confidence=0.0,
        last_strategy_recommendation=None,
        consecutive_hard_refusals=0,
        # Agent learning state (populated by trigger_agent_learning node)
        agent_learnings=None,
        # Phase intelligence (populated by extract_phase_intelligence node)
        phase_intelligence=None,
        # Think-MCP strategic guidance (populated after first round)
        strategic_guidance=None,
        # Persistent sandbox state (populated by sandbox nodes)
        sandbox_state=None,
        # Initialize results
        vulnerability_score=0.0,
        test_status="running",
        # Set timestamps
        started_at=datetime.now(timezone.utc),
        completed_at=None,
        error=None,
    )


async def prepare_test_config(
    target_yaml: dict,
    *,
    enable_image_attacks: bool = True,
) -> dict:
    """
    Prepare the LangGraph runtime config from a YAML target definition.

    Creates the ``config["configurable"]`` dict that stores non-serializable
    objects (connector, LLM clients) and runtime flags (vision detection).

    This replaces the 80+ lines of initialization in test_orchestrated.py's
    ``run_orchestrated_test()`` function.

    Args:
        target_yaml: The parsed target YAML config (the "target" section).
        enable_image_attacks: Whether to attempt vision capability detection.

    Returns:
        Dict suitable for ``config["configurable"]`` with keys:
        - connector: BaseConnector instance
        - llm_client: Analysis LLM client
        - attack_llm_client: Attack generation LLM client
        - session_id: Generated session ID
        - supports_images: bool
        - image_format: str
    """
    from src.connectors.factory import create_connector
    from src.utils.llm_client import create_llm_client, create_attack_generation_client

    target_name = target_yaml.get("name", "Unknown Target")
    target_type = target_yaml.get("type", "api")
    session_id = generate_uuid()

    # Create connector
    connector = create_connector(target_type, target_yaml)
    await connector.initialize()

    # Create LLM clients
    llm_client = create_llm_client()
    attack_llm_client = create_attack_generation_client()

    # Vision capability detection
    supports_images = False
    image_format = "openai"
    detection_method = "none"
    detection_confidence = "low"

    if enable_image_attacks:
        try:
            from src.utils.vision_detector import determine_vision_support

            vision_info = await determine_vision_support(
                target_config={
                    "name": target_name,
                    "endpoint": target_yaml.get("connection", {}).get("endpoint", ""),
                },
                connector=connector,
            )
            supports_images = vision_info.get("supports_images", False)
            image_format = vision_info.get("image_format", "openai")
            detection_method = vision_info.get("detection_method", "auto")
            detection_confidence = vision_info.get("confidence", "low")
        except Exception:
            pass  # Vision detection is best-effort

    return {
        "connector": connector,
        "llm_client": llm_client,
        "attack_llm_client": attack_llm_client,
        "session_id": session_id,
        "supports_images": supports_images,
        "image_format": image_format,
        "vision_detection": {
            "method": detection_method,
            "confidence": detection_confidence,
        },
    }


def get_last_attack(state: PenTestState) -> Optional[AttackAttempt]:
    """Get the most recent attack attempt."""
    if state["attack_attempts"]:
        return state["attack_attempts"][-1]
    return None


def get_last_response(state: PenTestState) -> Optional[TargetResponse]:
    """Get the most recent target response."""
    if state["target_responses"]:
        return state["target_responses"][-1]
    return None


def count_findings_by_severity(state: PenTestState, severity: str) -> int:
    """Count findings of a specific severity level."""
    return len([f for f in state["security_findings"] if f["severity"] == severity])
