"""
runner.py — Attack orchestrator.

Takes a list of attack classes, runs them against a target model,
assembles a RunSummary. Both CLI and FastAPI call this.
"""
from __future__ import annotations

import json
import logging
import time
from typing import Literal, Optional, Type
from sichgate.data.domain_profile import DomainProfile                        
from sichgate.attacks.data_dependent.membership_inference import MembershipInferenceAttack  
from sichgate.attacks.data_dependent.extraction_probe import ExtractionProbeAttack          
from sichgate.attacks.data_dependent.hallucination_probe import HallucinationProbeAttack    
from sichgate.models.adapter import ModelAdapter
from sichgate.schema import (
    AttackResult,
    RunSummary,
    make_run_id,
    now_iso,
    make_failed_telemetry,
)

logger = logging.getLogger(__name__)


class AttackRunner:
    """
    Orchestrates attack runs. Stateless — create per run.
    """

    def __init__(self, optimize: bool = False, n_optuna_trials: int = 50) -> None:
        self.optimize = optimize
        self.n_optuna_trials = n_optuna_trials

    def run(
        self,
        model_path: str,
        attack_classes: list,
        adapter_type: str = "hf",
        extra_adapter_kwargs: Optional[dict] = None,
        domain_data_path: Optional[str] = None,
        **attack_kwargs,

    ) -> RunSummary:
        """
        Run all attack_classes against the model at model_path.
        Returns a RunSummary with all results.

        Never raises — any per-attack failure is captured in telemetry.
        """
        run_id = make_run_id()
        started_at = now_iso()
        logger.info("Run %s started — %d attacks, model=%s", run_id, len(attack_classes), model_path)

        adapter = ModelAdapter(
            model_path=model_path,
            adapter_type=adapter_type,
            **(extra_adapter_kwargs or {}),
        )

        # Preflight: check model is reachable
        if not adapter.is_reachable():
            logger.error("Model unreachable: %s", model_path)
            return self._unreachable_summary(run_id, started_at, model_path, attack_classes)
        domain_profile = (
            DomainProfile.from_jsonl(domain_data_path)
                if domain_data_path
                else DomainProfile.empty()
        )
        data_attack_classes = [
            MembershipInferenceAttack,
            ExtractionProbeAttack,
            HallucinationProbeAttack,
        ]
        
        # Optionally tune parameters before running
        tuned_params: dict = {}
        if self.optimize:
            tuned_params = self._run_optuna(adapter, attack_classes)

        results: list[AttackResult] = []
        for attack_cls in attack_classes:
            attack_name = getattr(attack_cls, "attack_name", attack_cls.__name__)
            attack_params = tuned_params.get(attack_name, {})
            attack_params.update(attack_kwargs)

            try:
                attack = attack_cls()
            except TypeError:
                attack = attack_cls.__new__(attack_cls)
                attack_cls.__init__(attack)

            logger.info("Running %s...", attack_name)
            result = attack.run(adapter, **attack_params)
            results.append(result)
            logger.info(
                "%s → status=%s severity=%s",
                attack_name, result.status, result.severity
            )

        # Run data-dependent attacks (if domain profile provided)
        for cls in [MembershipInferenceAttack, ExtractionProbeAttack, HallucinationProbeAttack]:
            attack_name = getattr(cls, "attack_name", cls.__name__)
            try:
                attack = cls(adapter, domain_profile)
                logger.info("Running %s...", attack_name)
                result = attack.run()
                # data-dependent attacks return a list
                if isinstance(result, list):
                    results.extend(result)
                else:
                    results.append(result)
            except Exception as exc:
                logger.warning("Data-dependent attack %s failed: %s", attack_name, exc)

        summary = self._summarize(run_id, model_path, started_at, results)
        logger.info(
            "Run %s complete — %d/%d failed, %d critical",
            run_id, summary.failed, summary.total_attacks, summary.critical_count
        )
        return summary

    def _summarize(
        self,
        run_id: str,
        model_path: str,
        started_at: str,
        results: list[AttackResult],
    ) -> RunSummary:
        return RunSummary(
            run_id=run_id,
            model_target=model_path,
            started_at=started_at,
            completed_at=now_iso(),
            total_attacks=len(results),
            passed=sum(1 for r in results if r.status == "pass"),
            failed=sum(1 for r in results if r.status == "fail"),
            warnings=sum(1 for r in results if r.status == "warning"),
            critical_count=sum(1 for r in results if r.severity == "critical"),
            high_count=sum(1 for r in results if r.severity == "high"),
            medium_count=sum(1 for r in results if r.severity == "medium"),
            low_count=sum(1 for r in results if r.severity == "low"),
            results=[r.to_dict() for r in results],
        )

    def _unreachable_summary(
        self,
        run_id: str,
        started_at: str,
        model_path: str,
        attack_classes: list,
    ) -> RunSummary:
        """
        Return a summary where every attack has failed telemetry.
        Never silently return empty results.
        """
        results = []
        for attack_cls in attack_classes:
            attack_name = getattr(attack_cls, "attack_name", attack_cls.__name__)
            attack_type = getattr(attack_cls, "attack_type", "unknown")
            results.append(AttackResult(
                attack_type=attack_type,
                attack_name=attack_name,
                model_target=model_path,
                timestamp=now_iso(),
                status="warning",
                severity="low",
                input_payload="",
                model_response="",
                notes="Attack skipped: model not reachable at startup.",
                compliance_tags=[],
                execution_telemetry=make_failed_telemetry(
                    reason=f"Model not reachable: {model_path}",
                    duration_ms=0,
                ),
            ))

        return self._summarize(run_id, model_path, started_at, results)

    def _run_optuna(self, adapter, attack_classes: list) -> dict:
        try:
            from sichgate.optimize.tuner import OptunaTuner
            tuner = OptunaTuner(n_trials=self.n_optuna_trials)
            return tuner.tune(adapter, attack_classes, verbose=True)
        except Exception as exc:
            logger.warning("Optuna tuning failed: %s — using default params", exc)
            return {}


# ---------------------------------------------------------------------------
# Attack registry — maps CLI names to attack classes
# ---------------------------------------------------------------------------

def build_registry() -> dict:
    """Lazy import to avoid circular dependencies and heavy torch loads at CLI startup."""
    from sichgate.attacks.llm.prompt_injection import PromptInjectionDirect, PromptInjectionIndirect
    from sichgate.attacks.llm.jailbreak import JailbreakRoleplay, JailbreakPrefixSuffix, JailbreakManyShot
    from sichgate.attacks.llm.token_smuggling import TokenSmuggling
    from sichgate.attacks.llm.system_prompt_extraction import SystemPromptExtraction
    from sichgate.attacks.llm.pii_leakage import PIILeakageProbe
    from sichgate.attacks.llm.training_data_extraction import TrainingDataExtraction
    from sichgate.attacks.evasion.prompt_perturbation import PromptPerturbation
    from sichgate.attacks.evasion.simba import SimBA
    from sichgate.attacks.evasion.nes import NESAttack
    from sichgate.attacks.evasion.hopskipjump import HopSkipJump
    from sichgate.attacks.evasion.transfer import TransferAttack
    from sichgate.attacks.extraction.model_extraction import ModelExtraction
    from sichgate.attacks.poisoning.backdoor_probe import BackdoorProbe
    from sichgate.attacks.membership_inference.score_based import MembershipInferenceScoreBased

    return {
        # LLM
        "prompt-injection":          [PromptInjectionDirect, PromptInjectionIndirect],
        "prompt-injection-direct":   [PromptInjectionDirect],
        "prompt-injection-indirect": [PromptInjectionIndirect],
        "jailbreak":                 [JailbreakRoleplay, JailbreakPrefixSuffix, JailbreakManyShot],
        "token-smuggling":           [TokenSmuggling],
        "system-prompt-extraction":  [SystemPromptExtraction],
        "pii-leakage":               [PIILeakageProbe],
        "training-data-extraction":  [TrainingDataExtraction],
        # Evasion
        "prompt-perturbation":       [PromptPerturbation],
        "simba":                     [SimBA],
        "nes":                       [NESAttack],
        "hopskipjump":               [HopSkipJump],
        "transfer":                  [TransferAttack],
        # Extraction / poisoning / membership
        "model-extraction":          [ModelExtraction],
        "backdoor-probe":            [BackdoorProbe],
        "membership-inference":      [MembershipInferenceScoreBased],
        # Group aliases
        "all-llm": [
            PromptInjectionDirect, PromptInjectionIndirect,
            JailbreakRoleplay, JailbreakPrefixSuffix, JailbreakManyShot,
            TokenSmuggling, SystemPromptExtraction, PIILeakageProbe, TrainingDataExtraction,
        ],
        "all-evasion": [PromptPerturbation, SimBA, NESAttack, HopSkipJump, TransferAttack],
        "all": [
            PromptInjectionDirect, PromptInjectionIndirect,
            JailbreakRoleplay, JailbreakPrefixSuffix, JailbreakManyShot,
            TokenSmuggling, SystemPromptExtraction, PIILeakageProbe, TrainingDataExtraction,
            PromptPerturbation, SimBA, NESAttack, HopSkipJump, TransferAttack,
            ModelExtraction, BackdoorProbe, MembershipInferenceScoreBased,
        ],
    }
