"""OpenAI client wrapper for automatic token usage tracking."""

from __future__ import annotations

import inspect
import time
from typing import Any, Dict, Optional

from tokengrip.types import WrapOptions

_MAX_CONTENT_LEN = 500


def wrap_openai(client: Any, tg: Any, opts: Optional[WrapOptions] = None) -> Any:
    """Wrap an OpenAI client to automatically track token usage.

    Intercepts:
    - ``client.chat.completions.create()``
    - ``client.responses.create()``

    Both sync and async clients are supported.
    """
    return _WrappedOpenAIClient(client, tg, opts or {})


# ---------------------------------------------------------------------------
# __getattr__ wrapper chain
# ---------------------------------------------------------------------------


class _Wrapper:
    """Base wrapper delegating attribute access to the wrapped object."""

    __slots__ = ("_wrapped",)

    def __init__(self, wrapped: Any) -> None:
        object.__setattr__(self, "_wrapped", wrapped)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._wrapped, name)

    def __repr__(self) -> str:
        return repr(self._wrapped)


class _WrappedOpenAIClient(_Wrapper):
    __slots__ = ("_wrapped", "_tg", "_opts")

    def __init__(self, client: Any, tg: Any, opts: WrapOptions) -> None:
        super().__init__(client)
        object.__setattr__(self, "_tg", tg)
        object.__setattr__(self, "_opts", opts)

    def __getattr__(self, name: str) -> Any:
        attr = getattr(self._wrapped, name)
        if name == "chat":
            return _WrappedChat(attr, self._tg, self._opts)
        if name == "responses":
            return _WrappedResponses(attr, self._tg, self._opts)
        return attr


class _WrappedChat(_Wrapper):
    __slots__ = ("_wrapped", "_tg", "_opts")

    def __init__(self, chat: Any, tg: Any, opts: WrapOptions) -> None:
        super().__init__(chat)
        object.__setattr__(self, "_tg", tg)
        object.__setattr__(self, "_opts", opts)

    def __getattr__(self, name: str) -> Any:
        attr = getattr(self._wrapped, name)
        if name == "completions":
            return _WrappedCompletions(attr, self._tg, self._opts)
        return attr


class _WrappedCompletions(_Wrapper):
    __slots__ = ("_wrapped", "_tg", "_opts")

    def __init__(self, completions: Any, tg: Any, opts: WrapOptions) -> None:
        super().__init__(completions)
        object.__setattr__(self, "_tg", tg)
        object.__setattr__(self, "_opts", opts)

    def create(self, **kwargs: Any) -> Any:
        return _tracked_call(
            self._wrapped.create,
            self._tg,
            self._opts,
            "openai",
            "chat.completions.create",
            kwargs,
        )


class _WrappedResponses(_Wrapper):
    __slots__ = ("_wrapped", "_tg", "_opts")

    def __init__(self, responses: Any, tg: Any, opts: WrapOptions) -> None:
        super().__init__(responses)
        object.__setattr__(self, "_tg", tg)
        object.__setattr__(self, "_opts", opts)

    def create(self, **kwargs: Any) -> Any:
        return _tracked_call(
            self._wrapped.create,
            self._tg,
            self._opts,
            "openai",
            "responses.create",
            kwargs,
        )


# ---------------------------------------------------------------------------
# Tracked call dispatch (sync / async)
# ---------------------------------------------------------------------------


def _tracked_call(
    original_fn: Any,
    tg: Any,
    opts: WrapOptions,
    provider: str,
    endpoint: str,
    kwargs: Dict[str, Any],
) -> Any:
    start_time = time.monotonic()
    is_streaming = kwargs.get("stream", False)

    if inspect.iscoroutinefunction(original_fn):
        return _tracked_call_async(
            original_fn, tg, opts, provider, endpoint, kwargs, start_time, is_streaming
        )

    response = original_fn(**kwargs)

    if is_streaming:
        return _wrap_stream(
            response, tg, opts, provider, endpoint, kwargs, start_time
        )

    _track_from_response(response, tg, opts, provider, endpoint, kwargs, start_time)
    return response


async def _tracked_call_async(
    original_fn: Any,
    tg: Any,
    opts: WrapOptions,
    provider: str,
    endpoint: str,
    kwargs: Dict[str, Any],
    start_time: float,
    is_streaming: bool,
) -> Any:
    response = await original_fn(**kwargs)

    if is_streaming:
        return _wrap_stream(
            response, tg, opts, provider, endpoint, kwargs, start_time
        )

    _track_from_response(response, tg, opts, provider, endpoint, kwargs, start_time)
    return response


# ---------------------------------------------------------------------------
# Non-streaming usage extraction
# ---------------------------------------------------------------------------


def _track_from_response(
    response: Any,
    tg: Any,
    opts: WrapOptions,
    provider: str,
    endpoint: str,
    request_kwargs: Dict[str, Any],
    start_time: float,
) -> None:
    try:
        latency_ms = int((time.monotonic() - start_time) * 1000)
        model = getattr(response, "model", None) or request_kwargs.get(
            "model", "unknown"
        )

        usage = getattr(response, "usage", None)
        if usage is None:
            return

        if endpoint == "responses.create":
            input_tokens = getattr(usage, "input_tokens", 0) or 0
            output_tokens = getattr(usage, "output_tokens", 0) or 0
        else:
            input_tokens = getattr(usage, "prompt_tokens", 0) or 0
            output_tokens = getattr(usage, "completion_tokens", 0) or 0

        if input_tokens == 0 and output_tokens == 0:
            return

        payload: Dict[str, Any] = {
            "provider": provider,
            "model": model,
            "inputTokens": input_tokens,
            "outputTokens": output_tokens,
            "latencyMs": latency_ms,
            "projectId": opts.get("project_id"),
            "agentName": opts.get("agent_name"),
            "taskType": opts.get("task_type"),
        }

        if opts.get("capture_content"):
            payload["promptContent"] = _extract_prompt(request_kwargs, endpoint)
            payload["responseContent"] = _extract_response_content(
                response, endpoint
            )

        tg.track(payload)
    except Exception:
        pass


# ---------------------------------------------------------------------------
# Streaming wrappers
# ---------------------------------------------------------------------------


def _wrap_stream(
    stream: Any,
    tg: Any,
    opts: WrapOptions,
    provider: str,
    endpoint: str,
    request_kwargs: Dict[str, Any],
    start_time: float,
) -> Any:
    if hasattr(stream, "__aiter__"):
        return _WrappedAsyncStream(
            stream, tg, opts, provider, endpoint, request_kwargs, start_time
        )
    if hasattr(stream, "__iter__"):
        return _WrappedSyncStream(
            stream, tg, opts, provider, endpoint, request_kwargs, start_time
        )
    return stream


class _StreamMixin:
    """Shared chunk processing and tracking logic for stream wrappers."""

    _tg: Any
    _opts: WrapOptions
    _provider: str
    _endpoint: str
    _start_time: float
    _input_tokens: int
    _output_tokens: int
    _model: str
    _tracked: bool

    def _process_chunk(self, chunk: Any) -> None:
        try:
            model = getattr(chunk, "model", None)
            if model:
                self._model = model

            usage = getattr(chunk, "usage", None)
            if usage is not None:
                if self._endpoint == "responses.create":
                    self._input_tokens = (
                        getattr(usage, "input_tokens", None) or self._input_tokens
                    )
                    self._output_tokens = (
                        getattr(usage, "output_tokens", None) or self._output_tokens
                    )
                else:
                    self._input_tokens = (
                        getattr(usage, "prompt_tokens", None) or self._input_tokens
                    )
                    self._output_tokens = (
                        getattr(usage, "completion_tokens", None)
                        or self._output_tokens
                    )
        except Exception:
            pass

    def _do_track(self) -> None:
        if self._tracked:
            return
        self._tracked = True
        try:
            if self._input_tokens > 0 or self._output_tokens > 0:
                latency_ms = int((time.monotonic() - self._start_time) * 1000)
                self._tg.track(
                    {
                        "provider": self._provider,
                        "model": self._model,
                        "inputTokens": self._input_tokens,
                        "outputTokens": self._output_tokens,
                        "latencyMs": latency_ms,
                        "projectId": self._opts.get("project_id"),
                        "agentName": self._opts.get("agent_name"),
                        "taskType": self._opts.get("task_type"),
                    }
                )
        except Exception:
            pass


class _WrappedSyncStream(_StreamMixin):
    def __init__(
        self,
        stream: Any,
        tg: Any,
        opts: WrapOptions,
        provider: str,
        endpoint: str,
        request_kwargs: Dict[str, Any],
        start_time: float,
    ) -> None:
        self._stream = stream
        self._iter = iter(stream)
        self._tg = tg
        self._opts = opts
        self._provider = provider
        self._endpoint = endpoint
        self._start_time = start_time
        self._input_tokens = 0
        self._output_tokens = 0
        self._model = request_kwargs.get("model", "unknown")
        self._tracked = False

    def __iter__(self):  # type: ignore[override]
        return self

    def __next__(self) -> Any:
        try:
            chunk = next(self._iter)
        except StopIteration:
            self._do_track()
            raise
        self._process_chunk(chunk)
        return chunk

    def __enter__(self) -> _WrappedSyncStream:
        if hasattr(self._stream, "__enter__"):
            self._stream.__enter__()
        return self

    def __exit__(self, *args: Any) -> Any:
        self._do_track()
        if hasattr(self._stream, "__exit__"):
            return self._stream.__exit__(*args)
        return None

    def __getattr__(self, name: str) -> Any:
        return getattr(self._stream, name)


class _WrappedAsyncStream(_StreamMixin):
    def __init__(
        self,
        stream: Any,
        tg: Any,
        opts: WrapOptions,
        provider: str,
        endpoint: str,
        request_kwargs: Dict[str, Any],
        start_time: float,
    ) -> None:
        self._stream = stream
        self._aiter: Any = None
        self._tg = tg
        self._opts = opts
        self._provider = provider
        self._endpoint = endpoint
        self._start_time = start_time
        self._input_tokens = 0
        self._output_tokens = 0
        self._model = request_kwargs.get("model", "unknown")
        self._tracked = False

    def __aiter__(self) -> _WrappedAsyncStream:
        return self

    async def __anext__(self) -> Any:
        if self._aiter is None:
            self._aiter = self._stream.__aiter__()
        try:
            chunk = await self._aiter.__anext__()
        except StopAsyncIteration:
            self._do_track()
            raise
        self._process_chunk(chunk)
        return chunk

    async def __aenter__(self) -> _WrappedAsyncStream:
        if hasattr(self._stream, "__aenter__"):
            await self._stream.__aenter__()
        return self

    async def __aexit__(self, *args: Any) -> Any:
        self._do_track()
        if hasattr(self._stream, "__aexit__"):
            return await self._stream.__aexit__(*args)
        return None

    def __getattr__(self, name: str) -> Any:
        return getattr(self._stream, name)


# ---------------------------------------------------------------------------
# Content extraction helpers
# ---------------------------------------------------------------------------


def _truncate(s: str, max_len: int = _MAX_CONTENT_LEN) -> str:
    return s[:max_len] if len(s) > max_len else s


def _extract_prompt(
    request_kwargs: Dict[str, Any], endpoint: str
) -> Optional[str]:
    try:
        if endpoint == "responses.create":
            inp = request_kwargs.get("input")
            if isinstance(inp, str):
                return _truncate(inp)
            return None
        messages = request_kwargs.get("messages")
        if not messages:
            return None
        last = messages[-1]
        content = (
            last.get("content") if isinstance(last, dict) else getattr(last, "content", None)
        )
        if isinstance(content, str):
            return _truncate(content)
        return None
    except Exception:
        return None


def _extract_response_content(response: Any, endpoint: str) -> Optional[str]:
    try:
        if endpoint == "responses.create":
            output_text = getattr(response, "output_text", None)
            if isinstance(output_text, str):
                return _truncate(output_text)
            return None
        choices = getattr(response, "choices", None)
        if not choices:
            return None
        message = getattr(choices[0], "message", None)
        content = getattr(message, "content", None) if message else None
        if isinstance(content, str):
            return _truncate(content)
        return None
    except Exception:
        return None
