"""Langfuse integration — adds an OTel SpanProcessor to Langfuse's TracerProvider
so all Langfuse trace/generation/span data is also forwarded to Context0.

Works with Langfuse Python SDK v3+ (OTel-native).
Catches ALL usage patterns: @observe, LangfuseOpenAI, propagate_attributes(),
start_as_current_observation(), framework integrations, etc.

Usage:
    from llm_observatory import configure, wrap_langfuse

    configure(api_key="your-api-key")
    wrap_langfuse()
"""

from __future__ import annotations

import json
import logging
import threading
import uuid
from datetime import datetime, timezone
from typing import Any

from llm_observatory.config import get_config
from llm_observatory.transport import get_transport

logger = logging.getLogger("llm_observatory")

# ── State ────────────────────────────────────────────────────────

_installed = False
_lock = threading.Lock()
_emitted_traces: set[str] = set()
_MAX_EMITTED_TRACES = 10_000
_MAX_ID_MAPPINGS = 10_000

# Maps non-UUID raw IDs to their generated UUIDs so the same raw ID always
# resolves to the same UUID (preserves parent-child linkage).
_id_mapping: dict[str, str] = {}


# ── Public API ───────────────────────────────────────────────────

def wrap_langfuse(client: Any = None) -> Any:
    """Hook into Langfuse's OTel pipeline so all trace/generation/span data
    is also sent to Context0.

    Returns the client as-is for chaining. Installation is idempotent.
    Works with every Langfuse usage pattern — no code changes needed.

    Requires Langfuse Python SDK v3.0.0+ (OTel-native).
    """
    try:
        _try_install_processor()
    except Exception as exc:
        logger.debug("wrap_langfuse: failed to install processor: %s", exc)
    return client


def _reset_patch_state() -> None:
    """Reset internal state — for tests."""
    global _installed
    with _lock:
        _installed = False
        _emitted_traces.clear()
        _id_mapping.clear()


# ── Processor installation ───────────────────────────────────────

def _try_install_processor() -> None:
    global _installed
    if _installed:
        return

    processor = Context0SpanProcessor()

    # Strategy 1: Access Langfuse's TracerProvider via get_client()
    try:
        from langfuse._client.get_client import get_client  # type: ignore[import-untyped]
        from opentelemetry.sdk.trace import TracerProvider  # type: ignore[import-untyped]

        langfuse_client = get_client()
        if langfuse_client is not None:
            resources = getattr(langfuse_client, "_resources", None)
            if resources is not None:
                tp = getattr(resources, "tracer_provider", None)
                if tp is not None and isinstance(tp, TracerProvider):
                    tp.add_span_processor(processor)
                    _installed = True
                    logger.debug("wrap_langfuse: installed via Langfuse TracerProvider")
                    return
    except Exception as exc:
        logger.debug("wrap_langfuse: strategy 1 failed: %s", exc)

    # Strategy 2: Navigate the global OTel TracerProvider chain
    try:
        from opentelemetry import trace as otel_trace  # type: ignore[import-untyped]
        from opentelemetry.sdk.trace import TracerProvider  # type: ignore[import-untyped]

        proxy = otel_trace.get_tracer_provider()
        # ProxyTracerProvider wraps the real provider
        tp = getattr(proxy, "_real_tracer_provider", proxy)
        if isinstance(tp, TracerProvider):
            tp.add_span_processor(processor)
            _installed = True
            logger.debug("wrap_langfuse: installed via global OTel TracerProvider")
            return
    except Exception as exc:
        logger.debug("wrap_langfuse: strategy 2 failed: %s", exc)

    logger.warning("wrap_langfuse: could not find a TracerProvider to attach to — Langfuse data will NOT be forwarded to Context0")


# ── OTel SpanProcessor ───────────────────────────────────────────

class Context0SpanProcessor:
    """OTel SpanProcessor that converts Langfuse spans to Context0 events."""

    def on_start(self, span: Any, parent_context: Any = None) -> None:
        pass

    def _on_ending(self, span: Any) -> None:
        pass

    def on_end(self, span: Any) -> None:
        try:
            cfg = get_config()
            if not cfg.get("enabled"):
                return
            event = _span_to_event(span)
            if event is not None:
                get_transport().enqueue(event)
        except Exception as exc:
            logger.debug("Context0SpanProcessor.on_end: %s", exc)

    def shutdown(self) -> None:
        pass

    def force_flush(self, timeout_millis: int | None = None) -> bool:
        return True


# ── Span → Event conversion ─────────────────────────────────────

def _span_to_event(span: Any) -> dict[str, Any] | None:
    """Convert an OTel ReadableSpan (with Langfuse attributes) to a Context0 event."""
    attrs = dict(span.attributes or {})

    # Determine observation type
    obs_type = attrs.get("langfuse.observation.type")  # "generation", "span", etc.

    # Read IDs
    ctx = getattr(span, "context", None)
    if ctx is None:
        return None
    trace_id_hex = format(ctx.trace_id, "032x")
    span_id_hex = format(ctx.span_id, "016x")
    parent_id_hex = format(span.parent.span_id, "016x") if span.parent else None

    safe_trace_id = _ensure_uuid(trace_id_hex)
    safe_event_id = _ensure_uuid(span_id_hex)
    safe_parent_id = _ensure_uuid(parent_id_hex) if parent_id_hex else safe_trace_id

    # Timing (OTel uses nanoseconds)
    start_ns = span.start_time or 0
    end_ns = span.end_time or start_ns
    start_time = datetime.fromtimestamp(start_ns / 1e9, tz=timezone.utc).isoformat()
    end_time = datetime.fromtimestamp(end_ns / 1e9, tz=timezone.utc).isoformat()
    latency_ms = int((end_ns - start_ns) / 1e6)

    # Metadata — with compat key fallbacks (matches TS SDK)
    user_id = attrs.get("user.id") or attrs.get("langfuse.user.id")
    session_id = attrs.get("session.id") or attrs.get("langfuse.session.id")
    raw_tags = attrs.get("langfuse.trace.tags")
    if raw_tags is not None:
        tags = [str(t) for t in raw_tags] if isinstance(raw_tags, (list, tuple)) else [str(raw_tags)]
    else:
        tags = []

    metadata: dict[str, Any] = {}
    if user_id:
        metadata["user_id"] = str(user_id)
    if session_id:
        metadata["session_id"] = str(session_id)
    if tags:
        metadata["tags"] = tags

    # Custom metadata extraction — single JSON blob OR flattened keys
    custom: dict[str, str] = {}

    # Strategy 1: single JSON blob (e.g. langfuse.observation.metadata = '{"key":"val"}')
    custom_meta_raw = attrs.get("langfuse.observation.metadata")
    if custom_meta_raw is not None:
        custom_meta = custom_meta_raw
        if isinstance(custom_meta_raw, str):
            try:
                custom_meta = json.loads(custom_meta_raw)
            except (json.JSONDecodeError, ValueError):
                custom_meta = None
        if isinstance(custom_meta, dict) and custom_meta:
            custom.update({k: str(v) if not isinstance(v, str) else v for k, v in custom_meta.items() if v is not None})

    # Strategy 2: flattened keys (Langfuse Python v3 sets langfuse.observation.metadata.{key} and
    # langfuse.trace.metadata.{key} as individual OTel attributes)
    for prefix in ("langfuse.observation.metadata.", "langfuse.trace.metadata."):
        for attr_key, attr_val in attrs.items():
            if attr_key.startswith(prefix) and attr_val is not None:
                short_key = attr_key[len(prefix):]
                if short_key:
                    custom[short_key] = str(attr_val) if not isinstance(attr_val, str) else attr_val

    if custom:
        metadata["custom"] = custom

    # Emit a synthetic trace event (first time per trace_id)
    _maybe_emit_trace(safe_trace_id, span.name, start_time, metadata)

    # Root spans without an observation type only emit a trace event.
    # Root spans WITH obs_type (e.g. a standalone LangfuseOpenAI call) are
    # emitted as generations — these are single-span traces where the root
    # IS the generation.
    if obs_type is None and parent_id_hex is None:
        return None

    # Span type mapping
    span_type = "generation" if obs_type in ("generation", "embedding") else "span"

    # Level / status
    level = attrs.get("langfuse.observation.level")
    span_status = "error" if level == "ERROR" else "ok"

    # Model (with gen_ai fallback) and usage
    model = attrs.get("langfuse.observation.model.name") or attrs.get("gen_ai.request.model") or ""
    usage = _extract_usage_from_attrs(attrs)

    # Input/Output from attributes
    inp_raw = attrs.get("langfuse.observation.input") or attrs.get("langfuse.trace.input")
    out_raw = attrs.get("langfuse.observation.output") or attrs.get("langfuse.trace.output")

    # Parse input — for generations, this is typically the messages array
    messages: list[Any] = []
    inp_str: str | None = None
    if inp_raw is not None:
        inp_str = inp_raw if isinstance(inp_raw, str) else json.dumps(inp_raw)
        try:
            parsed_input = json.loads(inp_str) if isinstance(inp_raw, str) else inp_raw
            if isinstance(parsed_input, list):
                messages = parsed_input
        except (json.JSONDecodeError, ValueError):
            pass

    # Parse output — for generations, this is typically {role, content}
    response: dict[str, Any] = {}
    out_str: str | None = None
    if out_raw is not None:
        out_str = out_raw if isinstance(out_raw, str) else json.dumps(out_raw)
        try:
            parsed_output = json.loads(out_str) if isinstance(out_raw, str) else out_raw
            if isinstance(parsed_output, dict) and "content" in parsed_output:
                response = {
                    "id": "",
                    "model": model,
                    "choices": [{"index": 0, "message": parsed_output, "finish_reason": "stop"}],
                }
            else:
                # Non-dict JSON (array, scalar) — wrap as plain text
                response = {
                    "id": "",
                    "model": model,
                    "choices": [{"index": 0, "message": {"role": "assistant", "content": out_str}, "finish_reason": "stop"}],
                }
        except (json.JSONDecodeError, ValueError):
            # Not valid JSON — plain text response
            response = {
                "id": "",
                "model": model,
                "choices": [{"index": 0, "message": {"role": "assistant", "content": out_str}, "finish_reason": "stop"}],
            }

    # Build event
    now = datetime.now(timezone.utc).isoformat()
    event: dict[str, Any] = {
        "event_id": safe_event_id,
        "timestamp": now,
        "trace_id": safe_trace_id,
        "parent_id": safe_parent_id,
        "span_type": span_type,
        "span_name": span.name or f"langfuse-{span_type}",
        "span_start_time": start_time,
        "span_end_time": end_time,
        "span_status": span_status,
        "request": {"model": model, "messages": messages, "stream": False},
        "response": response,
        "usage": usage,
        "metrics": {"latency_ms": latency_ms},
        "metadata": metadata,
    }

    if inp_str is not None:
        event["span_input"] = inp_str
    if out_str is not None:
        event["span_output"] = out_str

    return event


# ── Trace emission ───────────────────────────────────────────────

def _maybe_emit_trace(trace_id: str, name: str, start_time: str, metadata: dict[str, Any]) -> None:
    """Emit a synthetic trace-level event the first time we see a trace_id."""
    with _lock:
        if trace_id in _emitted_traces:
            return

        cfg = get_config()
        if not cfg.get("enabled"):
            return

        if len(_emitted_traces) >= _MAX_EMITTED_TRACES:
            _emitted_traces.clear()
        _emitted_traces.add(trace_id)

    now = datetime.now(timezone.utc).isoformat()
    event: dict[str, Any] = {
        "event_id": trace_id,
        "timestamp": now,
        "trace_id": trace_id,
        "span_type": "trace",
        "span_name": name or "langfuse-trace",
        "span_start_time": start_time,
        "span_end_time": start_time,
        "span_status": "ok",
        "request": {"model": "", "messages": [], "stream": False},
        "response": {},
        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
        "metrics": {"latency_ms": 0},
        "metadata": metadata,
    }

    get_transport().enqueue(event)


# ── Helpers ──────────────────────────────────────────────────────

def _first_of(d: dict, *keys: str, default: Any = None) -> Any:
    """Return the first value from *d* whose key exists and is not None."""
    for k in keys:
        v = d.get(k)
        if v is not None:
            return v
    return default


def _extract_usage_from_attrs(attrs: dict[str, Any]) -> dict[str, int]:
    """Extract usage from OTel span attributes."""
    # Try Langfuse usage_details first
    raw = attrs.get("langfuse.observation.usage_details")
    if raw:
        if isinstance(raw, str):
            try:
                raw = json.loads(raw)
            except (json.JSONDecodeError, ValueError):
                raw = None

        if isinstance(raw, dict):
            prompt = _first_of(raw, "prompt_tokens", "promptTokens", "input", "input_tokens", default=0)
            completion = _first_of(raw, "completion_tokens", "completionTokens", "output", "output_tokens", default=0)
            total = _first_of(raw, "total_tokens", "totalTokens", "total")
            if total is None:
                total = (prompt if isinstance(prompt, (int, float)) else 0) + (completion if isinstance(completion, (int, float)) else 0)
            return {
                "prompt_tokens": int(prompt) if isinstance(prompt, (int, float)) else 0,
                "completion_tokens": int(completion) if isinstance(completion, (int, float)) else 0,
                "total_tokens": int(total) if isinstance(total, (int, float)) else 0,
            }

    # Fallback to GenAI semantic conventions
    p = _first_of(attrs, "gen_ai.usage.prompt_tokens", "gen_ai.usage.input_tokens", default=0)
    c = _first_of(attrs, "gen_ai.usage.completion_tokens", "gen_ai.usage.output_tokens", default=0)
    return {
        "prompt_tokens": int(p) if isinstance(p, (int, float)) else 0,
        "completion_tokens": int(c) if isinstance(c, (int, float)) else 0,
        "total_tokens": (int(p) if isinstance(p, (int, float)) else 0) + (int(c) if isinstance(c, (int, float)) else 0),
    }


def _ensure_uuid(val: str | None) -> str:
    """Ensure a value is a valid UUID string.

    If *val* is already a valid UUID, return it as-is.  If it is a non-UUID
    string, generate a stable UUID via ``_id_mapping`` so the same raw ID
    always resolves to the same UUID (critical for parent-child linkage).
    """
    if val is None:
        return str(uuid.uuid4())
    try:
        uuid.UUID(val)
        return val
    except (ValueError, AttributeError):
        with _lock:
            existing = _id_mapping.get(val)
            if existing is not None:
                return existing
            generated = str(uuid.uuid4())
            if len(_id_mapping) >= _MAX_ID_MAPPINGS:
                _id_mapping.clear()
            _id_mapping[val] = generated
            return generated


# ── Test helpers ─────────────────────────────────────────────────

def _process_span(span: Any) -> dict[str, Any] | None:
    """Process a span directly — for tests."""
    return _span_to_event(span)
