"""OpenAI SDK monkey-patches for auto-instrumentation."""

from __future__ import annotations

import json
import time
from datetime import UTC, datetime
from typing import Any

import openai

from caliper._annotation import generate_request_id, set_last_request_id
from caliper._features import get_linked_request_id
from caliper._patch_helpers import (
    capture_tools_enabled,
    get_sdk_version,
    merge_metadata,
    pop_caliper_metadata,
    push_record,
)
from caliper._patch_registry import get_original, save_and_patch
from caliper._record import Record

# ---------------------------------------------------------------------------
# OpenAI-specific helpers
# ---------------------------------------------------------------------------


def _classify_error(exc: Exception) -> tuple[str, int | None, str | None]:
    """Map OpenAI exception types to (status, http_status, error_code)."""
    if isinstance(exc, openai.RateLimitError):
        return "rate_limited", 429, getattr(exc, "code", None) or "rate_limit_error"
    if isinstance(exc, openai.APITimeoutError):
        return "timeout", None, None
    if isinstance(exc, openai.APIConnectionError):
        return "timeout", None, None
    if isinstance(exc, openai.APIStatusError):
        return "error", exc.status_code, getattr(exc, "code", None)
    return "error", None, None


def _extract_provider_request_id(response: Any) -> str | None:
    """Extract provider's correlation ID from response."""
    return getattr(response, "_request_id", None)


def _extract_tools_provided(kwargs: dict) -> list[str] | None:
    """Extract tool names from the request kwargs (OpenAI function-call format)."""
    tools = kwargs.get("tools")
    if not tools:
        return None
    names = []
    for t in tools:
        if isinstance(t, dict):
            fn = t.get("function", {})
            name = fn.get("name", "") if isinstance(fn, dict) else ""
            names.append(name)
    return names if names else None


def _extract_tool_calls(response: Any) -> str | None:
    """Extract tool calls from response.choices[0].message.tool_calls as JSON string."""
    choices = getattr(response, "choices", None)
    if not choices:
        return None
    message = getattr(choices[0], "message", None)
    if message is None:
        return None
    tool_calls = getattr(message, "tool_calls", None)
    if not tool_calls:
        return None
    result = []
    for tc in tool_calls:
        fn = getattr(tc, "function", None)
        result.append(
            {
                "name": getattr(fn, "name", "") if fn else "",
                "id": getattr(tc, "id", ""),
                "arguments": getattr(fn, "arguments", "") if fn else "",
            }
        )
    return json.dumps(result) if result else None


# ---------------------------------------------------------------------------
# Record builders (OpenAI-specific)
# ---------------------------------------------------------------------------


def _build_record_from_response(
    response: Any,
    kwargs: dict,
    duration_ms: int,
    request_meta: dict[str, str],
    request_id: str,
) -> Record:
    """Build a Record from a successful non-streaming OpenAI response."""
    usage = getattr(response, "usage", None)
    input_tokens = usage.prompt_tokens if usage else 0
    output_tokens = usage.completion_tokens if usage else 0

    # Detect cached tokens via prompt_tokens_details.cached_tokens
    cache_read = 0
    if usage:
        details = getattr(usage, "prompt_tokens_details", None)
        if details:
            cache_read = getattr(details, "cached_tokens", 0) or 0

    tools_provided = _extract_tools_provided(kwargs) if capture_tools_enabled() else None
    tool_calls = _extract_tool_calls(response) if capture_tools_enabled() else None

    return Record(
        request_id=request_id,
        timestamp=datetime.now(UTC).isoformat(),
        provider="openai",
        model=response.model,
        endpoint="chat",
        duration_ms=duration_ms,
        tokens_input=input_tokens,
        tokens_output=output_tokens,
        tokens_total=input_tokens + output_tokens,
        status="success",
        http_status=200,
        cached=cache_read > 0,
        provider_request_id=_extract_provider_request_id(response),
        linked_request_id=get_linked_request_id(),
        tools_provided=tools_provided,
        tool_calls=tool_calls,
        metadata=merge_metadata(request_meta),
        sdk_version=get_sdk_version(),
    )


def _build_record_from_error(
    exc: Exception,
    kwargs: dict,
    duration_ms: int,
    request_meta: dict[str, str],
    request_id: str,
) -> Record:
    """Build a Record from a failed OpenAI call."""
    status, http_status, error_code = _classify_error(exc)

    return Record(
        request_id=request_id,
        timestamp=datetime.now(UTC).isoformat(),
        provider="openai",
        model=kwargs.get("model", "unknown"),
        endpoint="chat",
        duration_ms=duration_ms,
        tokens_input=0,
        tokens_output=0,
        tokens_total=0,
        status=status,
        http_status=http_status,
        error_code=error_code,
        cached=False,
        linked_request_id=get_linked_request_id(),
        metadata=merge_metadata(request_meta),
        sdk_version=get_sdk_version(),
    )


# ---------------------------------------------------------------------------
# Tracked wrappers — cannot use make_tracked_call because create() must
# branch on stream=True/False
# ---------------------------------------------------------------------------


def _tracked_create(self, **kwargs):
    """Sync tracked wrapper for chat.completions.create."""
    if kwargs.get("stream"):
        return _tracked_stream_create(self, **kwargs)

    request_meta = pop_caliper_metadata(kwargs)
    request_id = generate_request_id()
    start = time.monotonic()
    error = None
    response = None
    try:
        original_fn = get_original(openai.resources.chat.Completions, "create")
        response = original_fn(self, **kwargs)
        return response
    except Exception as exc:
        error = exc
        raise
    finally:
        try:
            duration_ms = int((time.monotonic() - start) * 1000)
            if error is not None:
                record = _build_record_from_error(error, kwargs, duration_ms, request_meta, request_id)
            elif response is not None:
                record = _build_record_from_response(response, kwargs, duration_ms, request_meta, request_id)
            else:
                record = None
            if record is not None:
                push_record(record)
            set_last_request_id(request_id)
        except Exception:
            pass


async def _tracked_async_create(self, **kwargs):
    """Async tracked wrapper for chat.completions.create."""
    if kwargs.get("stream"):
        return await _tracked_async_stream_create(self, **kwargs)

    request_meta = pop_caliper_metadata(kwargs)
    request_id = generate_request_id()
    start = time.monotonic()
    error = None
    response = None
    try:
        original_fn = get_original(openai.resources.chat.AsyncCompletions, "create")
        response = await original_fn(self, **kwargs)
        return response
    except Exception as exc:
        error = exc
        raise
    finally:
        try:
            duration_ms = int((time.monotonic() - start) * 1000)
            if error is not None:
                record = _build_record_from_error(error, kwargs, duration_ms, request_meta, request_id)
            elif response is not None:
                record = _build_record_from_response(response, kwargs, duration_ms, request_meta, request_id)
            else:
                record = None
            if record is not None:
                push_record(record)
            set_last_request_id(request_id)
        except Exception:
            pass


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


def _tracked_stream_create(self, **kwargs):
    """Sync streaming: call original create(stream=True), wrap in _PatchedStream."""
    request_meta = pop_caliper_metadata(kwargs)
    request_id = generate_request_id()
    original_fn = get_original(openai.resources.chat.Completions, "create")
    stream = original_fn(self, **kwargs)
    return _PatchedStream(stream, kwargs, request_meta, request_id)


async def _tracked_async_stream_create(self, **kwargs):
    """Async streaming: call original create(stream=True), wrap in _AsyncPatchedStream."""
    request_meta = pop_caliper_metadata(kwargs)
    request_id = generate_request_id()
    original_fn = get_original(openai.resources.chat.AsyncCompletions, "create")
    stream = await original_fn(self, **kwargs)
    return _AsyncPatchedStream(stream, kwargs, request_meta, request_id)


# ---------------------------------------------------------------------------
# Chunk processing
# ---------------------------------------------------------------------------


def _process_chunk(chunk: Any, state: dict) -> str | None:
    """Process a ChatCompletionChunk, update state, return text delta if any.

    OpenAI streams tool calls interleaved by index — we accumulate them
    in a dict keyed by index and finalize at the end.
    """
    # Extract model from first chunk
    model = getattr(chunk, "model", None)
    if model and "model" not in state:
        state["model"] = model

    # Extract provider request ID from first chunk
    if "provider_request_id" not in state:
        req_id = getattr(chunk, "_request_id", None)
        if req_id:
            state["provider_request_id"] = req_id

    choices = getattr(chunk, "choices", None)
    if not choices:
        # Final chunk may carry usage without choices
        usage = getattr(chunk, "usage", None)
        if usage:
            state["input_tokens"] = getattr(usage, "prompt_tokens", 0) or 0
            state["output_tokens"] = getattr(usage, "completion_tokens", 0) or 0
        return None

    choice = choices[0]
    delta = getattr(choice, "delta", None)

    # Text content
    text = getattr(delta, "content", None) if delta else None
    if text:
        if state.get("first_token_time") is None:
            state["first_token_time"] = time.monotonic()
        return text

    # Tool calls (streamed by index)
    if capture_tools_enabled() and delta:
        tool_calls = getattr(delta, "tool_calls", None)
        if tool_calls:
            tool_dict = state.setdefault("_tool_calls", {})
            for tc in tool_calls:
                idx = getattr(tc, "index", 0)
                fn = getattr(tc, "function", None)
                if idx not in tool_dict:
                    tool_dict[idx] = {
                        "name": getattr(fn, "name", "") if fn else "",
                        "id": getattr(tc, "id", "") or "",
                        "_arg_parts": [],
                    }
                if fn:
                    args_chunk = getattr(fn, "arguments", "")
                    if args_chunk:
                        tool_dict[idx]["_arg_parts"].append(args_chunk)

    # Finish reason
    finish_reason = getattr(choice, "finish_reason", None)
    if finish_reason:
        state["finish_reason"] = finish_reason

    # Usage on final chunk (requires stream_options={"include_usage": True})
    usage = getattr(chunk, "usage", None)
    if usage:
        state["input_tokens"] = getattr(usage, "prompt_tokens", 0) or 0
        state["output_tokens"] = getattr(usage, "completion_tokens", 0) or 0

    return None


# ---------------------------------------------------------------------------
# _PatchedStream — wraps OpenAI Stream (not a context manager by necessity)
# ---------------------------------------------------------------------------


class _PatchedStream:
    """Wraps an OpenAI Stream to capture streaming metrics."""

    def __init__(self, stream: Any, kwargs: dict, request_meta: dict[str, str], request_id: str) -> None:
        self._stream = stream
        self._kwargs = kwargs
        self._request_meta = request_meta
        self._request_id = request_id
        self._state: dict = {}
        self._start: float = time.monotonic()
        self._finalized: bool = False
        set_last_request_id(self._request_id)

    def __enter__(self):
        self._stream.__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            self._finalize()
        except Exception:
            pass
        return self._stream.__exit__(exc_type, exc_val, exc_tb)

    def __iter__(self):
        try:
            for chunk in self._stream:
                _process_chunk(chunk, self._state)
                yield chunk
        finally:
            try:
                self._finalize()
            except Exception:
                pass

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

    def _finalize(self) -> None:
        if self._finalized:
            return
        self._finalized = True

        elapsed_ms = int((time.monotonic() - self._start) * 1000)
        first_token_time = self._state.get("first_token_time")
        ttft_ms = int((first_token_time - self._start) * 1000) if first_token_time else None
        input_tokens = self._state.get("input_tokens", 0)
        output_tokens = self._state.get("output_tokens", 0)

        # Detect caching from usage details if available
        cached = False

        tools_provided = _extract_tools_provided(self._kwargs) if capture_tools_enabled() else None
        tool_calls = self._build_tool_calls_json() if capture_tools_enabled() else None

        record = Record(
            request_id=self._request_id,
            timestamp=datetime.now(UTC).isoformat(),
            provider="openai",
            model=self._state.get("model", self._kwargs.get("model", "unknown")),
            endpoint="chat",
            duration_ms=elapsed_ms,
            ttft_ms=ttft_ms,
            tokens_input=input_tokens,
            tokens_output=output_tokens,
            tokens_total=input_tokens + output_tokens,
            status="success",
            http_status=200,
            cached=cached,
            provider_request_id=self._state.get("provider_request_id"),
            linked_request_id=get_linked_request_id(),
            tools_provided=tools_provided,
            tool_calls=tool_calls,
            metadata=merge_metadata(self._request_meta),
            sdk_version=get_sdk_version(),
        )
        push_record(record)

    def _build_tool_calls_json(self) -> str | None:
        tool_dict = self._state.get("_tool_calls")
        if not tool_dict:
            return None
        result = []
        for idx in sorted(tool_dict):
            entry = tool_dict[idx]
            args_raw = "".join(entry.pop("_arg_parts", []))
            result.append(
                {
                    "name": entry.get("name", ""),
                    "id": entry.get("id", ""),
                    "arguments": args_raw,
                }
            )
        return json.dumps(result) if result else None


# ---------------------------------------------------------------------------
# _AsyncPatchedStream — async equivalent
# ---------------------------------------------------------------------------


class _AsyncPatchedStream:
    """Wraps an OpenAI AsyncStream to capture streaming metrics."""

    def __init__(self, stream: Any, kwargs: dict, request_meta: dict[str, str], request_id: str) -> None:
        self._stream = stream
        self._kwargs = kwargs
        self._request_meta = request_meta
        self._request_id = request_id
        self._state: dict = {}
        self._start: float = time.monotonic()
        self._finalized: bool = False
        set_last_request_id(self._request_id)

    async def __aenter__(self):
        await self._stream.__aenter__()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        try:
            self._finalize()
        except Exception:
            pass
        return await self._stream.__aexit__(exc_type, exc_val, exc_tb)

    async def __aiter__(self):
        try:
            async for chunk in self._stream:
                _process_chunk(chunk, self._state)
                yield chunk
        finally:
            try:
                self._finalize()
            except Exception:
                pass

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

    def _finalize(self) -> None:
        if self._finalized:
            return
        self._finalized = True

        elapsed_ms = int((time.monotonic() - self._start) * 1000)
        first_token_time = self._state.get("first_token_time")
        ttft_ms = int((first_token_time - self._start) * 1000) if first_token_time else None
        input_tokens = self._state.get("input_tokens", 0)
        output_tokens = self._state.get("output_tokens", 0)

        cached = False

        tools_provided = _extract_tools_provided(self._kwargs) if capture_tools_enabled() else None
        tool_calls = self._build_tool_calls_json() if capture_tools_enabled() else None

        record = Record(
            request_id=self._request_id,
            timestamp=datetime.now(UTC).isoformat(),
            provider="openai",
            model=self._state.get("model", self._kwargs.get("model", "unknown")),
            endpoint="chat",
            duration_ms=elapsed_ms,
            ttft_ms=ttft_ms,
            tokens_input=input_tokens,
            tokens_output=output_tokens,
            tokens_total=input_tokens + output_tokens,
            status="success",
            http_status=200,
            cached=cached,
            provider_request_id=self._state.get("provider_request_id"),
            linked_request_id=get_linked_request_id(),
            tools_provided=tools_provided,
            tool_calls=tool_calls,
            metadata=merge_metadata(self._request_meta),
            sdk_version=get_sdk_version(),
        )
        push_record(record)

    def _build_tool_calls_json(self) -> str | None:
        tool_dict = self._state.get("_tool_calls")
        if not tool_dict:
            return None
        result = []
        for idx in sorted(tool_dict):
            entry = tool_dict[idx]
            args_raw = "".join(entry.pop("_arg_parts", []))
            result.append(
                {
                    "name": entry.get("name", ""),
                    "id": entry.get("id", ""),
                    "arguments": args_raw,
                }
            )
        return json.dumps(result) if result else None


# ---------------------------------------------------------------------------
# Patch entry point
# ---------------------------------------------------------------------------


def patch() -> None:
    """Patch OpenAI chat.completions.create (sync + async)."""
    save_and_patch(openai.resources.chat.Completions, "create", _tracked_create)
    save_and_patch(openai.resources.chat.AsyncCompletions, "create", _tracked_async_create)
