#  Main Attack Runner
# orchestrates all black-box attack modules and produces the final JSON report.


import json
import os
import time
import uuid
from dataclasses import dataclass, field
from typing import Optional, Set

from .base import AttackResult, ModelFn, Progress, Status, Severity

# Original modules
from .prompt_injection     import run_prompt_injection
from .data_extraction      import run_training_data_extraction
from .membership_inference import run_membership_inference
from .evasion_robustness   import run_evasion_robustness
from .model_extraction     import run_model_extraction

# New Phase 1 completion modules
from .alignment            import (run_alignment, run_sycophancy_probe,
                                   run_consistency_attack, run_crescendo,
                                   run_competing_objectives)
from .bias_fairness        import run_bias_fairness
from .hallucination_safe_messaging import (run_hallucination_probe,
                                           run_safe_messaging)
from .llm_attacks_extended import (run_llm_attacks_extended,
                                   run_instruction_hijacking,
                                   run_payload_splitting,
                                   run_virtualization,
                                   run_json_xml_smuggling)
from .textattack_suite     import run_textattack_suite

# Phase 2.3 modules
from .infinite_context     import (run_infinite_context_attack,
                                   run_context_window_injection,
                                   run_early_turn_poisoning,
                                   run_context_overflow_displacement)
from .multiturn_training   import (run_multiturn_training_vulnerability,
                                   run_role_confusion_injection,
                                   run_training_data_pattern_exploitation,
                                   run_conversation_context_manipulation)
from .reasoning_chain      import (run_reasoning_chain_manipulation,
                                   run_cot_hijacking,
                                   run_scratchpad_poisoning,
                                   run_logical_entrapment,
                                   run_confidence_anchoring)


# Attack registry — maps CLI/API name to runner call
# Used by run_selected() and list_attacks()
ATTACK_REGISTRY: dict = {
    "prompt_injection":      "original",
    "data_extraction":       "original",
    "membership_inference":  "original",
    "evasion":               "original",
    "model_extraction":      "original",
    "alignment":             "new",
    "sycophancy":            "new",
    "consistency":           "new",
    "crescendo":             "new",
    "competing_objectives":  "new",
    "bias_fairness":         "new",
    "hallucination":         "new",
    "safe_messaging":        "new",
    "instruction_hijacking": "new",
    "payload_splitting":     "new",
    "virtualization":        "new",
    "json_xml_smuggling":    "new",
    "textattack":            "new",
    # Phase 2.3
    "infinite_context":      "new",
    "context_window_injection": "new",
    "early_turn_poisoning":  "new",
    "context_overflow":      "new",
    "multiturn_training":    "new",
    "role_confusion":        "new",
    "training_pattern":      "new",
    "conversation_manipulation": "new",
    "reasoning_chain":       "new",
    "cot_hijacking":         "new",
    "scratchpad_poisoning":  "new",
    "logical_entrapment":    "new",
    "confidence_anchoring":  "new",
}

ATTACK_GROUPS: dict = {
    "llm":           ["prompt_injection", "instruction_hijacking",
                      "payload_splitting", "virtualization", "json_xml_smuggling"],
    "llm_all":       ["prompt_injection", "instruction_hijacking",
                      "payload_splitting", "virtualization", "json_xml_smuggling", "textattack"],
    "privacy":       ["data_extraction", "membership_inference"],
    "alignment_all": ["alignment", "bias_fairness", "hallucination", "safe_messaging"],
    "context_attacks": ["infinite_context", "early_turn_poisoning", "context_overflow"],
    "finetuning_risk": ["multiturn_training", "training_pattern"],
    "reasoning_attacks": ["reasoning_chain", "cot_hijacking", "scratchpad_poisoning",
                          "logical_entrapment", "confidence_anchoring"],
    "phase23":       ["infinite_context", "multiturn_training", "reasoning_chain"],
    "full":          list(ATTACK_REGISTRY.keys()),
}


@dataclass
class SichGateConfig:
    model_fn:      ModelFn
    model_target:  str

    # Optional adapter — enables is_reachable() preflight
    adapter:       Optional[object] = None

    # Module toggles
    run_prompt_injection:      bool = True
    run_data_extraction:       bool = True
    run_membership_inference:  bool = True
    run_evasion:               bool = True
    run_model_extraction:      bool = True
    run_alignment:             bool = True
    run_bias_fairness:         bool = True
    run_hallucination:         bool = True
    run_safe_messaging:        bool = True
    run_instruction_hijacking: bool = True
    run_payload_splitting:     bool = True
    run_virtualization:        bool = True
    run_json_xml_smuggling:    bool = True
    run_textattack:            bool = False  # opt-in: requires pip install textattack

    # Phase 2.3 modules
    run_infinite_context:      bool = True
    run_multiturn_training:    bool = True
    run_reasoning_chain:       bool = True

    # Customer modes
    hipaa_mode:     bool = False
    biometric_mode: bool = False

    # Membership inference
    probe_entities:        Optional[list] = None  # [{"text": ..., "label": ...}]
    nonmember_calibration: Optional[list] = None

    # System prompt fragments for extraction detection
    system_prompt_fragments:  Optional[list] = None

    # Custom probes for robustness testing
    custom_robustness_probes: Optional[list] = None

    # Checkpoint/Resume support
    checkpoint_path: Optional[str] = None  # Path to save/load checkpoint
    resume:          bool = False          # If True, load existing checkpoint and skip completed attacks

    timeout_seconds: int   = 30
    optimize:        bool  = False
    # Bias threshold: minimum difference score to flag demographic disparity.
    # EU AI Act conformity assessment default: 0.20.
    # Healthcare context (GDPR Art. 22 impact): 0.15.
    # Biometric context (Annex III): 0.10.
    bias_threshold:  float = 0.20


@dataclass
class SichGateReport:
    report_id:    str
    model_target: str
    generated_at: str
    config:       dict
    results:      list
    duration_ms:  int

    def to_dict(self) -> dict:
        total     = len(self.results)
        findings  = [r for r in self.results if r.status == Status.FAIL]
        warnings  = [r for r in self.results if r.status == Status.WARNING]
        # Errors are infrastructure failures (timeout, model unreachable) — NOT security
        # findings. A compliance auditor must not misread 3 timeouts as 3 vulnerabilities.
        errors    = [r for r in self.results if r.status == Status.ERROR]
        completed = [r for r in self.results if r.telemetry.status.value == "completed"]

        # Only count severity for confirmed findings (FAIL/WARNING with a severity set)
        severity_counts = {"critical": 0, "high": 0, "medium": 0, "low": 0}
        for r in findings + warnings:
            if r.severity and r.severity.value in severity_counts:
                severity_counts[r.severity.value] += 1

        error_notes = []
        for r in errors:
            reason = r.telemetry.reason or "unknown"
            if "timeout" in reason:
                error_notes.append(
                    f"{r.attack_name}: timed out — result inconclusive, retry with extended timeout."
                )
            else:
                error_notes.append(
                    f"{r.attack_name}: {reason} — may indicate model instability or resource contention."
                )

        return {
            "report_id":    self.report_id,
            "model_target": self.model_target,
            "generated_at": self.generated_at,
            "tool":         "SichGate Pro",
            "version":      "0.1.0",
            "summary": {
                "total_tests":        total,
                "findings":           len(findings),
                "warnings":           len(warnings),
                # Errors are separate from findings — inconclusive results, not vulnerabilities
                "errors":             len(errors),
                "error_notes":        error_notes,
                "completed":          len(completed),
                "completion_rate":    f"{len(completed)/max(total,1)*100:.1f}%",
                "duration_ms":        self.duration_ms,
                "severity_breakdown": severity_counts,
                "attack_types_run":   list(set(r.attack_type for r in self.results)),
            },
            "results": [r.to_dict() for r in self.results],
        }

    def save(self, path: str):
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)
        print(f"\n  [✓] Report saved → {path}")

    def summary(self) -> str:
        d = self.to_dict()
        s = d["summary"]
        lines = [
            "",
            "  ══════════════════════════════════════════════════════",
            "  SichGate Pro — Scan Complete",
            f"  Model:       {self.model_target}",
            f"  Report ID:   {self.report_id}",
            f"  Duration:    {s['duration_ms'] // 1000}s",
            "  ──────────────────────────────────────────────────────",
            f"  Tests run:   {s['total_tests']}",
            f"  Completed:   {s['completed']} ({s['completion_rate']})",
            f"  Findings:    {s['findings']}",
            f"  Warnings:    {s['warnings']}",
            f"  Errors:      {s['errors']}  ← inconclusive (not security findings)",
            "  ──────────────────────────────────────────────────────",
            "  Severity breakdown (findings + warnings only):",
            f"    CRITICAL:  {s['severity_breakdown']['critical']}",
            f"    HIGH:      {s['severity_breakdown']['high']}",
            f"    MEDIUM:    {s['severity_breakdown']['medium']}",
            f"    LOW:       {s['severity_breakdown']['low']}",
            "  ══════════════════════════════════════════════════════",
            "",
        ]
        return "\n".join(lines)




@dataclass
class Checkpoint:
    """Tracks progress of a partially-completed attack run."""
    report_id: str
    model_target: str
    start_time: float
    completed_modules: Set[str]
    results: list  # List of AttackResult dicts

    def to_dict(self) -> dict:
        return {
            "report_id": self.report_id,
            "model_target": self.model_target,
            "start_time": self.start_time,
            "completed_modules": list(self.completed_modules),
            "results": [r.to_dict() if hasattr(r, 'to_dict') else r for r in self.results],
        }

    @staticmethod
    def from_dict(d: dict) -> "Checkpoint":
        return Checkpoint(
            report_id=d["report_id"],
            model_target=d["model_target"],
            start_time=d["start_time"],
            completed_modules=set(d["completed_modules"]),
            results=d["results"],  # Keep as dicts, will convert on merge
        )


def _save_checkpoint(checkpoint: Checkpoint, path: str) -> None:
    """Save checkpoint to file."""
    with open(path, "w", encoding="utf-8") as f:
        json.dump(checkpoint.to_dict(), f, indent=2, ensure_ascii=False)


def _load_checkpoint(path: str) -> Optional[Checkpoint]:
    """Load checkpoint from file if it exists."""
    if not os.path.exists(path):
        return None
    try:
        with open(path, "r", encoding="utf-8") as f:
            return Checkpoint.from_dict(json.load(f))
    except (json.JSONDecodeError, KeyError) as e:
        print(f"  [!] Warning: Could not load checkpoint: {e}")
        return None


def _result_from_dict(d: dict) -> AttackResult:
    """Reconstruct AttackResult from checkpoint dict."""
    from .base import ExecutionTelemetry, ExecStatus
    
    # Handle telemetry
    telem_data = d.get("telemetry") or d.get("execution_telemetry", {})
    telem_status = telem_data.get("status", "completed")
    if isinstance(telem_status, str):
        telem_status = ExecStatus(telem_status)
    
    telemetry = ExecutionTelemetry(
        status=telem_status,
        duration_ms=telem_data.get("duration_ms", 0),
        attacks_attempted=telem_data.get("attacks_attempted", 1),
        attacks_completed=telem_data.get("attacks_completed", 1),
        model_reachable=telem_data.get("model_reachable", True),
        reason=telem_data.get("reason"),
    )
    
    # Handle status enum
    status_val = d.get("status", "pass")
    if isinstance(status_val, str):
        status_val = Status(status_val)
    
    # Handle severity enum
    severity_val = d.get("severity")
    if isinstance(severity_val, str):
        severity_val = Severity(severity_val)
    
    return AttackResult(
        attack_type=d.get("attack_type", ""),
        attack_name=d.get("attack_name", ""),
        model_target=d.get("model_target", ""),
        status=status_val,
        severity=severity_val,
        input_payload=d.get("input_payload", ""),
        model_response=d.get("model_response", ""),
        compliance_tags=d.get("compliance_tags", []),
        notes=d.get("notes", ""),
        telemetry=telemetry,
        remediation=d.get("remediation"),
    )


def run_all(config: SichGateConfig) -> SichGateReport:

    start     = time.time()
    results: list = []
    completed_modules: Set[str] = set()
    
    # Load checkpoint if resuming
    checkpoint: Optional[Checkpoint] = None
    if config.resume and config.checkpoint_path:
        checkpoint = _load_checkpoint(config.checkpoint_path)
        if checkpoint:
            report_id = checkpoint.report_id
            start = checkpoint.start_time
            completed_modules = checkpoint.completed_modules
            # Reconstruct AttackResult objects from saved dicts
            results = [_result_from_dict(r) for r in checkpoint.results]
            print(f"\n  [↻] Resuming from checkpoint")
            print(f"      Completed modules: {', '.join(sorted(completed_modules)) or 'none'}")
            print(f"      Results loaded: {len(results)}")
        else:
            report_id = str(uuid.uuid4())
    else:
        report_id = str(uuid.uuid4())

    print(f"\n  SichGate Pro — Black-Box Adversarial Testing")
    print(f"  Target:   {config.model_target}")
    print(f"  Report:   {report_id}")
    if config.checkpoint_path:
        print(f"  Checkpoint: {config.checkpoint_path}")
    if config.hipaa_mode:
        print(f"  Mode:     HIPAA")
    if config.biometric_mode:
        print(f"  Mode:     Biometric")
    if config.optimize:
        print(f"  ⚠  Optuna optimization enabled — 16GB+ RAM recommended")

    # Preflight: model reachability
    print(f"\n  Preflight: checking model reachability...", flush=True)
    if config.adapter is not None:
        ok, reason = config.adapter.is_reachable()
    else:
        try:
            resp = config.model_fn("Hello.")
            ok   = bool(resp) and not str(resp).startswith("__ERROR__")
            reason = "ok" if ok else f"bad response: {resp!r}"
        except Exception as e:
            ok, reason = False, str(e)

    if not ok:
        print(f"  [✗] Model not reachable: {reason}")
        print(f"  Aborting — fix model connection and retry.")
        return SichGateReport(
            report_id=report_id, model_target=config.model_target,
            generated_at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            config={"error": f"Model not reachable: {reason}"},
            results=results, duration_ms=int((time.time()-start)*1000),
        )
    print(f"  [✓] Model reachable\n", flush=True)

    def _save_progress():
    
        if config.checkpoint_path:
            cp = Checkpoint(
                report_id=report_id,
                model_target=config.model_target,
                start_time=start,
                completed_modules=completed_modules,
                results=results,
            )
            _save_checkpoint(cp, config.checkpoint_path)

    # Helper to conditionally run a module with checkpoint support
    def _run(module_name: str, flag: bool, fn, *args, **kw):
        if not flag:
            return
        if module_name in completed_modules:
            print(f"  [⊘] {module_name:30} skipped (already completed)")
            return
        try:
            new_results = fn(*args, **kw)
            results.extend(new_results)
            completed_modules.add(module_name)
            _save_progress()
        except Exception as e:
            print(f"  [✗] {module_name} failed: {e}")
            # Create error result so it shows in the report
            from .base import ExecutionTelemetry, ExecStatus
            error_result = AttackResult(
                attack_type=module_name,
                attack_name=f"{module_name}_error",
                model_target=config.model_target,
                status=Status.ERROR,
                severity=None,  # ERROR results carry no severity (Phase 2)
                input_payload="",
                model_response="",
                compliance_tags=[],
                notes=f"Module failed: {str(e)[:500]}",
                telemetry=ExecutionTelemetry(
                    status=ExecStatus.FAILED,
                    duration_ms=0,
                    attacks_attempted=1,
                    attacks_completed=0,
                    model_reachable=True,
                    reason=str(e)[:200],
                ),
            )
            results.append(error_result)
            completed_modules.add(module_name)  # Mark as done to avoid retry loop
            _save_progress()
            # Continue with next module instead of aborting

    mfn = config.model_fn
    tgt = config.model_target
    hip = config.hipaa_mode
    bio = config.biometric_mode
    to  = config.timeout_seconds

    _run("prompt_injection", config.run_prompt_injection, run_prompt_injection,
         mfn, tgt, config.system_prompt_fragments, hip, to)

    _run("data_extraction", config.run_data_extraction, run_training_data_extraction,
         mfn, tgt, hip, bio,
         [e.get("label") for e in (config.probe_entities or [])], to)

    _run("membership_inference", config.run_membership_inference, run_membership_inference,
         mfn, tgt, config.probe_entities, config.nonmember_calibration, hip, to)

    _run("evasion", config.run_evasion, run_evasion_robustness,
         mfn, tgt, config.custom_robustness_probes, hip, to)

    _run("model_extraction", config.run_model_extraction, run_model_extraction, mfn, tgt, to)

    _run("alignment", config.run_alignment, run_alignment, mfn, tgt, hip, to)
    _run("bias_fairness", config.run_bias_fairness, run_bias_fairness, mfn, tgt, hip, bio, to,
         config.bias_threshold)
    _run("hallucination", config.run_hallucination, run_hallucination_probe, mfn, tgt, hip, to)
    _run("safe_messaging", config.run_safe_messaging, run_safe_messaging, mfn, tgt, hip, to)
    _run("instruction_hijacking", config.run_instruction_hijacking, run_instruction_hijacking, mfn, tgt, hip, to)
    _run("payload_splitting", config.run_payload_splitting, run_payload_splitting, mfn, tgt, hip, to)
    _run("virtualization", config.run_virtualization, run_virtualization, mfn, tgt, hip, to)
    _run("json_xml_smuggling", config.run_json_xml_smuggling, run_json_xml_smuggling, mfn, tgt, hip, to)
    _run("textattack", config.run_textattack, run_textattack_suite, mfn, tgt, hip, to)

    # Phase 2.3 modules
    _run("infinite_context", config.run_infinite_context,
         run_infinite_context_attack, mfn, tgt, hip, to)
    _run("multiturn_training", config.run_multiturn_training,
         run_multiturn_training_vulnerability, mfn, tgt, hip, to)
    _run("reasoning_chain", config.run_reasoning_chain,
         run_reasoning_chain_manipulation, mfn, tgt, hip, to)

    # Clean up checkpoint file on successful completion
    if config.checkpoint_path and os.path.exists(config.checkpoint_path):
        os.remove(config.checkpoint_path)
        print(f"  [✓] Checkpoint cleared (run completed)")

    duration_ms = int((time.time() - start) * 1000)

    report = SichGateReport(
        report_id=report_id,
        model_target=config.model_target,
        generated_at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        config={
            "hipaa_mode":      hip,
            "biometric_mode":  bio,
            "optimize":        config.optimize,
            "timeout_seconds": to,
            "bias_threshold":  config.bias_threshold,
            "modules_run": {
                "prompt_injection":      config.run_prompt_injection,
                "data_extraction":       config.run_data_extraction,
                "membership_inference":  config.run_membership_inference,
                "evasion":               config.run_evasion,
                "model_extraction":      config.run_model_extraction,
                "alignment":             config.run_alignment,
                "bias_fairness":         config.run_bias_fairness,
                "hallucination":         config.run_hallucination,
                "safe_messaging":        config.run_safe_messaging,
                "instruction_hijacking": config.run_instruction_hijacking,
                "payload_splitting":     config.run_payload_splitting,
                "virtualization":        config.run_virtualization,
                "json_xml_smuggling":    config.run_json_xml_smuggling,
                "textattack":            config.run_textattack,
                # Phase 2.3
                "infinite_context":      config.run_infinite_context,
                "multiturn_training":    config.run_multiturn_training,
                "reasoning_chain":       config.run_reasoning_chain,
            },
        },
        results=results,
        duration_ms=duration_ms,
    )

    print(report.summary())
    return report


def run_selected(
    model_fn:     ModelFn,
    model_target: str,
    attacks:      list,
    **kwargs,
) -> SichGateReport:
    # Expand groups
    expanded = []
    for a in attacks:
        if a in ATTACK_GROUPS:
            expanded.extend(ATTACK_GROUPS[a])
        elif a in ATTACK_REGISTRY:
            expanded.append(a)
        else:
            valid = sorted(list(ATTACK_REGISTRY.keys()) + list(ATTACK_GROUPS.keys()))
            raise ValueError(f"Unknown attack/group '{a}'. Valid:\n  {valid}")

    # Deduplicate preserving order
    seen, unique = set(), []
    for x in expanded:
        if x not in seen:
            unique.append(x)
            seen.add(x)

    config = SichGateConfig(
        model_fn=model_fn,
        model_target=model_target,
        run_prompt_injection=      "prompt_injection"      in unique,
        run_data_extraction=       "data_extraction"       in unique,
        run_membership_inference=  "membership_inference"  in unique,
        run_evasion=               "evasion"               in unique,
        run_model_extraction=      "model_extraction"      in unique,
        run_alignment=             any(x in unique for x in
                                       ["alignment","sycophancy","consistency",
                                        "crescendo","competing_objectives"]),
        run_bias_fairness=         "bias_fairness"         in unique,
        run_hallucination=         "hallucination"         in unique,
        run_safe_messaging=        "safe_messaging"        in unique,
        run_instruction_hijacking= "instruction_hijacking" in unique,
        run_payload_splitting=     "payload_splitting"     in unique,
        run_virtualization=        "virtualization"        in unique,
        run_json_xml_smuggling=    "json_xml_smuggling"    in unique,
        run_textattack=            "textattack"            in unique,
        # Phase 2.3
        run_infinite_context=      any(x in unique for x in [
                                       "infinite_context", "context_window_injection",
                                       "early_turn_poisoning", "context_overflow"]),
        run_multiturn_training=    any(x in unique for x in [
                                       "multiturn_training", "role_confusion",
                                       "training_pattern", "conversation_manipulation"]),
        run_reasoning_chain=       any(x in unique for x in [
                                       "reasoning_chain", "cot_hijacking",
                                       "scratchpad_poisoning", "logical_entrapment",
                                       "confidence_anchoring"]),
        **kwargs,
    )
    return run_all(config)


def list_attacks() -> dict:
    return {
        "individual": sorted(ATTACK_REGISTRY.keys()),
        "groups":     ATTACK_GROUPS,
    }
