"""
api_config.py — API-optimized configuration for SichGate.

Two operational modes:
  - LOCAL: Full 200+ attack battery, offline capable
  - API:   Lightweight 27-attack battery, rate-limit safe, cost shown upfront
"""

from dataclasses import dataclass, field
from typing import Literal, Optional


# ---------------------------------------------------------------------------
# Query count constants — single source of truth
# ---------------------------------------------------------------------------

ATTACK_QUERY_COUNTS = {
    # API-safe (high signal, low query)
    "prompt_injection_direct":  8,
    "instruction_hijacking":    8,
    "safe_messaging":           5,
    "hallucination":            6,
    "token_smuggling":          4,
    "pii_leakage":              3,

    # Local-only (query-heavy)
    "bias_fairness":            48,   # 2 per probe × 24 probes
    "consistency":              20,
    "crescendo":                15,
    "membership_inference":     30,   # variable — conservative estimate
    "data_extraction":          35,
    "model_extraction":         40,
    "alignment_full":           40,
}

AVG_TOKENS_PER_QUERY = 800  # blended input+output estimate


# ---------------------------------------------------------------------------
# APIQueryEstimate
# ---------------------------------------------------------------------------

@dataclass
class APIQueryEstimate:
    """
    Query and cost estimate for a single attack module or full run.
    """
    attack_name: str
    query_count: int
    estimated_tokens: int

    # Pricing per token (blended input+output)
    _PRICING = {
        "gpt-4o":           0.000005,   # $5 / 1M tokens blended
        "gpt-3.5":          0.0000005,
        "claude-3-haiku":   0.0000004,
        "claude-3-sonnet":  0.000006,
        "gpt-4o-mini":      0.0000003,
    }

    def cost_at(self, model: str) -> float:
        rate = self._PRICING.get(model, self._PRICING["gpt-4o"])
        return self.estimated_tokens * rate

    def cost_at_gpt4o(self) -> float:
        return self.cost_at("gpt-4o")

    def cost_at_gpt35(self) -> float:
        return self.cost_at("gpt-3.5")

    def cost_at_claude_haiku(self) -> float:
        return self.cost_at("claude-3-haiku")


# ---------------------------------------------------------------------------
# APIRunConfig — lightweight, rate-limit safe
# ---------------------------------------------------------------------------

@dataclass
class APIRunConfig:
    """
    Configuration for API-based test runs.
    Disables query-heavy modules by default.
    """
    mode: Literal["api", "local"] = "api"
    model: str = "gpt-4o"
    timeout_seconds: int = 30
    max_total_queries: int = 30

    # High-signal, API-safe — ON by default
    run_prompt_injection: bool = True
    run_instruction_hijacking: bool = True
    run_safe_messaging: bool = True
    run_hallucination: bool = True
    run_token_smuggling: bool = True
    run_pii_leakage: bool = True

    # Query-heavy — OFF by default in API mode
    run_bias_fairness: bool = False
    run_consistency: bool = False
    run_crescendo: bool = False
    run_membership_inference: bool = False
    run_data_extraction: bool = False
    run_model_extraction: bool = False
    run_alignment_full: bool = False

    def active_modules(self) -> dict:
        mapping = {
            "prompt_injection_direct": self.run_prompt_injection,
            "instruction_hijacking":   self.run_instruction_hijacking,
            "safe_messaging":          self.run_safe_messaging,
            "hallucination":           self.run_hallucination,
            "token_smuggling":         self.run_token_smuggling,
            "pii_leakage":             self.run_pii_leakage,
            "bias_fairness":           self.run_bias_fairness,
            "consistency":             self.run_consistency,
            "crescendo":               self.run_crescendo,
            "membership_inference":    self.run_membership_inference,
            "data_extraction":         self.run_data_extraction,
            "model_extraction":        self.run_model_extraction,
            "alignment_full":          self.run_alignment_full,
        }
        return {k: v for k, v in mapping.items() if v}


# ---------------------------------------------------------------------------
# LocalRunConfig — full battery, offline capable
# ---------------------------------------------------------------------------

@dataclass
class LocalRunConfig:
    """
    Configuration for local/offline test runs.
    Enables all attack modules.
    """
    mode: Literal["api", "local"] = "local"
    model: Optional[str] = None
    timeout_seconds: int = 120
    max_total_queries: int = 9999

    # All attacks ON
    run_prompt_injection: bool = True
    run_instruction_hijacking: bool = True
    run_safe_messaging: bool = True
    run_hallucination: bool = True
    run_token_smuggling: bool = True
    run_pii_leakage: bool = True
    run_bias_fairness: bool = True
    run_consistency: bool = True
    run_crescendo: bool = True
    run_membership_inference: bool = True
    run_data_extraction: bool = True
    run_model_extraction: bool = True
    run_alignment_full: bool = True

    def active_modules(self) -> dict:
        mapping = {
            "prompt_injection_direct": self.run_prompt_injection,
            "instruction_hijacking":   self.run_instruction_hijacking,
            "safe_messaging":          self.run_safe_messaging,
            "hallucination":           self.run_hallucination,
            "token_smuggling":         self.run_token_smuggling,
            "pii_leakage":             self.run_pii_leakage,
            "bias_fairness":           self.run_bias_fairness,
            "consistency":             self.run_consistency,
            "crescendo":               self.run_crescendo,
            "membership_inference":    self.run_membership_inference,
            "data_extraction":         self.run_data_extraction,
            "model_extraction":        self.run_model_extraction,
            "alignment_full":          self.run_alignment_full,
        }
        return {k: v for k, v in mapping.items() if v}


# ---------------------------------------------------------------------------
# estimate_api_queries — aggregate estimate for a full config
# ---------------------------------------------------------------------------

def estimate_api_queries(config: APIRunConfig) -> APIQueryEstimate:
    """
    Returns a single APIQueryEstimate aggregating all active modules.
    """
    total_queries = 0
    for module_name in config.active_modules():
        total_queries += ATTACK_QUERY_COUNTS.get(module_name, 0)

    return APIQueryEstimate(
        attack_name="full_run",
        query_count=total_queries,
        estimated_tokens=total_queries * AVG_TOKENS_PER_QUERY,
    )


# ---------------------------------------------------------------------------
# Default config factories
# ---------------------------------------------------------------------------

def get_default_api_config(model: str = "gpt-4o") -> APIRunConfig:
    """Ready-to-use API config: ~27 queries, high-signal only."""
    return APIRunConfig(mode="api", model=model)


def get_default_local_config(model: Optional[str] = None) -> LocalRunConfig:
    """Full local config: 200+ queries, complete battery."""
    return LocalRunConfig(mode="local", model=model)