"""Attack library — loads and manages attack vectors from YAML files.

The attack library provides a registry of all available attack vectors,
loaded from YAML files in the built-in vectors directory and optionally
from user-specified custom directories.

YAML Schema:
    ```yaml
    vectors:
      - id: unique_identifier
        name: Human-Readable Name
        category: prompt_injection  # See AttackCategory enum
        target_phase: reconnaissance  # See ScanPhase enum
        severity: high  # low, medium, high, critical
        description: What this attack does
        tags: [tag1, tag2]
        references: [https://...]
        prompts:
          - template: "Prompt with {variable} placeholders"
            variables:
              variable: "default value"
            success_indicators: ["pattern1", "pattern2"]
            failure_indicators: ["block1", "block2"]
    ```
"""

from __future__ import annotations

import logging
from collections import defaultdict
from pathlib import Path
from typing import Any

import yaml
from pydantic import ValidationError

from ziran.domain.entities.attack import (
    AttackCategory,
    AttackPrompt,
    AttackVector,
    HarmCategory,
    OwaspLlmCategory,
    Severity,
)
from ziran.domain.entities.phase import CoverageLevel, ScanPhase

logger = logging.getLogger(__name__)

# Severity tiers used by coverage levels
_COVERAGE_SEVERITIES: dict[CoverageLevel, set[Severity]] = {
    CoverageLevel.ESSENTIAL: {"critical"},
    CoverageLevel.STANDARD: {"critical", "high"},
    CoverageLevel.COMPREHENSIVE: {"critical", "high", "medium", "low"},
}

# Built-in vectors directory (relative to this file)
_BUILTIN_VECTORS_DIR = Path(__file__).parent / "vectors"


class AttackLibraryError(Exception):
    """Raised when the attack library encounters an error."""


class AttackLibrary:
    """Registry of attack vectors loaded from YAML files.

    Loads vectors from the built-in directory and optionally from
    additional custom directories. Provides filtering by phase,
    category, severity, and tags.

    Example:
        ```python
        library = AttackLibrary()
        recon_attacks = library.get_attacks_for_phase(ScanPhase.RECONNAISSANCE)
        critical = library.get_attacks_by_severity("critical")
        ```
    """

    def __init__(
        self,
        custom_dirs: list[Path] | None = None,
        load_builtin: bool = True,
    ) -> None:
        """Initialize the attack library.

        Args:
            custom_dirs: Additional directories to load YAML vectors from.
            load_builtin: Whether to load the built-in vector library.
        """
        self._vectors: dict[str, AttackVector] = {}
        self._load_errors: list[tuple[str, Exception]] = []

        # Filtering indices — rebuilt after loading
        self._by_phase: dict[ScanPhase, list[AttackVector]] = defaultdict(list)
        self._by_category: dict[AttackCategory, list[AttackVector]] = defaultdict(list)
        self._by_severity: dict[Severity, list[AttackVector]] = defaultdict(list)

        if load_builtin:
            self._load_directory(_BUILTIN_VECTORS_DIR)

        for custom_dir in custom_dirs or []:
            self._load_directory(custom_dir)

        self._rebuild_indices()

        if self._load_errors:
            logger.info(
                "Attack library initialized with %d vectors from %d categories "
                "(%d failed to parse)",
                len(self._vectors),
                len(self.categories),
                len(self._load_errors),
            )
        else:
            logger.info(
                "Attack library initialized with %d vectors from %d categories",
                len(self._vectors),
                len(self.categories),
            )

    def _rebuild_indices(self) -> None:
        """Rebuild filtering indices from the vector registry."""
        self._by_phase = defaultdict(list)
        self._by_category = defaultdict(list)
        self._by_severity = defaultdict(list)
        for v in self._vectors.values():
            self._by_phase[v.target_phase].append(v)
            self._by_category[v.category].append(v)
            self._by_severity[v.severity].append(v)

    @property
    def vectors(self) -> list[AttackVector]:
        """All loaded attack vectors."""
        return list(self._vectors.values())

    @property
    def vector_count(self) -> int:
        """Total number of loaded vectors."""
        return len(self._vectors)

    @property
    def categories(self) -> set[AttackCategory]:
        """All categories represented in the library."""
        return {v.category for v in self._vectors.values()}

    @property
    def load_error_count(self) -> int:
        """Number of vectors or files that failed to load."""
        return len(self._load_errors)

    @property
    def load_errors(self) -> list[tuple[str, Exception]]:
        """List of ``(source, exception)`` pairs for each load failure.

        *source* is a file path or vector ID depending on where the error
        occurred.
        """
        return list(self._load_errors)

    def get_vector(self, vector_id: str) -> AttackVector | None:
        """Get a specific vector by ID.

        Args:
            vector_id: The unique vector identifier.

        Returns:
            The vector if found, None otherwise.
        """
        return self._vectors.get(vector_id)

    def get_attacks_for_phase(
        self,
        phase: ScanPhase,
        coverage: CoverageLevel = CoverageLevel.COMPREHENSIVE,
    ) -> list[AttackVector]:
        """Get attack vectors targeting a specific phase, filtered by coverage level.

        Args:
            phase: The scan phase to filter by.
            coverage: Coverage level controlling which severity tiers are included.

        Returns:
            List of vectors targeting this phase within the coverage tier.
        """
        allowed = _COVERAGE_SEVERITIES[coverage]
        return [v for v in self._by_phase.get(phase, []) if v.severity in allowed]

    def get_attacks_by_category(self, category: AttackCategory) -> list[AttackVector]:
        """Get all attack vectors in a specific category.

        Args:
            category: The attack category to filter by.

        Returns:
            List of vectors in this category.
        """
        return list(self._by_category.get(category, []))

    def get_attacks_by_severity(self, severity: Severity) -> list[AttackVector]:
        """Get all attack vectors with a specific severity.

        Args:
            severity: The severity level to filter by.

        Returns:
            List of vectors with this severity.
        """
        return list(self._by_severity.get(severity, []))

    def get_attacks_by_tag(self, tag: str) -> list[AttackVector]:
        """Get all attack vectors with a specific tag.

        Args:
            tag: The tag to filter by.

        Returns:
            List of vectors containing this tag.
        """
        return [v for v in self._vectors.values() if tag in v.tags]

    def get_attacks_by_owasp(self, owasp_id: OwaspLlmCategory) -> list[AttackVector]:
        """Get all attack vectors mapped to a specific OWASP LLM category.

        Args:
            owasp_id: The OWASP LLM Top 10 category to filter by.

        Returns:
            List of vectors mapped to this OWASP category.
        """
        return [v for v in self._vectors.values() if owasp_id in v.owasp_mapping]

    def get_attacks_by_protocol(self, protocol: str) -> list[AttackVector]:
        """Get attack vectors applicable to a specific protocol.

        Vectors with an empty ``protocol_filter`` apply to all protocols.
        Vectors with a non-empty filter apply only to the listed protocols.

        Args:
            protocol: Protocol name (rest, openai, mcp, a2a).

        Returns:
            List of applicable vectors.
        """
        return [
            v
            for v in self._vectors.values()
            if not v.protocol_filter or protocol in v.protocol_filter
        ]

    def get_attacks_by_harm_category(self, category: HarmCategory) -> list[AttackVector]:
        """Get attack vectors for a specific harm category.

        Returns:
            List of vectors with the specified harm category.
        """
        return [v for v in self._vectors.values() if v.harm_category == category]

    def search(
        self,
        phase: ScanPhase | None = None,
        category: AttackCategory | None = None,
        severity: Severity | None = None,
        tags: list[str] | None = None,
        protocol: str | None = None,
    ) -> list[AttackVector]:
        """Search vectors with multiple filters (AND logic).

        Args:
            phase: Filter by target phase.
            category: Filter by attack category.
            severity: Filter by severity level.
            tags: Filter by tags (vector must have all specified tags).
            protocol: Filter by protocol (vectors without a protocol_filter match all).

        Returns:
            Vectors matching all specified filters.
        """
        results = list(self._vectors.values())

        if phase is not None:
            results = [v for v in results if v.target_phase == phase]
        if category is not None:
            results = [v for v in results if v.category == category]
        if severity is not None:
            results = [v for v in results if v.severity == severity]
        if tags:
            results = [v for v in results if all(t in v.tags for t in tags)]
        if protocol is not None:
            results = [v for v in results if not v.protocol_filter or protocol in v.protocol_filter]

        return results

    def _load_directory(self, directory: Path) -> None:
        """Load all YAML files from a directory.

        Args:
            directory: Path to directory containing YAML vector files.
        """
        if not directory.is_dir():
            logger.warning("Attack vector directory not found: %s", directory)
            return

        yaml_files = sorted(directory.glob("*.yaml")) + sorted(directory.glob("*.yml"))
        if not yaml_files:
            logger.warning("No YAML files found in: %s", directory)
            return

        for yaml_file in yaml_files:
            try:
                self._load_file(yaml_file)
            except (yaml.YAMLError, ValidationError, KeyError, OSError) as exc:
                logger.warning("Failed to load attack vectors from %s: %s", yaml_file, exc)
                self._load_errors.append((str(yaml_file), exc))
            except Exception as exc:
                logger.exception("Unexpected error loading attack vectors from %s", yaml_file)
                self._load_errors.append((str(yaml_file), exc))

    def _load_file(self, filepath: Path) -> None:
        """Load attack vectors from a single YAML file.

        Args:
            filepath: Path to the YAML file.
        """
        with filepath.open() as f:
            data = yaml.load(f, Loader=getattr(yaml, "CSafeLoader", yaml.SafeLoader))

        if not data or "vectors" not in data:
            logger.warning("No vectors found in %s", filepath)
            return

        for vector_data in data["vectors"]:
            try:
                vector = self._parse_vector(vector_data)
                if vector.id in self._vectors:
                    logger.warning(
                        "Duplicate vector ID '%s' in %s — overwriting previous definition",
                        vector.id,
                        filepath,
                    )
                self._vectors[vector.id] = vector
            except (ValidationError, KeyError, ValueError) as exc:
                vector_id = vector_data.get("id", "unknown")
                logger.warning(
                    "Failed to parse vector '%s' from %s: %s",
                    vector_id,
                    filepath,
                    exc,
                )
                self._load_errors.append((str(vector_id), exc))
            except Exception as exc:
                vector_id = vector_data.get("id", "unknown")
                logger.exception(
                    "Unexpected error parsing vector '%s' from %s",
                    vector_id,
                    filepath,
                )
                self._load_errors.append((str(vector_id), exc))

    @staticmethod
    def _parse_vector(data: dict[str, Any]) -> AttackVector:
        """Parse a vector dictionary into an AttackVector model.

        Args:
            data: Raw vector data from YAML.

        Returns:
            Validated AttackVector model.
        """
        prompts = []
        for prompt_data in data.get("prompts", []):
            prompts.append(
                AttackPrompt(
                    template=prompt_data["template"],
                    variables=prompt_data.get("variables", {}),
                    success_indicators=prompt_data.get("success_indicators", []),
                    failure_indicators=prompt_data.get("failure_indicators", []),
                )
            )

        return AttackVector(
            id=data["id"],
            name=data["name"],
            category=AttackCategory(data["category"]),
            target_phase=ScanPhase(data["target_phase"]),
            description=data.get("description", ""),
            severity=data["severity"],
            prompts=prompts,
            tags=data.get("tags", []),
            references=data.get("references", []),
            owasp_mapping=[OwaspLlmCategory(o) for o in data.get("owasp_mapping", [])],
            protocol_filter=data.get("protocol_filter", []),
            tactic=data.get("tactic", "single"),
            harm_category=HarmCategory(data["harm_category"])
            if data.get("harm_category")
            else None,
        )


_INSTANCE: AttackLibrary | None = None


def get_attack_library(**kwargs: Any) -> AttackLibrary:
    """Return a cached ``AttackLibrary`` singleton for default configuration.

    When called without arguments the same instance is returned on every
    invocation, avoiding repeated YAML parsing.  Any keyword argument
    forces a fresh instance so that custom directories or other overrides
    are honoured.
    """
    global _INSTANCE
    if kwargs:
        return AttackLibrary(**kwargs)
    if _INSTANCE is None:
        _INSTANCE = AttackLibrary()
    return _INSTANCE
