"""Finding deduplication — LLM-embedding-based clustering of attack results.

Groups related ``AttackResult`` objects into ``DeduplicatedFinding``
clusters using cosine similarity on LLM-generated embeddings.
Provides cross-scan merging for incremental pentesting sessions.

Requires the ``pentest`` optional extra: ``pip install ziran[pentest]``.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from ziran.domain.entities.attack import AttackResult
    from ziran.infrastructure.llm.base import BaseLLMClient

from ziran.domain.entities.pentest import DeduplicatedFinding

logger = logging.getLogger(__name__)

# Severity ordering for dedup (highest wins when merging)
_SEVERITY_ORDER = {"critical": 4, "high": 3, "medium": 2, "low": 1}


class FindingDeduplicator:
    """Deduplicates attack results using LLM embeddings.

    Uses cosine similarity to cluster related findings so that
    multiple attack vectors targeting the same underlying vulnerability
    are reported as a single finding.

    Example::

        dedup = FindingDeduplicator(llm_client=llm_client)
        findings = await dedup.deduplicate(attack_results, threshold=0.85)
    """

    def __init__(
        self,
        llm_client: BaseLLMClient,
        *,
        embedding_model: str = "text-embedding-3-small",
    ) -> None:
        self.llm_client = llm_client
        self.embedding_model = embedding_model

    async def _embed_text(self, text: str) -> list[float]:
        """Generate an embedding for a text string.

        Uses litellm.aembedding() for multi-provider support.
        Falls back to a simple hash-based pseudo-embedding if
        the embedding call fails.
        """
        try:
            import litellm

            response = await litellm.aembedding(
                model=self.embedding_model,
                input=[text],
            )
            return list(response.data[0]["embedding"])
        except Exception as e:
            logger.warning("Embedding failed, using fallback: %s", e)
            return self._fallback_embedding(text)

    @staticmethod
    def _fallback_embedding(text: str, dim: int = 64) -> list[float]:
        """Simple hash-based pseudo-embedding for when LLM embeddings fail."""
        import hashlib

        h = hashlib.sha256(text.encode()).hexdigest()
        values = [int(h[i : i + 2], 16) / 255.0 for i in range(0, min(dim * 2, len(h)), 2)]
        # Pad to desired dimension
        while len(values) < dim:
            values.append(0.0)
        return values[:dim]

    def _finding_text(self, result: AttackResult) -> str:
        """Create a text representation of a finding for embedding."""
        parts = [
            f"Category: {result.category.value}",
            f"Vector: {result.vector_name}",
            f"Severity: {result.severity}",
        ]
        if result.owasp_mapping:
            parts.append(f"OWASP: {', '.join(o.value for o in result.owasp_mapping)}")
        if result.evidence:
            parts.append(f"Evidence keys: {', '.join(result.evidence.keys())}")
        if result.agent_response:
            parts.append(f"Response preview: {result.agent_response[:200]}")
        return " | ".join(parts)

    @staticmethod
    def _cosine_similarity(a: list[float], b: list[float]) -> float:
        """Compute cosine similarity between two vectors."""
        try:
            import numpy as np

            va = np.array(a)
            vb = np.array(b)
            dot = float(np.dot(va, vb))
            norm_a = float(np.linalg.norm(va))
            norm_b = float(np.linalg.norm(vb))
            if norm_a == 0 or norm_b == 0:
                return 0.0
            return float(dot / (norm_a * norm_b))
        except ImportError:
            # Pure Python fallback
            dot = sum(x * y for x, y in zip(a, b, strict=False))
            norm_a = sum(x * x for x in a) ** 0.5
            norm_b = sum(x * x for x in b) ** 0.5
            if norm_a == 0 or norm_b == 0:
                return 0.0
            return float(dot / (norm_a * norm_b))

    async def deduplicate(
        self,
        results: list[AttackResult],
        *,
        threshold: float = 0.85,
    ) -> list[DeduplicatedFinding]:
        """Cluster related attack results into deduplicated findings.

        Args:
            results: List of successful attack results to deduplicate.
            threshold: Cosine similarity threshold for clustering (0-1).

        Returns:
            List of deduplicated findings.
        """
        if not results:
            return []

        # Only deduplicate successful attacks
        successful = [r for r in results if r.successful]
        if not successful:
            return []

        # Generate embeddings
        texts = [self._finding_text(r) for r in successful]
        embeddings = []
        for text in texts:
            emb = await self._embed_text(text)
            embeddings.append(emb)

        # Cluster using greedy single-link clustering
        clusters: list[list[int]] = []
        assigned = set()

        for i in range(len(successful)):
            if i in assigned:
                continue
            cluster = [i]
            assigned.add(i)
            for j in range(i + 1, len(successful)):
                if j in assigned:
                    continue
                sim = self._cosine_similarity(embeddings[i], embeddings[j])
                if sim >= threshold:
                    cluster.append(j)
                    assigned.add(j)
            clusters.append(cluster)

        # Build deduplicated findings from clusters
        findings = []
        for cluster_indices in clusters:
            cluster_results = [successful[i] for i in cluster_indices]
            finding = self._build_finding(cluster_results, embeddings[cluster_indices[0]])
            findings.append(finding)

        return findings

    def _build_finding(
        self,
        results: list[AttackResult],
        embedding: list[float],
    ) -> DeduplicatedFinding:
        """Build a DeduplicatedFinding from a cluster of results."""
        # Take highest severity
        severities = [r.severity for r in results]
        max_severity = max(severities, key=lambda s: _SEVERITY_ORDER.get(s, 0))

        # Merge OWASP categories
        owasp_cats: list[str] = []
        seen_owasp: set[str] = set()
        for r in results:
            for o in r.owasp_mapping:
                if o.value not in seen_owasp:
                    owasp_cats.append(o.value)
                    seen_owasp.add(o.value)

        # Use the highest-severity result as canonical
        canonical = max(results, key=lambda r: _SEVERITY_ORDER.get(r.severity, 0))

        # Merge evidence
        evidence_parts = []
        for r in results:
            if r.evidence:
                evidence_parts.append(f"{r.vector_name}: {list(r.evidence.keys())}")

        return DeduplicatedFinding(
            canonical_title=canonical.vector_name,
            canonical_description=(
                f"Vulnerability cluster: {len(results)} related finding(s) "
                f"in category '{canonical.category.value}'. "
                f"Primary vector: {canonical.vector_name}."
            ),
            severity=max_severity,
            owasp_categories=owasp_cats,
            attack_result_ids=[r.vector_id for r in results],
            evidence_summary="; ".join(evidence_parts) if evidence_parts else "",
            embedding=embedding,
            confidence=1.0 if len(results) == 1 else 0.9,
        )

    @staticmethod
    def merge_findings(
        existing: list[DeduplicatedFinding],
        new: list[DeduplicatedFinding],
        *,
        threshold: float = 0.85,
    ) -> list[DeduplicatedFinding]:
        """Merge new findings into existing ones for cross-scan dedup.

        Args:
            existing: Previously deduplicated findings.
            new: Newly discovered findings.
            threshold: Similarity threshold for merging.

        Returns:
            Merged list of findings.
        """
        if not existing:
            return list(new)
        if not new:
            return list(existing)

        merged = list(existing)
        for new_finding in new:
            merged_with_existing = False
            for i, ex_finding in enumerate(merged):
                if not new_finding.embedding or not ex_finding.embedding:
                    continue
                sim = FindingDeduplicator._cosine_similarity(
                    new_finding.embedding, ex_finding.embedding
                )
                if sim >= threshold:
                    # Merge into existing finding
                    merged[i] = DeduplicatedFinding(
                        finding_id=ex_finding.finding_id,
                        canonical_title=ex_finding.canonical_title,
                        canonical_description=(
                            f"{ex_finding.canonical_description} "
                            f"(merged with {new_finding.canonical_title})"
                        ),
                        severity=max(
                            ex_finding.severity,
                            new_finding.severity,
                            key=lambda s: _SEVERITY_ORDER.get(s, 0),
                        ),
                        owasp_categories=list(
                            set(ex_finding.owasp_categories + new_finding.owasp_categories)
                        ),
                        attack_result_ids=(
                            ex_finding.attack_result_ids + new_finding.attack_result_ids
                        ),
                        evidence_summary=(
                            f"{ex_finding.evidence_summary}; {new_finding.evidence_summary}"
                        ),
                        embedding=ex_finding.embedding,
                        confidence=min(ex_finding.confidence, new_finding.confidence),
                    )
                    merged_with_existing = True
                    break

            if not merged_with_existing:
                merged.append(new_finding)

        return merged
