"""Zero-code-change instrumentation for the OpenAI Python SDK."""
from __future__ import annotations

import functools
import logging
import time
from typing import Any

from ..client import get_client, get_tracer
from ..context import reset_current_span, set_current_span
from ..types import (
    META_FREQUENCY_PENALTY,
    META_MAX_TOKENS,
    META_PRESENCE_PENALTY,
    META_PROVIDER,
    META_REASONING_TOKENS,
    META_RESPONSE_FORMAT,
    META_RESPONSE_ID,
    META_SEED,
    META_SERVICE_TIER,
    META_STOP_REASON,
    META_STOP_SEQUENCES,
    META_TEMPERATURE,
    META_TOOL_CHOICE,
    META_TOP_P,
    OPENAI_KEY_CONTENT,
    OPENAI_KEY_MESSAGES,
    OPENAI_KEY_MODEL,
    OPENAI_KEY_ROLE,
    OPENAI_KEY_TOOL_CALLS,
    OPENAI_KEY_TOOLS,
    SpanKind,
    SpanStatus,
)
from ._openai_streaming import AsyncOpenAIStreamProxy, OpenAIStreamProxy

logger = logging.getLogger("staso")

_patched = False


def patch_openai() -> None:
    """Monkey-patch ``openai.resources.chat.Completions.create`` (sync and async)
    to automatically emit LLM spans with comprehensive data capture.

    Safe to call multiple times -- only patches once.
    """
    global _patched
    if _patched:
        return

    try:
        from openai.resources.chat import AsyncCompletions, Completions
    except ImportError:
        raise ImportError(
            "openai package is required for this integration. "
            "Install it with: pip install 'staso[openai]'"
        )

    _wrap_sync(Completions)
    _wrap_async(AsyncCompletions)
    _patched = True
    logger.debug("staso: openai SDK patched")


def _extract_model(args: tuple, kwargs: dict) -> str:
    return str(kwargs.get("model", args[0] if args else ""))


def _capture_input(model: str, kwargs: dict) -> dict:
    """Build the span input dict from request kwargs."""
    inp: dict[str, Any] = {
        OPENAI_KEY_MODEL: model,
    }

    tools = kwargs.get("tools")
    if tools:
        inp[OPENAI_KEY_TOOLS] = tools

    client = get_client()
    if client and client.config.capture_messages:
        messages = kwargs.get("messages")
        if messages:
            inp[OPENAI_KEY_MESSAGES] = messages

    return inp


def _capture_metadata(kwargs: dict) -> dict[str, Any]:
    """Extract request parameters into metadata dict."""
    meta: dict[str, Any] = {META_PROVIDER: "openai"}

    _optional_keys = [
        ("temperature", META_TEMPERATURE),
        ("max_tokens", META_MAX_TOKENS),
        ("max_completion_tokens", META_MAX_TOKENS),
        ("top_p", META_TOP_P),
        ("frequency_penalty", META_FREQUENCY_PENALTY),
        ("presence_penalty", META_PRESENCE_PENALTY),
        ("seed", META_SEED),
        ("stop", META_STOP_SEQUENCES),
        ("tool_choice", META_TOOL_CHOICE),
        ("response_format", META_RESPONSE_FORMAT),
        ("service_tier", META_SERVICE_TIER),
    ]
    for kwarg_key, meta_key in _optional_keys:
        val = kwargs.get(kwarg_key)
        if val is not None:
            meta[meta_key] = val

    return meta


def _capture_response(span: Any, response: Any) -> None:
    """Extract response data onto span fields."""
    choices = getattr(response, "choices", None)
    if choices and len(choices) > 0:
        message = getattr(choices[0], "message", None)
        if message:
            content = getattr(message, "content", None) or ""
            role = getattr(message, "role", "assistant")
            tool_calls = getattr(message, "tool_calls", None)

            output: dict[str, Any] = {
                OPENAI_KEY_CONTENT: content,
                OPENAI_KEY_ROLE: role,
            }
            if tool_calls:
                tc_list = []
                for tc in tool_calls:
                    if hasattr(tc, "model_dump"):
                        tc_list.append(tc.model_dump())
                    elif isinstance(tc, dict):
                        tc_list.append(tc)
                    else:
                        tc_list.append(str(tc))
                output[OPENAI_KEY_TOOL_CALLS] = tc_list
            span.output = output

        finish_reason = getattr(choices[0], "finish_reason", None)
        if finish_reason:
            span.metadata[META_STOP_REASON] = finish_reason

    response_id = getattr(response, "id", None)
    if response_id:
        span.metadata[META_RESPONSE_ID] = response_id

    usage = getattr(response, "usage", None)
    if usage:
        span.input_tokens = getattr(usage, "prompt_tokens", 0) or 0
        span.output_tokens = getattr(usage, "completion_tokens", 0) or 0
        span.total_tokens = getattr(usage, "total_tokens", 0) or 0

        details = getattr(usage, "completion_tokens_details", None)
        if details:
            reasoning = getattr(details, "reasoning_tokens", None)
            if reasoning:
                span.metadata[META_REASONING_TOKENS] = reasoning


def _wrap_sync(cls: type) -> None:
    original = cls.create

    @functools.wraps(original)
    def _traced(self: Any, *args: Any, **kwargs: Any) -> Any:
        tracer = get_tracer()
        model = _extract_model(args, kwargs)

        if kwargs.get("stream"):
            span = tracer.start_span(
                "openai.chat.completions.create", kind=SpanKind.LLM
            )
            span.model = model
            span.input = _capture_input(model, kwargs)
            span.metadata.update(_capture_metadata(kwargs))
            token = set_current_span(span)
            start_time = time.time()
            try:
                stream = original(self, *args, **kwargs)
                return OpenAIStreamProxy(stream, span, start_time, token)
            except Exception as exc:
                span.record_exception(exc)
                span.end()
                reset_current_span(token)
                raise
        else:
            with tracer.start_as_current_span(
                "openai.chat.completions.create", kind=SpanKind.LLM
            ) as span:
                span.model = model
                span.input = _capture_input(model, kwargs)
                span.metadata.update(_capture_metadata(kwargs))
                try:
                    response = original(self, *args, **kwargs)
                    _capture_response(span, response)
                    span.status = SpanStatus.OK
                    return response
                except Exception as exc:
                    span.record_exception(exc)
                    raise

    cls.create = _traced  # type: ignore[assignment]


def _wrap_async(cls: type) -> None:
    original = cls.create

    @functools.wraps(original)
    async def _traced(self: Any, *args: Any, **kwargs: Any) -> Any:
        tracer = get_tracer()
        model = _extract_model(args, kwargs)

        if kwargs.get("stream"):
            span = tracer.start_span(
                "openai.chat.completions.create", kind=SpanKind.LLM
            )
            span.model = model
            span.input = _capture_input(model, kwargs)
            span.metadata.update(_capture_metadata(kwargs))
            token = set_current_span(span)
            start_time = time.time()
            try:
                stream = await original(self, *args, **kwargs)
                return AsyncOpenAIStreamProxy(stream, span, start_time, token)
            except Exception as exc:
                span.record_exception(exc)
                span.end()
                reset_current_span(token)
                raise
        else:
            with tracer.start_as_current_span(
                "openai.chat.completions.create", kind=SpanKind.LLM
            ) as span:
                span.model = model
                span.input = _capture_input(model, kwargs)
                span.metadata.update(_capture_metadata(kwargs))
                try:
                    response = await original(self, *args, **kwargs)
                    _capture_response(span, response)
                    span.status = SpanStatus.OK
                    return response
                except Exception as exc:
                    span.record_exception(exc)
                    raise

    cls.create = _traced  # type: ignore[assignment]
