"""OpenAI adapter — instruments ``client.chat.completions.create``.

Recognizes an ``openai.OpenAI`` / ``AsyncOpenAI`` client by module, and reads
the model, token usage, and output text out of each chat-completions response.
The ``openai`` package is NOT a hard dependency: the adapter duck-types on the
client, so it works with the real SDK or a stand-in.
"""
from __future__ import annotations

import functools
from typing import Any

from .base import AbomRun, ModelCall, _sha256, register

OPENAI_ENDPOINT = "api.openai.com"


def _output_text(response: Any) -> str:
    """Best-effort extraction of the assistant message text for hashing."""
    choices = getattr(response, "choices", None)
    if choices is None and isinstance(response, dict):
        choices = response.get("choices")
    if not choices:
        return ""
    parts = []
    for ch in choices:
        msg = getattr(ch, "message", None)
        if msg is None and isinstance(ch, dict):
            msg = ch.get("message")
        if msg is None:
            continue
        text = getattr(msg, "content", None)
        if text is None and isinstance(msg, dict):
            text = msg.get("content")
        if text:
            parts.append(text)
    return "".join(parts)


def _tokens(response: Any) -> int | None:
    usage = getattr(response, "usage", None)
    if usage is None and isinstance(response, dict):
        usage = response.get("usage")
    if usage is None:
        return None
    total = getattr(usage, "total_tokens", None)
    if total is None and isinstance(usage, dict):
        total = usage.get("total_tokens")
    if total is not None:
        return total
    # fall back to prompt+completion if total isn't present
    p = getattr(usage, "prompt_tokens", None)
    c = getattr(usage, "completion_tokens", None)
    if p is None and isinstance(usage, dict):
        p, c = usage.get("prompt_tokens"), usage.get("completion_tokens")
    if p is None and c is None:
        return None
    return (p or 0) + (c or 0)


def parse(model: str, response: Any) -> ModelCall:
    """Normalize an OpenAI chat-completions response into a :class:`ModelCall`.

    Records the *requested* model (what was gated), not the snapshot id the API
    may echo back (e.g. ``gpt-4o-mini-2024-07-18``) — keeping the provenance
    record consistent with the manifest the gate checked against.
    """
    text = _output_text(response)
    return ModelCall(
        model=model,
        tokens=_tokens(response),
        output_sha256=_sha256(text) if text else None,
        endpoint=OPENAI_ENDPOINT,
    )


class OpenAIAdapter:
    provider = "openai"

    def matches(self, client: Any) -> bool:
        mod = type(client).__module__ or ""
        name = type(client).__name__
        return mod.split(".")[0] == "openai" or name in ("OpenAI", "AsyncOpenAI")

    def wrap(self, client: Any, run: AbomRun) -> Any:
        completions = getattr(getattr(client, "chat", None), "completions", None)
        if completions is None or not hasattr(completions, "create"):
            raise TypeError(
                "OpenAI adapter expects a client with .chat.completions.create; "
                f"got {type(client).__name__}"
            )
        original = completions.create

        @functools.wraps(original)
        def create(*args: Any, **kwargs: Any) -> Any:
            model = kwargs.get("model") or (args[0] if args else "unknown")
            # gate BEFORE the call leaves the boundary (deny-by-default)
            run.enforce(ModelCall(model=model, endpoint=OPENAI_ENDPOINT))
            response = original(*args, **kwargs)
            run.record(parse(model, response))
            return response

        completions.create = create  # type: ignore[attr-defined]
        return client


register(OpenAIAdapter())
