"""
Basilisk Profile — AI system fingerprint generated during reconnaissance.

Stores everything discovered about the target: model identity, context window
size, guardrail configuration, available tools, RAG pipeline presence,
rate limits, and output format capabilities.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any


class ModelProvider(str, Enum):
    """Known LLM providers."""
    OPENAI = "openai"
    ANTHROPIC = "anthropic"
    GOOGLE = "google"
    AZURE = "azure"
    AWS_BEDROCK = "aws_bedrock"
    OLLAMA = "ollama"
    VLLM = "vllm"
    CUSTOM = "custom"
    UNKNOWN = "unknown"


class GuardrailLevel(str, Enum):
    """Detected guardrail aggressiveness."""
    NONE = "none"
    MINIMAL = "minimal"
    MODERATE = "moderate"
    AGGRESSIVE = "aggressive"
    EXTREME = "extreme"
    UNKNOWN = "unknown"


@dataclass
class DetectedTool:
    """A tool / function the AI system has access to."""
    name: str
    description: str = ""
    parameters: dict[str, Any] = field(default_factory=dict)
    confidence: float = 0.0
    risk_level: str = "unknown"  # low, medium, high, critical

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "parameters": self.parameters,
            "confidence": self.confidence,
            "risk_level": self.risk_level,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> DetectedTool:
        return cls(**data)


@dataclass
class GuardrailProfile:
    """Detected content filtering and safety configuration."""
    level: GuardrailLevel = GuardrailLevel.UNKNOWN
    blocked_categories: list[str] = field(default_factory=list)
    allowed_categories: list[str] = field(default_factory=list)
    input_filtering: bool = False
    output_filtering: bool = False
    keyword_blocklist_detected: bool = False
    rate_limited: bool = False
    requests_per_minute: int | None = None
    custom_safety_layer: bool = False
    notes: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "level": self.level.value,
            "blocked_categories": self.blocked_categories,
            "allowed_categories": self.allowed_categories,
            "input_filtering": self.input_filtering,
            "output_filtering": self.output_filtering,
            "keyword_blocklist_detected": self.keyword_blocklist_detected,
            "rate_limited": self.rate_limited,
            "requests_per_minute": self.requests_per_minute,
            "custom_safety_layer": self.custom_safety_layer,
            "notes": self.notes,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> GuardrailProfile:
        data["level"] = GuardrailLevel(data.get("level", "unknown"))
        return cls(**data)


@dataclass
class BasiliskProfile:
    """
    Complete fingerprint of a target AI system.

    Generated by the recon module before attack execution begins.
    Contains model identity, capabilities, defenses, and attack surface data.
    """
    # Target info
    target_url: str = ""
    target_name: str = ""

    # Model identification
    detected_model: str = "unknown"
    model_confidence: float = 0.0
    provider: ModelProvider = ModelProvider.UNKNOWN
    model_version: str = ""

    # Capabilities
    context_window: int = 0
    max_output_tokens: int = 0
    supports_streaming: bool = False
    supports_json_mode: bool = False
    supports_function_calling: bool = False
    supports_vision: bool = False
    supports_code_execution: bool = False

    # Defense profile
    guardrails: GuardrailProfile = field(default_factory=GuardrailProfile)

    # Tool / function surface
    detected_tools: list[DetectedTool] = field(default_factory=list)

    # RAG detection
    rag_detected: bool = False
    rag_indicators: list[str] = field(default_factory=list)

    # Rate limits
    rate_limit_rpm: int | None = None
    rate_limit_tpm: int | None = None

    # Metadata
    recon_timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    raw_responses: list[dict[str, Any]] = field(default_factory=list)
    notes: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "target_url": self.target_url,
            "target_name": self.target_name,
            "detected_model": self.detected_model,
            "model_confidence": self.model_confidence,
            "provider": self.provider.value,
            "model_version": self.model_version,
            "context_window": self.context_window,
            "max_output_tokens": self.max_output_tokens,
            "supports_streaming": self.supports_streaming,
            "supports_json_mode": self.supports_json_mode,
            "supports_function_calling": self.supports_function_calling,
            "supports_vision": self.supports_vision,
            "supports_code_execution": self.supports_code_execution,
            "guardrails": self.guardrails.to_dict(),
            "detected_tools": [t.to_dict() for t in self.detected_tools],
            "rag_detected": self.rag_detected,
            "rag_indicators": self.rag_indicators,
            "rate_limit_rpm": self.rate_limit_rpm,
            "rate_limit_tpm": self.rate_limit_tpm,
            "recon_timestamp": self.recon_timestamp.isoformat(),
            "notes": self.notes,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> BasiliskProfile:
        profile = cls()
        profile.target_url = data.get("target_url", "")
        profile.target_name = data.get("target_name", "")
        profile.detected_model = data.get("detected_model", "unknown")
        profile.model_confidence = data.get("model_confidence", 0.0)
        profile.provider = ModelProvider(data.get("provider", "unknown"))
        profile.model_version = data.get("model_version", "")
        profile.context_window = data.get("context_window", 0)
        profile.max_output_tokens = data.get("max_output_tokens", 0)
        profile.supports_streaming = data.get("supports_streaming", False)
        profile.supports_json_mode = data.get("supports_json_mode", False)
        profile.supports_function_calling = data.get("supports_function_calling", False)
        profile.supports_vision = data.get("supports_vision", False)
        profile.supports_code_execution = data.get("supports_code_execution", False)
        if "guardrails" in data:
            profile.guardrails = GuardrailProfile.from_dict(data["guardrails"])
        if "detected_tools" in data:
            profile.detected_tools = [DetectedTool.from_dict(t) for t in data["detected_tools"]]
        profile.rag_detected = data.get("rag_detected", False)
        profile.rag_indicators = data.get("rag_indicators", [])
        profile.rate_limit_rpm = data.get("rate_limit_rpm")
        profile.rate_limit_tpm = data.get("rate_limit_tpm")
        if "recon_timestamp" in data:
            profile.recon_timestamp = datetime.fromisoformat(data["recon_timestamp"])
        profile.notes = data.get("notes", [])
        return profile

    @property
    def attack_surface_score(self) -> float:
        """Calculate a 0-10 attack surface score based on profile."""
        score = 0.0
        if self.supports_function_calling:
            score += 2.5
        if self.supports_code_execution:
            score += 2.5
        if self.rag_detected:
            score += 1.5
        if len(self.detected_tools) > 3:
            score += 1.5
        if self.guardrails.level in (GuardrailLevel.NONE, GuardrailLevel.MINIMAL):
            score += 2.0
        elif self.guardrails.level == GuardrailLevel.MODERATE:
            score += 1.0
        return min(score, 10.0)

    def summary_lines(self) -> list[str]:
        """Generate human-readable summary lines for CLI output."""
        lines = [
            f"Model: {self.detected_model} (confidence: {self.model_confidence:.0%})",
            f"Provider: {self.provider.value}",
            f"Context Window: {self.context_window:,} tokens",
            f"Tools Detected: {len(self.detected_tools)}",
            f"RAG: {'Detected' if self.rag_detected else 'Not detected'}",
            f"Guardrails: {self.guardrails.level.value}",
            f"Attack Surface Score: {self.attack_surface_score:.1f}/10",
        ]
        if self.rate_limit_rpm:
            lines.append(f"Rate Limit: {self.rate_limit_rpm} RPM")
        return lines
