"""SemanticChunkingStrategy — embedding-based context chunking.

Implements a context strategy that uses embedding similarity to retain
the most semantically relevant context chunks for complex multi-turn
agent workflows where topic coherence matters.

The strategy:
1. Computes embeddings for each fragment's content and the anchor message.
2. Ranks fragments by cosine similarity to the anchor.
3. Selects the top-K most relevant chunks within the token budget.
4. Caches embeddings to avoid redundant API calls.

Implements the ``ContextStrategy`` protocol defined in ``acms_service.py``
and is registered in the ``ACMSPipeline`` under key ``"semantic_chunking"``.

Based on ``docs/specification.md`` §25207-25216.

ISSUES CLOSED: #9996
"""

from __future__ import annotations

import logging
import math
import re
from collections.abc import Callable, Sequence
from typing import Any

from cleveragents.application.services.acms_service import (
    StrategyCapabilities,
    _pack_budget,
)
from cleveragents.domain.models.core.context_fragment import (
    ContextBudget,
    ContextFragment,
)

__all__ = [
    "DEFAULT_EMBEDDING_MODEL",
    "DEFAULT_TOP_K",
    "EmbeddingFn",
    "SemanticChunkingStrategy",
]

logger = logging.getLogger(__name__)

# Default embedding model name (used when no real embedding provider is wired).
DEFAULT_EMBEDDING_MODEL: str = "text-embedding-ada-002"

# Default number of top-K chunks to retain.
DEFAULT_TOP_K: int = 10

# Word tokenisation pattern (reused from context_strategies.py).
_WORD_RE = re.compile(r"\w+", re.UNICODE)

# Type alias for an embedding function: text -> list[float]
EmbeddingFn = Callable[[str], list[float]]


# ---------------------------------------------------------------------------
# Default (built-in) embedding function
# ---------------------------------------------------------------------------


def _default_embedding(text: str) -> list[float]:
    """Compute a simple character-frequency embedding for *text*.

    This is a bag-of-characters vector over the 64 most common ASCII
    characters (ordinals 32-95), normalised to unit length.  It provides
    a lightweight, dependency-free approximation of semantic similarity
    suitable for testing and fallback use.

    Production deployments should inject a real embedding function via
    the ``embedding_fn`` constructor parameter.
    """
    # Defensive guard: public API never calls with empty text.
    if not text:  # pragma: no cover
        return [0.0] * 64
    vec = [0.0] * 64
    for ch in text.lower():
        idx = ord(ch) - 32
        if 0 <= idx < 64:
            vec[idx] += 1.0
    magnitude = math.sqrt(sum(v * v for v in vec))
    if magnitude > 0.0:
        vec = [v / magnitude for v in vec]
    return vec


# ---------------------------------------------------------------------------
# Cosine similarity helper
# ---------------------------------------------------------------------------


def _cosine_similarity(vec_a: list[float], vec_b: list[float]) -> float:
    """Compute cosine similarity between two equal-length vectors.

    Returns a value in [-1.0, 1.0].  Returns 0.0 when either vector is
    the zero vector to avoid division-by-zero.
    """
    # Defensive guard: embedding_fn always returns same-length vectors.
    if len(vec_a) != len(vec_b):  # pragma: no cover
        return 0.0
    dot = sum(a * b for a, b in zip(vec_a, vec_b, strict=True))
    mag_a = math.sqrt(sum(a * a for a in vec_a))
    mag_b = math.sqrt(sum(b * b for b in vec_b))
    if mag_a == 0.0 or mag_b == 0.0:
        return 0.0
    return dot / (mag_a * mag_b)


# ---------------------------------------------------------------------------
# SemanticChunkingStrategy
# ---------------------------------------------------------------------------


class SemanticChunkingStrategy:
    """Embedding-based semantic chunking strategy.

    Ranks context fragments by cosine similarity of their content
    embeddings to an anchor message (typically the current query or
    last message in the conversation).  The top-K most similar chunks
    are retained within the token budget.

    Embeddings are cached by content string to avoid redundant API calls
    when multiple fragments share identical content.

    **Configuration**:

    - ``embedding_model``: Name of the embedding model to use.  Passed
      to the ``embedding_fn`` for informational purposes; the actual
      model selection is the responsibility of the injected function.
    - ``top_k``: Maximum number of fragments to retain before budget
      packing.  Defaults to 10.
    - ``embedding_fn``: Callable ``(text: str) -> list[float]``.
      Defaults to a lightweight character-frequency approximation.
      Inject a real embedding provider for production use.

    **Fallback behaviour**: When no anchor is provided (empty string),
    the strategy falls back to relevance-score ordering.

    Implements ``ContextStrategy`` protocol from ``acms_service.py``.

    Based on ``docs/specification.md`` §25207-25216.
    """

    def __init__(
        self,
        *,
        embedding_model: str = DEFAULT_EMBEDDING_MODEL,
        top_k: int = DEFAULT_TOP_K,
        embedding_fn: EmbeddingFn | None = None,
    ) -> None:
        self._embedding_model = embedding_model
        self._top_k = top_k
        self._embedding_fn: EmbeddingFn = embedding_fn or _default_embedding
        # Cache: content string -> embedding vector
        self._embedding_cache: dict[str, list[float]] = {}
        self._anchor: str = ""

    # ------------------------------------------------------------------
    # Configuration accessors
    # ------------------------------------------------------------------

    @property
    def embedding_model(self) -> str:
        """Return the configured embedding model name."""
        return self._embedding_model

    @property
    def top_k(self) -> int:
        """Return the configured top-K limit."""
        return self._top_k

    def set_anchor(self, anchor: str) -> None:
        """Set the anchor message for similarity ranking (optional)."""
        self._anchor = anchor

    # ------------------------------------------------------------------
    # ContextStrategy protocol
    # ------------------------------------------------------------------

    @property
    def name(self) -> str:
        return "semantic_chunking"

    @property
    def capabilities(self) -> StrategyCapabilities:
        return StrategyCapabilities(supports_semantic_search=True)

    def can_handle(self, request: dict[str, Any]) -> float:
        """Return 0.75 when a query is present, else 0.1.

        The strategy is most useful when there is a query or anchor
        message to compute similarity against.  Without one it falls
        back to relevance ordering, which is less distinctive.
        """
        query = str(request.get("query", "") or "")
        self._anchor = query
        return 0.75 if query else 0.1

    def assemble(
        self,
        fragments: Sequence[ContextFragment],
        budget: ContextBudget,
    ) -> Sequence[ContextFragment]:
        """Rank fragments by cosine similarity to the anchor, then pack.

        Steps:
        1. If no anchor is set, fall back to relevance-score ordering.
        2. Compute (cached) embeddings for the anchor and each fragment.
        3. Rank fragments by cosine similarity (descending).
        4. Apply top-K limit.
        5. Pack within the token budget.
        """
        if not fragments:
            return list(fragments)

        if not self._anchor:
            # No anchor — fall back to relevance ordering
            sorted_frags = sorted(
                fragments, key=lambda f: f.relevance_score, reverse=True
            )
            return _pack_budget(sorted_frags, budget)

        # Compute anchor embedding (cached)
        anchor_vec = self._get_embedding(self._anchor)

        # Score each fragment by cosine similarity to anchor
        scored: list[tuple[ContextFragment, float]] = []
        for frag in fragments:
            frag_vec = self._get_embedding(frag.content)
            sim = _cosine_similarity(anchor_vec, frag_vec)
            scored.append((frag, sim))

        # Sort by similarity descending, then by relevance_score as tiebreaker
        scored.sort(
            key=lambda pair: (pair[1], pair[0].relevance_score),
            reverse=True,
        )

        # Apply top-K limit
        top_frags = [frag for frag, _ in scored[: self._top_k]]

        logger.info(
            "SemanticChunking ranked fragments",
            extra={
                "anchor_length": len(self._anchor),
                "fragment_count": len(fragments),
                "top_k": self._top_k,
                "selected_count": len(top_frags),
                "embedding_model": self._embedding_model,
            },
        )

        return _pack_budget(top_frags, budget)

    def explain(self) -> str:
        return (
            f"Embedding-based semantic chunking strategy. "
            f"Ranks context fragments by cosine similarity of their "
            f"content embeddings to the anchor message. "
            f"Retains the top-{self._top_k} most relevant chunks within "
            f"the token budget. Uses embedding model '{self._embedding_model}'. "
            f"Caches embeddings to avoid redundant API calls. "
            f"Falls back to relevance ordering when no anchor is provided."
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _get_embedding(self, text: str) -> list[float]:
        """Return the embedding for *text*, using the cache when possible."""
        if text not in self._embedding_cache:
            self._embedding_cache[text] = self._embedding_fn(text)
        return self._embedding_cache[text]
