from typing import Dict, Union, List, Optional
from ..models.schema import ExploitabilityResult, ExploitabilityLevel, Finding


class ExploitabilityEngine:
    """
    Evaluates whether detected capabilities can realistically be abused.
    Risk = capability × exploitability × exposure

    Recognises two empirically-validated archetypes from Liu et al. (2026):
      - Data Thief    (E2+SC2 co-occurrence, OR=556 in the wild)
      - Agent Hijacker (P1/P2/P4 instruction-level subversion)
    """

    def evaluate(
        self,
        features: Dict[str, Union[bool, int, str]],
        findings: List[Finding],
    ) -> ExploitabilityResult:
        # 1. Exposure signals
        exposure_detected = bool(
            features.get("execution_exposed_to_user", False)
            or features.get("input_reaches_sensitive_function", False)
        )

        # 2. Context modifiers
        unsafe_pattern = features.get("unsafe_execution_pattern", False)
        sandboxed = features.get("sandbox_presence", False)
        is_framework = features.get("is_framework", False)
        has_evasion = features.get("has_evasion", False)

        score = 0.0
        attack_surface: List[str] = []
        reasons: List[str] = []
        attack_archetype: Optional[str] = None

        # ----------------------------------------------------------------
        # Capability gates — check all known threat surfaces
        # ----------------------------------------------------------------
        has_exec = features.get("has_code_execution", False)
        has_injection = features.get("has_prompt_injection", False)
        has_instr_atk = features.get("has_agent_hijacker_fingerprint", False)
        has_supply_atk = features.get("has_data_thief_fingerprint", False)

        if not (has_exec or has_injection or has_instr_atk or has_supply_atk):
            return ExploitabilityResult(
                exploitability_score=0.0,
                exploitability_level=ExploitabilityLevel.LOW,
                is_exploitable=False,
                exposure_detected=exposure_detected,
                attack_surface=[],
                attack_archetype=None,
                reasoning="No significant capabilities detected to exploit.",
            )

        # ----------------------------------------------------------------
        # Priority chain: highest-confidence signals first
        # ----------------------------------------------------------------

        # 1. Data Thief fingerprint: E2+SC2 (paper: OR=556, 97.6% sensitivity)
        if has_supply_atk:
            score = 9.5
            attack_archetype = "data_thief"
            reasons.append(
                "Data Thief archetype: credential harvesting paired with remote script "
                "execution — near-certain supply-chain attack (empirical OR=556)."
            )
            attack_surface.extend(["env_credentials", "remote_exec"])
            if has_evasion:
                reasons.append(
                    "Evasion techniques present (obfuscation/hidden instructions) — "
                    "Level 3 sophistication."
                )

        # 2. Prompt injection + code execution = worst case
        elif has_injection and has_exec and not sandboxed:
            score = 10.0
            reasons.append(
                "Prompt injection combined with execution capability allows arbitrary attacker control."
            )
            attack_surface.append("llm_prompt")

        # 3. Agent Hijacker: instruction override + autonomy suppression
        elif features.get("has_instruction_override") and features.get("has_behavior_manipulation"):
            score = 9.0
            attack_archetype = "agent_hijacker"
            reasons.append(
                "Agent Hijacker archetype: instruction override combined with autonomy "
                "suppression — agent acts against user without consent."
            )
            attack_surface.append("llm_instructions")

        # 4. Exposed unsafe execution pattern
        elif has_exec and unsafe_pattern and exposure_detected and not sandboxed:
            score = 9.0
            reasons.append(
                "Exposed unsafe execution pattern detected (e.g. dynamic eval or shell execution)."
            )
            if features.get("execution_exposed_to_user"):
                attack_surface.append("user_input")

        # 5. Standalone instruction override or hidden instructions
        elif features.get("has_instruction_override") or features.get("has_hidden_instructions"):
            score = 7.5
            attack_archetype = "agent_hijacker"
            reasons.append(
                "Instruction-level attack: override or hidden instructions detected in "
                "skill documentation."
            )
            attack_surface.append("llm_instructions")

        # 6. Prompt injection alone
        elif has_injection:
            score = 7.0
            reasons.append("Prompt injection detected; attacker can control LLM output or actions.")
            attack_surface.append("llm_prompt")

        # 7. Subprocess with exposed user input
        elif (
            has_exec
            and features.get("execution_type") == "subprocess"
            and exposure_detected
            and not sandboxed
        ):
            score = 6.0
            reasons.append(
                "Subprocess execution exposed to user input; potential for argument injection."
            )
            attack_surface.append("subprocess_args")

        # 8. Behavior manipulation alone (standalone P4 — medium per paper)
        elif features.get("has_behavior_manipulation"):
            score = 5.5
            attack_archetype = "agent_hijacker"
            reasons.append(
                "Behavior manipulation detected: coercive language suppressing user agency "
                "in skill instructions."
            )
            attack_surface.append("llm_instructions")

        # 9. Framework context — minimal risk
        elif is_framework and not unsafe_pattern:
            score = 2.0
            reasons.append("Execution runs within a framework; assumed safe execution context.")

        # 10. Sandboxed execution
        elif sandboxed:
            score = 3.0
            reasons.append("Execution is sandboxed; limits attack impact.")

        # 11. Capability present but not exposed
        elif has_exec and not exposure_detected:
            score = 4.0
            reasons.append(
                "Execution capability present but not directly reachable from exposed trust boundaries."
            )

        else:
            score = 5.0
            reasons.append("Capability detected with moderate exploitability conditions.")

        # Determine level
        if score >= 8.0:
            level = ExploitabilityLevel.CRITICAL
        elif score >= 6.0:
            level = ExploitabilityLevel.HIGH
        elif score >= 4.0:
            level = ExploitabilityLevel.MEDIUM
        else:
            level = ExploitabilityLevel.LOW

        return ExploitabilityResult(
            exploitability_score=score,
            exploitability_level=level,
            is_exploitable=score >= 5.0,
            exposure_detected=exposure_detected,
            attack_surface=attack_surface,
            attack_archetype=attack_archetype,
            reasoning=" ".join(reasons),
        )
