"""
OpenAI SDK monkey-patch for plyra-trace.

Wraps openai.chat.completions.create (sync + async) and
openai.embeddings.create to emit LLM spans with model, tokens, cost,
prompt content, and completion content.
"""

from __future__ import annotations

import json
import logging
from typing import Any

import wrapt
from opentelemetry import trace as trace_api

from plyra_trace.semconv import SpanAttributes, SpanKind

logger = logging.getLogger("plyra_instrument_openai")

# Store originals for uninstrumentation
_originals: dict[str, Any] = {}


def _patch_openai(tracer_provider: Any, capture_content: bool, capture_usage: bool) -> None:
    """Wrap openai chat completion sync + async + embeddings."""
    try:
        import openai.resources.chat.completions as chat_mod
        import openai.resources.embeddings as emb_mod
    except ImportError as exc:
        raise ImportError("openai >= 1.0.0 required for plyra-instrument-openai") from exc

    # ── Save originals ────────────────────────────────────────────────────────
    _originals["chat_create"] = chat_mod.Completions.create
    _originals["chat_acreate"] = chat_mod.AsyncCompletions.create
    _originals["emb_create"] = emb_mod.Embeddings.create

    # ── Build tracer ──────────────────────────────────────────────────────────
    if tracer_provider is not None:
        tracer = tracer_provider.get_tracer("plyra-instrument-openai")
    else:
        tracer = trace_api.get_tracer("plyra-instrument-openai")

    # ── Sync chat.completions.create ─────────────────────────────────────────
    @wrapt.decorator
    def _wrap_chat_create(wrapped: Any, instance: Any, args: tuple, kwargs: dict) -> Any:
        model = kwargs.get("model", args[0] if args else "unknown")
        span_name = f"llm.{model}"

        with tracer.start_as_current_span(span_name) as span:
            _set_llm_span_attrs(span, model, kwargs, capture_content)
            try:
                response = wrapped(*args, **kwargs)
            except Exception as e:
                span.set_status(trace_api.StatusCode.ERROR, str(e))
                raise

            _record_response(span, response, model, capture_content, capture_usage)
            return response

    # ── Async chat.completions.create ─────────────────────────────────────────
    @wrapt.decorator
    async def _wrap_chat_acreate(wrapped: Any, instance: Any, args: tuple, kwargs: dict) -> Any:
        model = kwargs.get("model", args[0] if args else "unknown")
        span_name = f"llm.{model}"

        with tracer.start_as_current_span(span_name) as span:
            _set_llm_span_attrs(span, model, kwargs, capture_content)
            try:
                response = await wrapped(*args, **kwargs)
            except Exception as e:
                span.set_status(trace_api.StatusCode.ERROR, str(e))
                raise

            _record_response(span, response, model, capture_content, capture_usage)
            return response

    # ── Sync embeddings.create ────────────────────────────────────────────────
    @wrapt.decorator
    def _wrap_emb_create(wrapped: Any, instance: Any, args: tuple, kwargs: dict) -> Any:
        model = kwargs.get("model", "text-embedding-ada-002")

        with tracer.start_as_current_span(f"embedding.{model}") as span:
            span.set_attribute(SpanAttributes.SPAN_KIND, SpanKind.EMBEDDING)
            span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, "EMBEDDING")
            span.set_attribute(SpanAttributes.EMBEDDING_MODEL_NAME, model)
            span.set_attribute(SpanAttributes.LLM_PROVIDER, "openai")
            if capture_content:
                input_data = kwargs.get("input", "")
                if isinstance(input_data, list):
                    input_data = str(input_data[:3])  # first 3 items max
                span.set_attribute(SpanAttributes.INPUT_VALUE, str(input_data)[:2000])

            try:
                response = wrapped(*args, **kwargs)
            except Exception as e:
                span.set_status(trace_api.StatusCode.ERROR, str(e))
                raise

            if capture_usage and hasattr(response, "usage") and response.usage:
                span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, response.usage.prompt_tokens or 0)
                span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, response.usage.total_tokens or 0)
            return response

    # ── Apply patches ─────────────────────────────────────────────────────────
    chat_mod.Completions.create = _wrap_chat_create(chat_mod.Completions.create)  # type: ignore[method-assign]
    chat_mod.AsyncCompletions.create = _wrap_chat_acreate(chat_mod.AsyncCompletions.create)  # type: ignore[method-assign]
    emb_mod.Embeddings.create = _wrap_emb_create(emb_mod.Embeddings.create)  # type: ignore[method-assign]

    logger.debug("plyra-instrument-openai: patches applied")


def _unpatch_openai() -> None:
    """Restore original OpenAI methods."""
    try:
        import openai.resources.chat.completions as chat_mod
        import openai.resources.embeddings as emb_mod

        if "chat_create" in _originals:
            chat_mod.Completions.create = _originals.pop("chat_create")  # type: ignore[method-assign]
        if "chat_acreate" in _originals:
            chat_mod.AsyncCompletions.create = _originals.pop("chat_acreate")  # type: ignore[method-assign]
        if "emb_create" in _originals:
            emb_mod.Embeddings.create = _originals.pop("emb_create")  # type: ignore[method-assign]
    except ImportError:
        pass
    logger.debug("plyra-instrument-openai: patches removed")


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


def _set_llm_span_attrs(span: Any, model: str, kwargs: dict, capture_content: bool) -> None:
    """Set standard LLM span attributes from request kwargs."""
    span.set_attribute(SpanAttributes.SPAN_KIND, SpanKind.LLM)
    span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, "LLM")
    span.set_attribute(SpanAttributes.LLM_MODEL_NAME, model)
    span.set_attribute(SpanAttributes.LLM_PROVIDER, "openai")
    span.set_attribute(SpanAttributes.LLM_SYSTEM, "openai")

    if capture_content:
        msgs = kwargs.get("messages", [])
        if msgs:
            span.set_attribute(SpanAttributes.INPUT_VALUE, json.dumps(msgs, default=str)[:8000])
            span.set_attribute(SpanAttributes.INPUT_MIME_TYPE, "application/json")

    # Invocation parameters
    params = {
        k: v
        for k, v in kwargs.items()
        if k
        in ("temperature", "max_tokens", "max_completion_tokens", "top_p", "stream", "n", "stop", "response_format")
        and v is not None
    }
    if params:
        span.set_attribute(SpanAttributes.LLM_INVOCATION_PARAMETERS, json.dumps(params, default=str))


def _record_response(span: Any, response: Any, model: str, capture_content: bool, capture_usage: bool) -> None:
    """Record response data on the span."""
    if capture_content and hasattr(response, "choices") and response.choices:
        choice = response.choices[0]
        content = ""
        if hasattr(choice, "message") and choice.message:
            content = choice.message.content or ""
            # Tool calls
            if not content and choice.message.tool_calls:
                content = json.dumps(
                    [
                        {"function": {"name": tc.function.name, "arguments": tc.function.arguments}}
                        for tc in choice.message.tool_calls
                    ]
                )
        if content:
            span.set_attribute(SpanAttributes.OUTPUT_VALUE, content[:8000])
            span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "text/plain")

    if capture_usage and hasattr(response, "usage") and response.usage:
        usage = response.usage
        prompt_toks = getattr(usage, "prompt_tokens", 0) or 0
        completion_toks = getattr(usage, "completion_tokens", 0) or 0
        total_toks = getattr(usage, "total_tokens", 0) or (prompt_toks + completion_toks)

        span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, prompt_toks)
        span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, completion_toks)
        span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, total_toks)

        # Estimate cost
        try:
            from plyra_trace.collector.cost_table import estimate_cost

            cost = estimate_cost(model, prompt_toks, completion_toks)
            if cost is not None:
                span.set_attribute("llm.cost_usd", cost)
        except Exception:  # noqa: BLE001
            pass
