"""Chat completion functions — openai SDK imported lazily only when needed."""

from litelm._client_cache import get_async_client, get_sync_client
from litelm._dispatch import get_handler
from litelm._exceptions import (
    APIConnectionError,
    APIStatusError,
    AuthenticationError,
    BadRequestError,
    ContextWindowExceededError,
    InternalServerError,
    LitelmError,
    NotFoundError,
    PermissionDeniedError,
    RateLimitError,
    Timeout,
    UnprocessableEntityError,
    is_context_window_error,
)
from litelm._providers import parse_model
from litelm._types import (
    ChatCompletion,
    ChatCompletionChunk,
    ChatCompletionMessage,
    ChatCompletionMessageToolCall,
    Choice,
    ChoiceDelta,
    ChunkChoice,
    CompletionUsage,
    Function,
    ModelResponse,
    ModelResponseStream,
)

# Build tuple of all openai SDK error types to catch
_openai_errors: tuple = ()
_bad_request_errors = [BadRequestError]
try:
    import openai as _openai

    _bad_request_errors.append(_openai.BadRequestError)
    _openai_errors = (_openai.APIError,)
except ImportError:
    pass
_bad_request_errors = tuple(_bad_request_errors)


def _map_openai_error(e):
    """Map openai SDK exceptions to litelm exception types. Re-raises as litelm error."""
    try:
        import openai
    except ImportError:
        raise LitelmError(message=str(e)) from e

    msg = str(e)

    if isinstance(e, openai.BadRequestError):
        if is_context_window_error(msg):
            raise ContextWindowExceededError(
                message=msg, response=getattr(e, "response", None), body=getattr(e, "body", None)
            ) from e
        raise BadRequestError(message=msg, response=getattr(e, "response", None), body=getattr(e, "body", None)) from e
    elif isinstance(e, openai.RateLimitError):
        raise RateLimitError(message=msg, response=getattr(e, "response", None), body=getattr(e, "body", None)) from e
    elif isinstance(e, openai.AuthenticationError):
        raise AuthenticationError(
            message=msg, response=getattr(e, "response", None), body=getattr(e, "body", None)
        ) from e
    elif isinstance(e, openai.APITimeoutError):
        raise Timeout(request=getattr(e, "request", None)) from e
    elif isinstance(e, openai.APIConnectionError):
        raise APIConnectionError(request=getattr(e, "request", None)) from e
    elif isinstance(e, openai.InternalServerError):
        raise InternalServerError(
            message=msg, response=getattr(e, "response", None), body=getattr(e, "body", None)
        ) from e
    elif isinstance(e, openai.NotFoundError):
        raise NotFoundError(message=msg, response=getattr(e, "response", None), body=getattr(e, "body", None)) from e
    elif isinstance(e, openai.PermissionDeniedError):
        raise PermissionDeniedError(
            message=msg, response=getattr(e, "response", None), body=getattr(e, "body", None)
        ) from e
    elif isinstance(e, openai.UnprocessableEntityError):
        raise UnprocessableEntityError(
            message=msg, response=getattr(e, "response", None), body=getattr(e, "body", None)
        ) from e
    elif isinstance(e, openai.APIStatusError):
        raise APIStatusError(message=msg, response=getattr(e, "response", None), body=getattr(e, "body", None)) from e
    raise LitelmError(message=msg) from e


# kwargs that are litellm-specific and must be stripped before passing to OpenAI SDK
_LITELLM_ONLY_KWARGS = {"cache", "num_retries", "retry_strategy", "caching"}


def _prepare_call(model, kwargs):
    """Parse model, build client kwargs, strip litellm-specific params."""
    num_retries = kwargs.pop("num_retries", 0)
    kwargs.pop("retry_strategy", None)
    kwargs.pop("cache", None)
    kwargs.pop("caching", None)
    kwargs.pop("custom_llm_provider", None)
    kwargs.pop("fallbacks", None)
    kwargs.pop("mock_timeout", None)

    api_key = kwargs.pop("api_key", None)
    api_base = kwargs.pop("api_base", None) or kwargs.pop("base_url", None)
    api_version = kwargs.pop("api_version", None)
    headers = kwargs.pop("headers", None)

    provider, model_name, base_url, resolved_api_key, resolved_api_version = parse_model(
        model, api_key=api_key, api_base=api_base, api_version=api_version
    )

    if headers:
        kwargs["extra_headers"] = headers

    return provider, model_name, base_url, resolved_api_key, resolved_api_version or api_version, num_retries, kwargs


def _wrap_context_window_error(e):
    """Convert BadRequestError to ContextWindowExceededError if about context length."""
    if is_context_window_error(str(e)):
        raise ContextWindowExceededError(
            message=str(e),
            response=getattr(e, "response", None),
            body=getattr(e, "body", None),
        ) from e
    raise e


def completion(model, messages=None, *, timeout=None, stream=False, shared_session=None, **kwargs):
    """Synchronous chat completion."""
    mock = kwargs.pop("mock_response", None)
    n = kwargs.pop("n", None) or 1
    provider, model_name, base_url, api_key, api_version, num_retries, kwargs = _prepare_call(model, kwargs)

    if mock is not None:
        content = str(mock) if mock is not True else "This is a mock request"
        if stream:
            return _mock_stream_sync(content, model_name)
        return ModelResponse(
            ChatCompletion(
                id="mock",
                choices=[
                    Choice(
                        index=i, message=ChatCompletionMessage(role="assistant", content=content), finish_reason="stop"
                    )
                    for i in range(n)
                ],
                created=0,
                model=model_name,
                object="chat.completion",
                usage=CompletionUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
            )
        )

    if n > 1:
        kwargs["n"] = n

    handler = get_handler(provider)
    if handler:
        return handler.completion(
            model_name, messages, stream=stream, api_key=api_key, base_url=base_url, timeout=timeout, **kwargs
        )

    client = get_sync_client(provider, base_url, api_key, max_retries=num_retries, api_version=api_version)

    try:
        sdk_kwargs = dict(model=model_name, messages=messages, stream=stream, **kwargs)
        if timeout is not None:
            sdk_kwargs["timeout"] = timeout
        response = client.chat.completions.create(**sdk_kwargs)
    except _bad_request_errors as e:
        _wrap_context_window_error(e)
    except _openai_errors as e:
        _map_openai_error(e)

    if stream:
        return _wrap_stream_sync(response)
    return ModelResponse(response)


async def acompletion(model, messages=None, *, timeout=None, stream=False, shared_session=None, **kwargs):
    """Async chat completion."""
    mock = kwargs.pop("mock_response", None)
    n = kwargs.pop("n", None) or 1
    provider, model_name, base_url, api_key, api_version, num_retries, kwargs = _prepare_call(model, kwargs)

    if mock is not None:
        content = str(mock) if mock is not True else "This is a mock request"
        if stream:
            return _mock_stream_async(content, model_name)
        return ModelResponse(
            ChatCompletion(
                id="mock",
                choices=[
                    Choice(
                        index=i, message=ChatCompletionMessage(role="assistant", content=content), finish_reason="stop"
                    )
                    for i in range(n)
                ],
                created=0,
                model=model_name,
                object="chat.completion",
                usage=CompletionUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
            )
        )

    if n > 1:
        kwargs["n"] = n

    handler = get_handler(provider)
    if handler:
        return await handler.acompletion(
            model_name, messages, stream=stream, api_key=api_key, base_url=base_url, timeout=timeout, **kwargs
        )

    client = get_async_client(provider, base_url, api_key, max_retries=num_retries, api_version=api_version)

    try:
        sdk_kwargs = dict(model=model_name, messages=messages, stream=stream, **kwargs)
        if timeout is not None:
            sdk_kwargs["timeout"] = timeout
        response = await client.chat.completions.create(**sdk_kwargs)
    except _bad_request_errors as e:
        _wrap_context_window_error(e)
    except _openai_errors as e:
        _map_openai_error(e)

    if stream:
        return _wrap_stream_async(response)
    return ModelResponse(response)


def mock_completion(model, messages, n=1, stream=False, **kwargs):
    """Create a mock completion response."""
    return completion(model, messages, stream=stream, mock_response=True, n=n, **kwargs)


def _mock_stream_sync(content, model):
    """Yield mock streaming chunks."""
    yield ModelResponseStream(
        ChatCompletionChunk(
            id="mock",
            model=model,
            choices=[ChunkChoice(delta=ChoiceDelta(role="assistant", content=content))],
        )
    )
    yield ModelResponseStream(
        ChatCompletionChunk(
            id="mock",
            model=model,
            choices=[ChunkChoice(finish_reason="stop")],
        )
    )


async def _mock_stream_async(content, model):
    """Yield mock streaming chunks (async)."""
    yield ModelResponseStream(
        ChatCompletionChunk(
            id="mock",
            model=model,
            choices=[ChunkChoice(delta=ChoiceDelta(role="assistant", content=content))],
        )
    )
    yield ModelResponseStream(
        ChatCompletionChunk(
            id="mock",
            model=model,
            choices=[ChunkChoice(finish_reason="stop")],
        )
    )


def _wrap_stream_sync(stream):
    """Wrap sync stream to yield ModelResponseStream objects."""
    for chunk in stream:
        yield ModelResponseStream(chunk)


async def _wrap_stream_async(stream):
    """Wrap async stream to yield ModelResponseStream objects."""
    async for chunk in stream:
        yield ModelResponseStream(chunk)


def stream_chunk_builder(chunks):
    """Build a ModelResponse from a list of ModelResponseStream chunks."""
    if not chunks:
        return ModelResponse(
            ChatCompletion(
                id="empty",
                choices=[],
                created=0,
                model="",
                object="chat.completion",
            )
        )

    # Gather content, tool calls, and usage from chunks
    content_parts = []
    role = "assistant"
    model = ""
    chunk_id = ""
    tool_calls_by_index = {}
    images_by_index = {}
    usage = None
    finish_reason = None
    reasoning_content_parts = []

    for c in chunks:
        chunk = c._chunk if hasattr(c, "_chunk") else c
        chunk_id = chunk_id or chunk.id
        model = model or chunk.model

        if hasattr(chunk, "usage") and chunk.usage is not None:
            if usage is None:
                usage = CompletionUsage(
                    prompt_tokens=chunk.usage.prompt_tokens,
                    completion_tokens=chunk.usage.completion_tokens,
                    total_tokens=chunk.usage.total_tokens,
                )
            else:
                # Accumulate: take max of each field (Anthropic splits across chunks)
                usage.prompt_tokens = max(usage.prompt_tokens, chunk.usage.prompt_tokens)
                usage.completion_tokens = max(usage.completion_tokens, chunk.usage.completion_tokens)
                # Preserve API's total if it includes hidden categories (e.g. reasoning tokens)
                computed = usage.prompt_tokens + usage.completion_tokens
                usage.total_tokens = max(usage.total_tokens, chunk.usage.total_tokens, computed)

        if chunk.choices:
            choice = chunk.choices[0]
            finish_reason = choice.finish_reason or finish_reason
            delta = choice.delta
            if delta:
                if delta.role:
                    role = delta.role
                if delta.content:
                    content_parts.append(delta.content)
                # Collect reasoning_content if present
                rc = getattr(delta, "reasoning_content", None)
                if rc:
                    reasoning_content_parts.append(rc)
                if getattr(delta, "images", None):
                    for img in delta.images:
                        images_by_index[img["index"]] = img
                if delta.tool_calls:
                    for tc in delta.tool_calls:
                        idx = tc.index
                        if idx not in tool_calls_by_index:
                            tool_calls_by_index[idx] = {
                                "id": tc.id or "",
                                "type": "function",
                                "function": {"name": tc.function.name or "", "arguments": ""},
                            }
                        else:
                            if tc.id:
                                tool_calls_by_index[idx]["id"] = tc.id
                            if tc.function.name:
                                tool_calls_by_index[idx]["function"]["name"] = tc.function.name
                        if tc.function.arguments:
                            tool_calls_by_index[idx]["function"]["arguments"] += tc.function.arguments

    content = "".join(content_parts) or None
    tool_calls = None
    if tool_calls_by_index:
        tool_calls = [
            ChatCompletionMessageToolCall(
                id=tc["id"],
                type=tc["type"],
                function=Function(name=tc["function"]["name"], arguments=tc["function"]["arguments"]),
            )
            for _, tc in sorted(tool_calls_by_index.items())
        ]

    images = [v for _, v in sorted(images_by_index.items())] or None

    message = ChatCompletionMessage(
        role=role,
        content=content,
        tool_calls=tool_calls,
        reasoning_content="".join(reasoning_content_parts) if reasoning_content_parts else None,
        images=images,
    )

    if usage is None:
        usage = CompletionUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0)

    completion_obj = ChatCompletion(
        id=chunk_id or "chatcmpl-stream",
        choices=[Choice(index=0, message=message, finish_reason=finish_reason or "stop")],
        created=0,
        model=model,
        object="chat.completion",
        usage=usage,
    )

    return ModelResponse(completion_obj)
