"""
Capability inference via embedding similarity.

Uses an embedding model to semantically match a task description against
pre-defined capability dimension anchors.  This is the *correct* way to
infer ``required_capabilities`` for adaptive routing — no keyword hacks.

Design:
  1. Define rich anchor texts for each of the 8 capability dimensions
     (bilingual EN + CN for robust multilingual matching).
  2. Embed the task prompt *and* all anchor texts in a single batch call.
  3. Compute cosine similarity between the task embedding and each anchor.
  4. Normalize scores → return an ``AgentCapability`` vector.

The embedding model is configurable (e.g. ``text-embedding-3-small``,
``gemini/text-embedding-004``, any LiteLLM-supported embedding model).
"""

from __future__ import annotations

import logging
import math
from typing import Any

from nexus.types import AgentCapability

logger = logging.getLogger("nexus.orchestration.inference")

# ---------------------------------------------------------------------------
# Anchor texts: one rich description per capability dimension.
# Bilingual (EN + CN) for multilingual embedding models.
# ---------------------------------------------------------------------------

CAPABILITY_ANCHORS: dict[str, str] = {
    "coding": (
        "software programming, writing code, implementing algorithms, "
        "debugging, building applications, API development, software engineering, "
        "refactoring, code review, unit testing, data structures, "
        "软件编程、写代码、实现算法、快速排序、数据结构、调试bug、"
        "开发应用、代码重构、程序设计"
    ),
    "reasoning": (
        "logical reasoning, analytical thinking, problem solving, deduction, "
        "inference, critical analysis, root cause analysis, explaining why, "
        "逻辑推理、分析思考、问题解决、因果分析、"
        "批判性思维、解释原因、推理判断"
    ),
    "writing": (
        "writing articles, creative writing, drafting documents, composing essays, "
        "blog posts, technical writing, copywriting, content creation, storytelling, "
        "写文章、创意写作、撰写文档、博客、技术写作、"
        "文案撰写、内容创作、故事叙述、小说"
    ),
    "research": (
        "research, investigation, searching for information, literature review, "
        "fact-checking, competitive analysis, market research, survey, "
        "调研、研究、搜索信息、文献综述、事实核查、"
        "竞品分析、市场调研、资料查找"
    ),
    "data_analysis": (
        "data analysis, data science, statistics, visualization, charts, graphs, "
        "data processing, machine learning, datasets, spreadsheets, SQL queries, "
        "数据分析、数据科学、统计学、可视化、图表、"
        "数据处理、机器学习、数据集、报表"
    ),
    "math": (
        "mathematics, calculations, equations, formulas, proofs, algebra, "
        "calculus, linear algebra, probability, geometry, numerical methods, "
        "数学、计算、方程、公式、证明、代数、"
        "微积分、线性代数、概率、几何"
    ),
    "creativity": (
        "creative thinking, brainstorming, ideation, design, innovation, "
        "artistic expression, imaginative solutions, concept art, UX design, "
        "创意思维、头脑风暴、设计、创新、"
        "艺术表达、想象力、概念设计"
    ),
    "communication": (
        "communication, summarization, email writing, presentations, "
        "explaining concepts, teaching, translating, negotiation, "
        "沟通、总结、邮件撰写、演示文稿、"
        "概念解释、教学、翻译、谈判"
    ),
}

# Ordered list of dimension names (matches AgentCapability field order)
_DIM_NAMES: list[str] = list(CAPABILITY_ANCHORS.keys())


def _cosine_similarity(a: list[float], b: list[float]) -> float:
    """Cosine similarity between two dense vectors."""
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    if norm_a == 0.0 or norm_b == 0.0:
        return 0.0
    return dot / (norm_a * norm_b)


class CapabilityInferenceService:
    """
    Infers ``AgentCapability`` from a natural-language task description
    using an embedding model.

    Usage::

        service = CapabilityInferenceService(
            embedding_model="text-embedding-3-small",
            api_key="sk-xxx",
        )
        caps = await service.infer("用 Python 实现快排算法")
        # → AgentCapability(coding=0.82, reasoning=0.45, ...)

    The service embeds the task prompt and all 8 capability anchor texts,
    then computes cosine similarity to build the capability vector.
    """

    def __init__(
        self,
        embedding_model: str = "text-embedding-3-small",
        api_key: str | None = None,
        api_base: str | None = None,
        threshold: float = 0.3,
        **kwargs: Any,
    ) -> None:
        """
        Args:
            embedding_model: Model identifier in LiteLLM format
                             (e.g. "text-embedding-3-small", "gemini/text-embedding-004")
            api_key: API key for the embedding provider
            api_base: Base URL override (for custom/proxy endpoints)
            threshold: Minimum similarity score — below this is treated as 0.0
            **kwargs: Extra params forwarded to litellm.aembedding()
        """
        self._model = embedding_model
        self._api_key = api_key
        self._api_base = api_base
        self._threshold = threshold
        self._extra_kwargs = kwargs

        # Cache anchor embeddings (computed lazily on first call)
        self._anchor_embeddings: dict[str, list[float]] | None = None
        self._litellm: Any = None

    def _get_litellm(self) -> Any:
        """Lazy import litellm."""
        if self._litellm is None:
            try:
                import litellm
                self._litellm = litellm
            except ImportError:
                raise ImportError(
                    "litellm is required for embedding-based capability inference. "
                    "Install with: pip install nexus-agents[litellm]"
                )
        return self._litellm

    async def _embed_batch(self, texts: list[str]) -> list[list[float]]:
        """Embed a batch of texts and return their vectors."""
        litellm = self._get_litellm()

        params: dict[str, Any] = {
            "model": self._model,
            "input": texts,
        }
        if self._api_key:
            params["api_key"] = self._api_key
        if self._api_base:
            params["api_base"] = self._api_base
        params.update(self._extra_kwargs)

        response = await litellm.aembedding(**params)

        # Sort by index to preserve order
        data = sorted(response.data, key=lambda x: x["index"])
        return [item["embedding"] for item in data]

    async def _ensure_anchors(self) -> dict[str, list[float]]:
        """Embed all anchor texts (cached after first call)."""
        if self._anchor_embeddings is not None:
            return self._anchor_embeddings

        anchor_texts = [CAPABILITY_ANCHORS[dim] for dim in _DIM_NAMES]
        vectors = await self._embed_batch(anchor_texts)

        self._anchor_embeddings = dict(zip(_DIM_NAMES, vectors))
        logger.debug(
            f"Anchor embeddings cached ({len(vectors)} dims, "
            f"{len(vectors[0])} dimensions each)"
        )
        return self._anchor_embeddings

    async def infer(self, text: str) -> AgentCapability:
        """
        Infer capability requirements from a task description.

        Embeds the text, computes cosine similarity against each
        capability anchor, and returns an ``AgentCapability`` vector.
        """
        # Get anchor embeddings (cached)
        anchors = await self._ensure_anchors()

        # Embed the task text
        task_vectors = await self._embed_batch([text])
        task_vec = task_vectors[0]

        # Compute similarity against each dimension
        scores: dict[str, float] = {}
        for dim, anchor_vec in anchors.items():
            sim = _cosine_similarity(task_vec, anchor_vec)
            # Apply threshold and clamp to [0, 1]
            scores[dim] = max(0.0, min(1.0, sim)) if sim >= self._threshold else 0.0

        # Normalize: scale so that the max score maps to ~1.0
        max_score = max(scores.values()) if scores else 0.0
        if max_score > 0:
            scores = {k: round(v / max_score, 3) for k, v in scores.items()}

        logger.info(
            f"🧠 Inferred capabilities: "
            + ", ".join(f"{k}={v:.2f}" for k, v in scores.items() if v > 0)
        )

        return AgentCapability(
            coding=scores.get("coding", 0.0),
            reasoning=scores.get("reasoning", 0.0),
            writing=scores.get("writing", 0.0),
            research=scores.get("research", 0.0),
            data_analysis=scores.get("data_analysis", 0.0),
            math=scores.get("math", 0.0),
            creativity=scores.get("creativity", 0.0),
            communication=scores.get("communication", 0.0),
        )
