import asyncio
import time
import traceback
from copy import deepcopy
from typing import Protocol, cast

import httpx
import litellm
from judy_core.configs import JudyCompletionOption
from litellm import Choices, CustomStreamWrapper
from litellm.exceptions import (
    InternalServerError,
    RateLimitError,
    ServiceUnavailableError,
)
from litellm.exceptions import Timeout as LitellmTimeout
from litellm.llms.vertex_ai.common_utils import _build_vertex_schema
from litellm.types.utils import Message, ModelResponse, ModelResponseStream


class AsyncLogger(Protocol):
    async def __call__(self, text: str) -> None: ...


BLANK_RESPONSE_RETRY_DEFAULT_MESSAGE = """Ignore this message. This message is only intended to trigger your reply. \
You should give a reasonable response based on the current conversation history, as if this message did not exist at all.
"""


def _remove_warnings_from_litellm_completion(messages):
    # 升级litellm版本后，出现大量如下类似警告，暂时用这种方式抑制，待litellmt解决这个问题后，这部分代码可以取出
    # pydantic/main.py:464:UserWarning: Pydantic serializer warnings:
    # PydanticSerializationUnexpectedValue(Expected `ChatCompletionMessageToolCall` - serialized value may not be as expected [field_name='tool_calls', input_value={'index': 0, 'provider_sp...+4', 'type': 'function'}, input_type=dict])
    new_messages = []
    for message in messages:
        if isinstance(message, Message):
            message = message.model_dump()
        new_messages.append(message)
    return new_messages


def adapt_tools_for_model(model: str, tools: list[dict] | None) -> list[dict] | None:
    if not tools:
        return tools

    if not ("openrouter" in model and "gemini" in model):
        return tools

    adapted_tools = []
    for tool in tools:
        adapted_tool = deepcopy(tool)
        function = adapted_tool.get("function")
        if function:
            parameters = function.get("parameters")
            if parameters:
                adapted_parameters = _build_vertex_schema(parameters)
                function["parameters"] = adapted_parameters
        adapted_tools.append(adapted_tool)
    return adapted_tools


def enable_claude_opus_thinking(args: dict) -> None:
    if "thinking" not in args:
        args["thinking"] = {"type": "adaptive"}
    allowed_openai_params = args.setdefault("allowed_openai_params", [])
    allowed_openai_params.append("thinking")


def adapt_completion_params(model: str, **kwargs) -> dict:
    if "gemini" in model:
        if "reasoning_effort" not in kwargs:
            kwargs["reasoning_effort"] = "high"
        # if "temperature" not in kwargs:
        #     kwargs["temperature"] = 1.0
        # kwargs["toolConfig"] = {
        #     "functionCallingConfig": {
        #         "stream_function_call_arguments": True,
        #     }
        # }
    elif "opus-4.6" in model or "opus-4-6" in model:
        if "max_tokens" not in kwargs:
            kwargs["max_tokens"] = 128 * 1024
        if "effort" not in kwargs:
            kwargs["effort"] = "max"
        if model.startswith("anthropic/"):
            enable_claude_opus_thinking(kwargs)
        else:
            extra_body = kwargs.setdefault("extra_body", {})
            enable_claude_opus_thinking(extra_body)
    return kwargs


class StreamChunkTimeout(Exception):
    def __init__(self, msg: str, ex: Exception):
        super().__init__(msg)
        self.ex = ex


JUDY_HIDDEN_PARAMS_KEY = "judy"
RESPONSE_TIME_COST = "response_time_cost"
FIRST_TOKEN_TIME_COST = "first_token_time_cost"
STREAM_TIME_COST = "stream_time_cost"
STREAM_CHUNKS_LEN = "stream_chunks_len"


async def _acompletion(
    model: str,
    *,
    messages,
    judy_option: JudyCompletionOption,
    stream: bool | None = None,
    stream_options: dict | None = None,
    **kwargs,
) -> ModelResponse:
    if stream:
        default_stream_options = {"include_usage": True}
        if stream_options is not None:
            default_stream_options.update(stream_options)
        stream_options = default_stream_options
    start_time = time.perf_counter()
    response = await litellm.acompletion(
        model,
        messages=messages,
        stream=stream,
        stream_options=stream_options,
        **kwargs,
    )
    response_time = time.perf_counter()
    response_time_cost = response_time - start_time
    judy_hidden_params = {
        RESPONSE_TIME_COST: response_time_cost,
    }
    if isinstance(response, CustomStreamWrapper):
        first_token_time = None
        chunks = []
        response_iter = aiter(response)
        while True:
            wait_start_time = time.perf_counter()
            try:
                async with asyncio.timeout(judy_option.stream_chunk_timeout):
                    chunk = await anext(response_iter)
                    if first_token_time is None:
                        first_token_time = time.perf_counter()
                        judy_hidden_params[FIRST_TOKEN_TIME_COST] = (
                            first_token_time - response_time
                        )
            except StopAsyncIteration:
                break
            except Exception as e:
                if not isinstance(e, httpx.ReadTimeout | TimeoutError):
                    raise
                current_time = time.perf_counter()
                errmsg = "stream chunk timeout, response_id: {}, response_time_cost:{}, stream_time_cost: {}, current_chunk_wait_time_cost:{}, stream_chunks: {}".format(
                    response.response_id,
                    response_time_cost,
                    current_time - response_time,
                    current_time - wait_start_time,
                    len(chunks),
                )
                raise StreamChunkTimeout(errmsg, e)
            chunk = cast(ModelResponseStream, chunk)
            chunks.append(chunk)
        stream_end_time = time.perf_counter()
        judy_hidden_params[STREAM_TIME_COST] = stream_end_time - response_time
        judy_hidden_params[STREAM_CHUNKS_LEN] = len(chunks)
        # chunks[-1] = response.chunks[-1]
        model_response = litellm.stream_chunk_builder(chunks, messages)
        assert isinstance(model_response, ModelResponse)
    else:
        model_response = cast(ModelResponse, response)

    if getattr(model_response, "_hidden_params", None) is None:
        model_response._hidden_params = {}
    model_response._hidden_params[JUDY_HIDDEN_PARAMS_KEY] = judy_hidden_params
    return model_response


async def _acompletion_with_retries(
    model: str,
    *,
    judy_option: JudyCompletionOption | None = None,
    judy_logger: AsyncLogger | None = None,
    **kwargs,
) -> ModelResponse:
    retry_count = 0
    if judy_option is None:
        judy_option = JudyCompletionOption()

    while True:
        start_time = time.perf_counter()
        try:
            model_response = await _acompletion(
                model,
                judy_option=judy_option,
                **kwargs,
            )
            return model_response
        except RateLimitError as e:
            ex = e
            errmsg = "rate limit"
        except LitellmTimeout as e:
            ex = e
            errmsg = "response timeout"
        except StreamChunkTimeout as e:
            ex = e.ex
            errmsg = str(e)
        except Exception as e:
            if not isinstance(
                e, ServiceUnavailableError | InternalServerError | httpx.ReadTimeout
            ):
                raise
            ex = e
            errmsg = "unknown exception"

        if judy_logger:
            time_cost = time.perf_counter() - start_time
            await judy_logger(
                "completion error, total_time_cost: {}, errmsg: {}, type:{}, ex: {}".format(
                    time_cost,
                    errmsg,
                    type(ex).__name__,
                    "\n".join(traceback.format_exception(ex)),
                )
            )

        if retry_count < judy_option.num_retries:
            retry_count += 1
        else:
            raise ex


async def acompletion(
    model: str,
    *,
    messages,
    tools: list[dict] | None = None,
    judy_option: JudyCompletionOption | None = None,
    judy_logger: AsyncLogger | None = None,
    blank_response_retry_message: str | None = None,
    **kwargs,
) -> ModelResponse | None:
    api_call_count = 0
    while True:
        try:
            adapted_tools = adapt_tools_for_model(model, tools)
            adapted_kwargs = adapt_completion_params(model, **kwargs)
            response = await _acompletion_with_retries(
                model,
                messages=_remove_warnings_from_litellm_completion(messages),
                tools=adapted_tools,
                judy_option=judy_option,
                judy_logger=judy_logger,
                stream=True,
                **adapted_kwargs,
            )
            response = cast(ModelResponse, response)
        except InternalServerError as e:
            if "Invalid response object" not in e.message:
                raise
            else:
                response = ModelResponse(choices=[])

        if len(response.choices) == 0:
            # gemini api有时返回内容为空（没有candidates)，出现这种情况后尝试重新生成
            # https://discuss.ai.google.dev/t/gemini-2-5-pro-with-empty-response-text/81175/356
            if judy_logger:
                await judy_logger(f"response has no choice: {response}")
            api_call_count += 1
            if api_call_count < 3:
                continue
            elif api_call_count == 3:
                if blank_response_retry_message is None:
                    blank_response_retry_message = BLANK_RESPONSE_RETRY_DEFAULT_MESSAGE
                messages.append(
                    {"role": "user", "content": blank_response_retry_message}
                )
            elif api_call_count < 5:
                continue
            else:
                return None
                # raise ModelBlankResponseError(
                #     f"No choices returned from LLM: {response}"
                # )
        else:
            break

    return response


async def acompletion_choice(**kwargs) -> Choices | None:
    response = await acompletion(**kwargs)
    if response is None:
        return None
    choice = cast(Choices, response.choices[0])
    return choice
