from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any, Dict, Optional


class LLMAdapter(ABC):
    """
    Provider adapter interface.

    The wrapper/orchestrator should be provider-agnostic and only depend on this
    adapter surface to parse input/output and usage.
    """

    @abstractmethod
    def extract_input(self, args: tuple[Any, ...], kwargs: Dict[str, Any]) -> Optional[str]:
        raise NotImplementedError

    @abstractmethod
    def get_model_name(
        self, args: tuple[Any, ...], kwargs: Dict[str, Any], response: Any | None = None
    ) -> Optional[str]:
        raise NotImplementedError

    @abstractmethod
    def extract_output(self, response: Any) -> str:
        raise NotImplementedError

    @abstractmethod
    def get_token_usage(self, response: Any) -> Optional[Dict[str, int]]:
        raise NotImplementedError


def _extract_last_user_message_text(messages: Any) -> Optional[str]:
    if not isinstance(messages, list):
        return None

    last_user = None
    for msg in messages:
        if isinstance(msg, dict) and msg.get("role") == "user":
            last_user = msg

    if not isinstance(last_user, dict):
        return None

    content = last_user.get("content")
    if isinstance(content, str):
        return content

    # OpenAI SDK can represent content as a list of parts (e.g., [{"type":"text","text":"..."}]).
    if isinstance(content, list):
        parts: list[str] = []
        for part in content:
            if isinstance(part, dict):
                text = part.get("text")
                if isinstance(text, str):
                    parts.append(text)
        if parts:
            return "".join(parts)

    return None


def _get_attr(obj: Any, name: str, default: Any = None) -> Any:
    if isinstance(obj, dict):
        return obj.get(name, default)
    return getattr(obj, name, default)


class OpenAIAdapter(LLMAdapter):
    """
    Adapter for the OpenAI Python SDK (OpenAI / AsyncOpenAI) chat.completions.create.
    """

    def extract_input(self, args: tuple[Any, ...], kwargs: Dict[str, Any]) -> Optional[str]:
        req: Any = kwargs
        if args and isinstance(args[0], dict):
            req = args[0]
        messages = req.get("messages") if isinstance(req, dict) else kwargs.get("messages")
        return _extract_last_user_message_text(messages)

    def get_model_name(
        self, args: tuple[Any, ...], kwargs: Dict[str, Any], response: Any | None = None
    ) -> Optional[str]:
        if response is not None:
            model = _get_attr(response, "model", None)
            if isinstance(model, str) and model:
                return model

        req: Any = kwargs
        if args and isinstance(args[0], dict):
            req = args[0]
        model = req.get("model") if isinstance(req, dict) else kwargs.get("model")
        return model if isinstance(model, str) and model else None

    def extract_output(self, response: Any) -> str:
        choices = _get_attr(response, "choices", None)
        if isinstance(choices, list) and choices:
            first = choices[0]
            message = _get_attr(first, "message", None)
            content = _get_attr(message, "content", None)
            if isinstance(content, str):
                return content
        return ""

    def get_token_usage(self, response: Any) -> Optional[Dict[str, int]]:
        usage = _get_attr(response, "usage", None)

        # OpenAI may return a dict-like or an object with attributes.
        prompt = _get_attr(usage, "prompt_tokens", None)
        completion = _get_attr(usage, "completion_tokens", None)
        total = _get_attr(usage, "total_tokens", None)

        def _to_int(val: Any) -> Optional[int]:
            try:
                num = int(val)
                return num if num >= 0 else None
            except Exception:
                return None

        prompt_i = _to_int(prompt)
        completion_i = _to_int(completion)
        total_i = _to_int(total)

        if prompt_i is None and completion_i is None and total_i is None:
            return None

        out: Dict[str, int] = {}
        if prompt_i is not None:
            out["prompt_tokens"] = prompt_i
        if completion_i is not None:
            out["completion_tokens"] = completion_i
        if total_i is not None:
            out["total_tokens"] = total_i
        return out

