"""OpenAI client wrapper for automatic logging."""

from __future__ import annotations

import os
import time
import uuid
from functools import wraps
from typing import TYPE_CHECKING, Any, Callable, Iterator

if TYPE_CHECKING:
    from openai import OpenAI
    from openai.types.chat import ChatCompletion, ChatCompletionChunk
    from openai.types.audio import Transcription, Translation
    from aitracer.client import AITracer


def wrap_openai_client(client: "OpenAI", tracer: "AITracer") -> "OpenAI":
    """
    Wrap an OpenAI client to automatically log all API calls.

    Args:
        client: OpenAI client instance.
        tracer: AITracer instance.

    Returns:
        Wrapped OpenAI client (same instance, modified in place).
    """
    # Wrap chat.completions.create
    _wrap_chat_completions(client, tracer)

    # Wrap audio.transcriptions.create and audio.translations.create
    _wrap_audio_transcriptions(client, tracer)
    _wrap_audio_translations(client, tracer)

    return client


def _wrap_chat_completions(client: "OpenAI", tracer: "AITracer") -> None:
    """Wrap chat.completions.create method."""
    original_create = client.chat.completions.create

    @wraps(original_create)
    def wrapped_create(*args: Any, **kwargs: Any) -> Any:
        """Wrapped chat.completions.create method."""
        start_time = time.time()
        span_id = str(uuid.uuid4())

        # Extract request data
        model = kwargs.get("model", args[0] if args else "unknown")
        messages = kwargs.get("messages", args[1] if len(args) > 1 else [])
        stream = kwargs.get("stream", False)

        try:
            response = original_create(*args, **kwargs)

            if stream:
                # Handle streaming response
                return _wrap_stream_response(
                    response=response,
                    tracer=tracer,
                    model=model,
                    messages=messages,
                    start_time=start_time,
                    span_id=span_id,
                )
            else:
                # Handle non-streaming response
                latency_ms = int((time.time() - start_time) * 1000)
                _log_completion(
                    tracer=tracer,
                    model=model,
                    messages=messages,
                    response=response,
                    latency_ms=latency_ms,
                    span_id=span_id,
                )
                return response

        except Exception as e:
            # Log error
            latency_ms = int((time.time() - start_time) * 1000)
            tracer.log(
                model=model,
                provider="openai",
                input_data={"messages": _serialize_messages(messages)},
                output_data=None,
                latency_ms=latency_ms,
                status="error",
                error_message=str(e),
                span_id=span_id,
            )
            raise

    # Replace method
    client.chat.completions.create = wrapped_create  # type: ignore


def _wrap_stream_response(
    response: Iterator["ChatCompletionChunk"],
    tracer: "AITracer",
    model: str,
    messages: list,
    start_time: float,
    span_id: str,
) -> Iterator["ChatCompletionChunk"]:
    """Wrap streaming response to log after completion."""
    chunks: list[Any] = []
    content_parts: list[str] = []
    input_tokens = 0
    output_tokens = 0

    try:
        for chunk in response:
            chunks.append(chunk)

            # Accumulate content
            if chunk.choices and chunk.choices[0].delta.content:
                content_parts.append(chunk.choices[0].delta.content)

            # Get usage if available (some models include it in final chunk)
            if hasattr(chunk, "usage") and chunk.usage:
                input_tokens = chunk.usage.prompt_tokens or 0
                output_tokens = chunk.usage.completion_tokens or 0

            yield chunk

        # Log after stream completes
        latency_ms = int((time.time() - start_time) * 1000)
        full_content = "".join(content_parts)

        tracer.log(
            model=model,
            provider="openai",
            input_data={"messages": _serialize_messages(messages)},
            output_data={"content": full_content},
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            latency_ms=latency_ms,
            status="success",
            span_id=span_id,
        )

    except Exception as e:
        latency_ms = int((time.time() - start_time) * 1000)
        tracer.log(
            model=model,
            provider="openai",
            input_data={"messages": _serialize_messages(messages)},
            output_data=None,
            latency_ms=latency_ms,
            status="error",
            error_message=str(e),
            span_id=span_id,
        )
        raise


def _log_completion(
    tracer: "AITracer",
    model: str,
    messages: list,
    response: "ChatCompletion",
    latency_ms: int,
    span_id: str,
) -> None:
    """Log a non-streaming completion."""
    # Extract response data
    output_content = None
    if response.choices and response.choices[0].message:
        output_content = response.choices[0].message.content

    input_tokens = 0
    output_tokens = 0
    if response.usage:
        input_tokens = response.usage.prompt_tokens
        output_tokens = response.usage.completion_tokens

    tracer.log(
        model=model,
        provider="openai",
        input_data={"messages": _serialize_messages(messages)},
        output_data={"content": output_content},
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        latency_ms=latency_ms,
        status="success",
        span_id=span_id,
    )


def _serialize_messages(messages: list) -> list[dict]:
    """Serialize messages to JSON-compatible format."""
    result = []
    for msg in messages:
        if isinstance(msg, dict):
            result.append(msg)
        elif hasattr(msg, "model_dump"):
            result.append(msg.model_dump())
        elif hasattr(msg, "__dict__"):
            result.append(msg.__dict__)
        else:
            result.append({"content": str(msg)})
    return result


def _wrap_audio_transcriptions(client: "OpenAI", tracer: "AITracer") -> None:
    """Wrap audio.transcriptions.create method for Whisper API."""
    original_create = client.audio.transcriptions.create

    @wraps(original_create)
    def wrapped_create(*args: Any, **kwargs: Any) -> Any:
        """Wrapped audio.transcriptions.create method."""
        start_time = time.time()
        span_id = str(uuid.uuid4())

        # Extract request data
        model = kwargs.get("model", "whisper-1")
        language = kwargs.get("language")
        file = kwargs.get("file")

        # Get file info
        file_info = _get_file_info(file)

        try:
            response = original_create(*args, **kwargs)
            latency_ms = int((time.time() - start_time) * 1000)

            # Extract response text
            response_text = ""
            if hasattr(response, "text"):
                response_text = response.text
            elif isinstance(response, str):
                response_text = response

            tracer.log(
                model=model,
                provider="openai",
                input_data={
                    "type": "audio_transcription",
                    "file_name": file_info.get("name"),
                    "file_size_bytes": file_info.get("size"),
                    "language": language,
                },
                output_data={
                    "text": response_text,
                    "text_length": len(response_text),
                },
                input_tokens=0,  # Whisper doesn't use tokens
                output_tokens=0,
                latency_ms=latency_ms,
                status="success",
                span_id=span_id,
                metadata={
                    "audio_duration_seconds": file_info.get("estimated_duration"),
                },
            )
            return response

        except Exception as e:
            latency_ms = int((time.time() - start_time) * 1000)
            tracer.log(
                model=model,
                provider="openai",
                input_data={
                    "type": "audio_transcription",
                    "file_name": file_info.get("name"),
                    "file_size_bytes": file_info.get("size"),
                    "language": language,
                },
                output_data=None,
                latency_ms=latency_ms,
                status="error",
                error_message=str(e),
                span_id=span_id,
            )
            raise

    client.audio.transcriptions.create = wrapped_create  # type: ignore


def _wrap_audio_translations(client: "OpenAI", tracer: "AITracer") -> None:
    """Wrap audio.translations.create method for Whisper API."""
    original_create = client.audio.translations.create

    @wraps(original_create)
    def wrapped_create(*args: Any, **kwargs: Any) -> Any:
        """Wrapped audio.translations.create method."""
        start_time = time.time()
        span_id = str(uuid.uuid4())

        # Extract request data
        model = kwargs.get("model", "whisper-1")
        file = kwargs.get("file")

        # Get file info
        file_info = _get_file_info(file)

        try:
            response = original_create(*args, **kwargs)
            latency_ms = int((time.time() - start_time) * 1000)

            # Extract response text
            response_text = ""
            if hasattr(response, "text"):
                response_text = response.text
            elif isinstance(response, str):
                response_text = response

            tracer.log(
                model=model,
                provider="openai",
                input_data={
                    "type": "audio_translation",
                    "file_name": file_info.get("name"),
                    "file_size_bytes": file_info.get("size"),
                },
                output_data={
                    "text": response_text,
                    "text_length": len(response_text),
                },
                input_tokens=0,
                output_tokens=0,
                latency_ms=latency_ms,
                status="success",
                span_id=span_id,
                metadata={
                    "audio_duration_seconds": file_info.get("estimated_duration"),
                },
            )
            return response

        except Exception as e:
            latency_ms = int((time.time() - start_time) * 1000)
            tracer.log(
                model=model,
                provider="openai",
                input_data={
                    "type": "audio_translation",
                    "file_name": file_info.get("name"),
                    "file_size_bytes": file_info.get("size"),
                },
                output_data=None,
                latency_ms=latency_ms,
                status="error",
                error_message=str(e),
                span_id=span_id,
            )
            raise

    client.audio.translations.create = wrapped_create  # type: ignore


def _get_file_info(file: Any) -> dict:
    """Extract file information from file object."""
    info: dict[str, Any] = {
        "name": None,
        "size": None,
        "estimated_duration": None,
    }

    if file is None:
        return info

    # Get file name
    if hasattr(file, "name"):
        info["name"] = os.path.basename(file.name)

    # Get file size
    try:
        if hasattr(file, "seek") and hasattr(file, "tell"):
            current_pos = file.tell()
            file.seek(0, 2)  # Seek to end
            info["size"] = file.tell()
            file.seek(current_pos)  # Restore position
    except (OSError, IOError):
        pass

    return info
