"""
Universal span normalization for plyra-trace collector.

Accepts spans from any OTel source — OpenInference, OpenLLMetry, LangChain,
raw OTLP, gen_ai semantic conventions — and normalizes them into plyra's
canonical schema.

Priority order: plyra attrs > OpenInference > OpenLLMetry > LangChain > gen_ai > inferred
"""

from __future__ import annotations

import re
from typing import Any

# ── Kind inference from span name patterns ─────────────────────────────────────

_LLM_PATTERNS = re.compile(r"llm|chat|completion|generate|gpt|claude|gemini|mistral|cohere", re.I)
_TOOL_PATTERNS = re.compile(r"tool|function|action|execute|call|invoke|run_tool", re.I)
_AGENT_PATTERNS = re.compile(r"agent|plan|think|reason|decide|react|step", re.I)
_RETRIEVER_PATTERNS = re.compile(r"embed|retriev|search|vector|similarity|lookup|recall", re.I)
_GUARD_PATTERNS = re.compile(r"guard|safety|policy|check|filter|moderate|censor", re.I)
_CHAIN_DEFAULT = "CHAIN"

# ── OpenLLMetry / Traceloop kind mapping ─────────────────────────────────────
_TRACELOOP_KIND_MAP: dict[str, str] = {
    "workflow": "CHAIN",
    "task": "TOOL",
    "agent": "AGENT",
    "tool": "TOOL",
    "llm": "LLM",
    "retriever": "RETRIEVER",
}

# ── LangChain span type mapping ───────────────────────────────────────────────
_LANGCHAIN_KIND_MAP: dict[str, str] = {
    "chain": "CHAIN",
    "agent": "AGENT",
    "tool": "TOOL",
    "llm": "LLM",
    "retriever": "RETRIEVER",
}

# ── OpenInference kind passthrough (already in plyra format) ─────────────────
_OPENINFERENCE_VALID_KINDS = {
    "LLM",
    "CHAIN",
    "AGENT",
    "TOOL",
    "RETRIEVER",
    "EMBEDDING",
    "RERANKER",
    "GUARD",
    "EVALUATOR",
    "PROMPT",
    "ROUTER",
}


class SpanNormalizer:
    """
    Normalizes spans from any OTel source into plyra's canonical schema.
    Called on every span during OTLP ingest.

    Priority order:
        1. plyra.span.kind (already normalized)
        2. openinference.span.kind
        3. traceloop.span.kind (OpenLLMetry)
        4. langchain.span_type
        5. gen_ai.* + llm.* attributes → infer LLM
        6. Infer from span name patterns
    """

    def normalize(self, raw_span: dict[str, Any]) -> dict[str, Any]:
        """
        Takes a raw span dict (from protobuf or JSON parse).
        Returns a normalized span dict with plyra.span.kind always set.
        """
        span = dict(raw_span)
        attrs: dict[str, Any] = dict(span.get("attributes") or {})

        # ── Determine canonical kind ──────────────────────────────────────────
        kind = self._resolve_kind(span.get("name", ""), attrs)
        attrs["plyra.span.kind"] = kind
        span["kind"] = kind

        # ── Normalize LLM attributes ──────────────────────────────────────────
        attrs = self._normalize_llm_attrs(attrs)

        # ── Normalize input/output ────────────────────────────────────────────
        attrs = self._normalize_io_attrs(attrs)

        # ── Normalize project_id from resource ───────────────────────────────
        if not span.get("project_id") or span.get("project_id") == "unknown":
            # Propagate plyra.project.id from OTLP resource attributes
            resource = span.get("resource") or {}
            if isinstance(resource, dict):
                pid = resource.get("plyra.project.id")
                if pid:
                    attrs["plyra.project.id"] = str(pid)
                    span["project_id"] = str(pid)
            if not span.get("project_id") or span.get("project_id") == "unknown":
                project_id = self._extract_project_id(attrs)
                if project_id:
                    span["project_id"] = project_id

        # ── Normalize environment ─────────────────────────────────────────────
        if not attrs.get("deployment.environment"):
            env = attrs.get("plyra.environment") or attrs.get("PLYRA_ENVIRONMENT")
            if env:
                attrs["deployment.environment"] = env

        span["attributes"] = attrs
        return span

    def _resolve_kind(self, name: str, attrs: dict[str, Any]) -> str:
        """Determine span kind from attributes, using priority order."""
        # 1. plyra native
        plyra_kind = attrs.get("plyra.span.kind")
        if plyra_kind and str(plyra_kind).upper() in _OPENINFERENCE_VALID_KINDS:
            return str(plyra_kind).upper()

        # 2. OpenInference
        oi_kind = attrs.get("openinference.span.kind")
        if oi_kind:
            normalized = str(oi_kind).upper()
            if normalized in _OPENINFERENCE_VALID_KINDS:
                return normalized

        # 3. OpenLLMetry / Traceloop
        tl_kind = attrs.get("traceloop.span.kind")
        if tl_kind:
            mapped = _TRACELOOP_KIND_MAP.get(str(tl_kind).lower())
            if mapped:
                return mapped

        # 4. LangChain
        lc_kind = attrs.get("langchain.span_type")
        if lc_kind:
            mapped = _LANGCHAIN_KIND_MAP.get(str(lc_kind).lower())
            if mapped:
                return mapped

        # 5. gen_ai / llm attributes → must be LLM
        if (
            attrs.get("gen_ai.request.model")
            or attrs.get("gen_ai.system")
            or attrs.get("llm.request.model")
            or attrs.get("llm.model_name")
        ):
            return "LLM"

        # 6. Infer from span name
        return self._infer_kind_from_name(name)

    def _infer_kind_from_name(self, name: str) -> str:
        """
        Infer span kind from span name patterns when no explicit kind is set.
        Used for raw OTLP spans from non-instrumented agents.
        """
        if not name:
            return _CHAIN_DEFAULT
        if _GUARD_PATTERNS.search(name):
            return "GUARD"
        if _LLM_PATTERNS.search(name):
            return "LLM"
        if _TOOL_PATTERNS.search(name):
            return "TOOL"
        if _AGENT_PATTERNS.search(name):
            return "AGENT"
        if _RETRIEVER_PATTERNS.search(name):
            return "RETRIEVER"
        return _CHAIN_DEFAULT

    def _normalize_llm_attrs(self, attrs: dict[str, Any]) -> dict[str, Any]:
        """
        Map source-specific LLM attributes to plyra canonical names.

        Handles: OpenLLMetry, gen_ai semantic conventions, OpenInference passthrough.
        """
        # Model name
        if not attrs.get("llm.model_name"):
            for src in ("gen_ai.request.model", "llm.request.model", "gen_ai.response.model", "model"):
                if attrs.get(src):
                    attrs["llm.model_name"] = attrs[src]
                    break

        # Provider
        if not attrs.get("llm.provider"):
            for src in ("gen_ai.system", "llm.system", "llm.provider"):
                if attrs.get(src):
                    attrs["llm.provider"] = attrs[src]
                    break

        # Token counts
        if not attrs.get("llm.token_count.prompt"):
            for src in ("llm.usage.prompt_tokens", "gen_ai.usage.input_tokens", "gen_ai.usage.prompt_tokens"):
                if attrs.get(src) is not None:
                    attrs["llm.token_count.prompt"] = int(attrs[src])
                    break

        if not attrs.get("llm.token_count.completion"):
            for src in ("llm.usage.completion_tokens", "gen_ai.usage.output_tokens", "gen_ai.usage.completion_tokens"):
                if attrs.get(src) is not None:
                    attrs["llm.token_count.completion"] = int(attrs[src])
                    break

        if not attrs.get("llm.token_count.total"):
            for src in ("llm.usage.total_tokens", "gen_ai.usage.total_tokens"):
                if attrs.get(src) is not None:
                    attrs["llm.token_count.total"] = int(attrs[src])
                    break
            # Compute from parts if not set
            if not attrs.get("llm.token_count.total"):
                p = attrs.get("llm.token_count.prompt", 0) or 0
                c = attrs.get("llm.token_count.completion", 0) or 0
                if p or c:
                    attrs["llm.token_count.total"] = int(p) + int(c)

        # Cost
        if not attrs.get("llm.cost_usd"):
            for src in ("gen_ai.usage.cost", "litellm.cost", "llm.cost"):
                if attrs.get(src) is not None:
                    attrs["llm.cost_usd"] = float(attrs[src])
                    break
            # Try estimation if not found
            if not attrs.get("llm.cost_usd"):
                cost = self._extract_llm_cost(attrs)
                if cost is not None:
                    attrs["llm.cost_usd"] = cost

        return attrs

    def _normalize_io_attrs(self, attrs: dict[str, Any]) -> dict[str, Any]:
        """
        Map source-specific input/output attributes to plyra canonical names.
        """
        # Input
        if not attrs.get("input.value"):
            for src in ("gen_ai.prompt.0.content", "langchain.inputs", "traceloop.entity.input", "llm.prompts"):
                if attrs.get(src):
                    attrs["input.value"] = attrs[src]
                    break

        # Output
        if not attrs.get("output.value"):
            for src in (
                "gen_ai.completion.0.content",
                "langchain.outputs",
                "traceloop.entity.output",
                "llm.completions",
            ):
                if attrs.get(src):
                    attrs["output.value"] = attrs[src]
                    break

        # Agent/tool names from OpenLLMetry
        if not attrs.get("agent.name") and attrs.get("traceloop.entity.name"):
            if attrs.get("plyra.span.kind") in ("AGENT", "CHAIN"):
                attrs["agent.name"] = attrs["traceloop.entity.name"]

        if not attrs.get("tool.name") and attrs.get("traceloop.entity.name"):
            if attrs.get("plyra.span.kind") == "TOOL":
                attrs["tool.name"] = attrs["traceloop.entity.name"]

        return attrs

    def _extract_llm_cost(self, attrs: dict[str, Any]) -> float | None:
        """
        Try to estimate LLM cost from model name + token counts.
        Returns cost in USD or None if unable to estimate.
        """
        from plyra_trace.collector.cost_table import estimate_cost

        model = attrs.get("llm.model_name") or attrs.get("gen_ai.request.model")
        prompt_tokens = attrs.get("llm.token_count.prompt")
        completion_tokens = attrs.get("llm.token_count.completion")

        if model and prompt_tokens is not None and completion_tokens is not None:
            return estimate_cost(str(model), int(prompt_tokens), int(completion_tokens))
        return None

    def _extract_project_id(self, attrs: dict[str, Any]) -> str | None:
        """Extract project_id from any known attribute."""
        for src in ("plyra.project", "service.name", "PLYRA_PROJECT"):
            if attrs.get(src):
                return str(attrs[src])
        return None
