from __future__ import annotations

import json
import math
from collections.abc import Callable
from typing import Any

import tiktoken
from pydantic import BaseModel

from nexus.core.state.embedding_fallback import deterministic_embedding
from nexus.core.state.hot_layer import HotEntry


class CompressionResult(BaseModel):
    """Result for one compressed cluster."""

    original_keys: list[str]
    compressed_content: str
    original_tokens: int
    compressed_tokens: int
    actual_ratio: float
    method: str


class SemanticCompressor:
    """Hierarchical semantic compressor with embedding-based clustering."""

    def __init__(
        self,
        llm_caller: Callable[[str], Any],
        embedding_model: str = "all-MiniLM-L6-v2",
        target_ratio: float = 4.0,
    ) -> None:
        self.llm_caller = llm_caller
        self.target_ratio = target_ratio
        self._encoder = tiktoken.get_encoding("cl100k_base")
        self._embedding_model = embedding_model
        self._embedder: object | None = None
        self._use_fallback = False

    def _get_embedder(self) -> object | None:
        if self._embedder is None and not self._use_fallback:
            try:
                from sentence_transformers import SentenceTransformer  # lazy import

                self._embedder = SentenceTransformer(self._embedding_model)
            except Exception:
                self._use_fallback = True
        return self._embedder

    def _embed_text(self, text: str) -> list[float]:
        embedder = self._get_embedder()
        if embedder is None or self._use_fallback:
            return deterministic_embedding(text)
        return [float(x) for x in embedder.encode(text).tolist()]  # type: ignore[attr-defined]

    async def compress(self, entries: list[HotEntry], target_ratio: float | None = None) -> list[CompressionResult]:
        """Cluster entries and compress each cluster to target ratio."""
        if not entries:
            return []
        ratio = target_ratio or self.target_ratio
        clusters = self._cluster_by_similarity(entries)
        out: list[CompressionResult] = []
        for cluster in clusters:
            result = await self._compress_cluster(cluster)
            if result.actual_ratio < ratio and result.compressed_tokens > 0:
                target_tokens = max(32, int(result.original_tokens / ratio))
                trim = self._encoder.decode(self._encoder.encode(result.compressed_content)[:target_tokens])
                trimmed_tokens = len(self._encoder.encode(trim))
                result = CompressionResult(
                    original_keys=result.original_keys,
                    compressed_content=trim,
                    original_tokens=result.original_tokens,
                    compressed_tokens=trimmed_tokens,
                    actual_ratio=(result.original_tokens / max(1, trimmed_tokens)),
                    method="truncate",
                )
            out.append(result)
        return out

    def _cluster_by_similarity(self, entries: list[HotEntry], threshold: float = 0.75) -> list[list[HotEntry]]:
        """Greedy clustering by nearest centroid cosine similarity."""
        texts = [str(e.value) for e in entries]
        vectors = [self._embed_text(t) for t in texts]

        clusters: list[list[tuple[HotEntry, list[float]]]] = []
        centroids: list[list[float]] = []

        for entry, vec in zip(entries, vectors):
            best_idx = -1
            best_sim = -1.0
            for idx, c in enumerate(centroids):
                sim = self._cosine(vec, c)
                if sim > best_sim:
                    best_sim = sim
                    best_idx = idx
            if best_idx >= 0 and best_sim >= threshold:
                clusters[best_idx].append((entry, vec))
                centroids[best_idx] = self._mean_vector([v for _, v in clusters[best_idx]])
            else:
                clusters.append([(entry, vec)])
                centroids.append(vec)

        return [[entry for entry, _ in cluster] for cluster in clusters]

    async def _compress_cluster(self, cluster: list[HotEntry]) -> CompressionResult:
        """Choose compression method based on value types."""
        keys = [e.key for e in cluster]
        values = [e.value for e in cluster]
        original_text = "\n".join(str(v) for v in values)
        original_tokens = len(self._encoder.encode(original_text))

        if all(isinstance(v, (dict, list)) for v in values):
            compressed = self._compact_json(cluster)
            method = "json_compact"
        else:
            prompt = (
                "Compress the following content while preserving factual accuracy and key actions."
                " Return concise plain text only.\n\n"
                f"{original_text}"
            )
            try:
                compressed = str(await self.llm_caller(prompt))
                method = "llm_summary"
            except Exception:
                compressed = original_text[: max(200, len(original_text) // 3)]
                method = "truncate"

        compressed_tokens = len(self._encoder.encode(compressed))
        return CompressionResult(
            original_keys=keys,
            compressed_content=compressed,
            original_tokens=original_tokens,
            compressed_tokens=compressed_tokens,
            actual_ratio=(original_tokens / max(1, compressed_tokens)),
            method=method,
        )

    def _compact_json(self, entries: list[HotEntry]) -> str:
        """Compact JSON representation for structured entries."""
        payload = {e.key: e.value for e in entries}
        return json.dumps(payload, separators=(",", ":"), ensure_ascii=False)

    def _cosine(self, a: list[float], b: list[float]) -> float:
        dot = sum(x * y for x, y in zip(a, b))
        na = math.sqrt(sum(x * x for x in a))
        nb = math.sqrt(sum(y * y for y in b))
        if na == 0.0 or nb == 0.0:
            return 0.0
        return dot / (na * nb)

    def _mean_vector(self, vectors: list[list[float]]) -> list[float]:
        if not vectors:
            return []
        dim = len(vectors[0])
        sums = [0.0] * dim
        for vec in vectors:
            for i, val in enumerate(vec):
                sums[i] += float(val)
        return [val / len(vectors) for val in sums]
