"""
OpenAI Patcher

Monkey-patches openai.resources.chat.completions.Completions.create (sync)
and openai.resources.chat.completions.AsyncCompletions.create (async) to
wrap each call in a GENERATION span with model name and token usage.

Only patches if the `openai` package is importable.
"""

import functools
import logging
from typing import Any, Callable, Optional

from opentelemetry import trace
from opentelemetry.trace import StatusCode

logger = logging.getLogger("pendo_sdk")

_SDK_NAME = "pendo_sdk"

# Stash originals for unpatching.
_original_sync_create: Optional[Callable] = None
_original_async_create: Optional[Callable] = None


def _extract_attributes(span: trace.Span, result: Any) -> None:
    """Extract model name and token usage from an OpenAI response."""
    model = getattr(result, "model", None)
    if model is not None:
        span.set_attribute("llm.model_name", str(model))

    usage = getattr(result, "usage", None)
    if usage is not None:
        # OpenAI uses prompt_tokens / completion_tokens
        prompt_tokens = getattr(usage, "prompt_tokens", None)
        if prompt_tokens is not None:
            span.set_attribute("llm.token_count.prompt", prompt_tokens)
        completion_tokens = getattr(usage, "completion_tokens", None)
        if completion_tokens is not None:
            span.set_attribute("llm.token_count.completion", completion_tokens)
        total_tokens = getattr(usage, "total_tokens", None)
        if total_tokens is not None:
            span.set_attribute("llm.token_count.total", total_tokens)


def _make_sync_wrapper(original: Callable) -> Callable:
    """Wrap the sync Completions.create method."""

    @functools.wraps(original)
    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        model = kwargs.get("model") or (args[0] if args else "unknown")
        span_name = "openai.chat.completions.create"
        tracer = trace.get_tracer(_SDK_NAME)

        with tracer.start_as_current_span(span_name) as span:
            span.set_attribute("pendo.span.type", "GENERATION")
            span.set_attribute("llm.model_name", str(model))
            try:
                result = original(self, *args, **kwargs)
            except Exception as exc:
                span.set_status(StatusCode.ERROR, str(exc))
                span.record_exception(exc)
                raise
            _extract_attributes(span, result)
            return result

    return wrapper


def _make_async_wrapper(original: Callable) -> Callable:
    """Wrap the async AsyncCompletions.create method."""

    @functools.wraps(original)
    async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        model = kwargs.get("model") or (args[0] if args else "unknown")
        span_name = "openai.chat.completions.create"
        tracer = trace.get_tracer(_SDK_NAME)

        with tracer.start_as_current_span(span_name) as span:
            span.set_attribute("pendo.span.type", "GENERATION")
            span.set_attribute("llm.model_name", str(model))
            try:
                result = await original(self, *args, **kwargs)
            except Exception as exc:
                span.set_status(StatusCode.ERROR, str(exc))
                span.record_exception(exc)
                raise
            _extract_attributes(span, result)
            return result

    return wrapper


def patch_openai() -> Optional[Callable[[], None]]:
    """Patch openai.chat.completions.Completions.create and async variant.

    Returns an unpatch callable, or None if openai is not installed.
    """
    global _original_sync_create, _original_async_create

    try:
        import openai.resources.chat.completions as completions_mod
    except ImportError:
        logger.debug("openai not installed, skipping patch")
        return None

    # Patch sync
    sync_cls = getattr(completions_mod, "Completions", None)
    if sync_cls is not None:
        _original_sync_create = sync_cls.create
        sync_cls.create = _make_sync_wrapper(_original_sync_create)

    # Patch async
    async_cls = getattr(completions_mod, "AsyncCompletions", None)
    if async_cls is not None:
        _original_async_create = async_cls.create
        async_cls.create = _make_async_wrapper(_original_async_create)

    logger.debug("OpenAI patched")

    def _unpatch() -> None:
        global _original_sync_create, _original_async_create
        if sync_cls is not None and _original_sync_create is not None:
            sync_cls.create = _original_sync_create
            _original_sync_create = None
        if async_cls is not None and _original_async_create is not None:
            async_cls.create = _original_async_create
            _original_async_create = None
        logger.debug("OpenAI unpatched")

    return _unpatch
