import json
import traceback
from collections.abc import AsyncGenerator
from copy import deepcopy
from datetime import datetime
from typing import Annotated, Any, cast

import contexttimer
from judy_core.utils.llm import acompletion
from litellm import Choices
from litellm.types.llms.vertex_ai import VertexToolName
from litellm.types.utils import ModelResponse
from pydantic import BaseModel, Field

CALL_LLM_BATCH_TOOL_NAME = "call_llm_batch"

CallLLMId = str | None

SYSTEM_INSTRUCTIOINS_FIELD_DESCRIPTION = """\
不需要填写role字段，直接填写content内容即可。
样例1：
"..."
样例2：
[
    {"type":"text", "text":"..."}, 
    {"type":"image_url", "image_url":{"url":"data:image/jpeg;base64,..."}}
]
\
"""

MESSAGES_FIELD_DESCRIPTION = """\
与litellm的message格式相同（与openai格式类似），至少需要包含role,content字段。样例：
[
    {
        "role": "user",
        "content": [
            {"type":"text", "text":"..."}, 
            {"type":"image_url", "image_url":{"url":"data:image/jpeg;base64,..."}}
        ]
    },
    ...
]\
"""


class CallLLMInput(BaseModel):
    id: CallLLMId

    system_instructions: Annotated[
        str | list[dict], Field(description=SYSTEM_INSTRUCTIOINS_FIELD_DESCRIPTION)
    ]
    messages: Annotated[
        list[dict],
        Field(description=MESSAGES_FIELD_DESCRIPTION),
    ]

    response_json_schema: Annotated[
        dict[str, Any] | None, Field(description="设置structured output的json schema")
    ] = None
    tool_names: list[str] | None = None


class CallLLMBatchInput(BaseModel):
    inputs: list[CallLLMInput]


class CallLLMOutput(BaseModel):
    error: Annotated[str | None, Field(description="错误信息")]
    id: CallLLMId
    text: str


class CallLLMOutputWithDebugInfo(BaseModel):
    output: CallLLMOutput
    time_cost: float
    timestamp: str
    raw_errmsg: str | None = None
    response: ModelResponse | None = None


class CallLLMBatchOutput(BaseModel):
    outputs: list[CallLLMOutput]


async def call_llm_batch(input_json_file_path: str, output_json_file_path: str) -> str:
    """你可以使用此工具来处理复杂的文本任务（如分类、信息提取等）

    input_json_file_path: json schema如下：{input_json_schema}
    output_json_file_path: json schema 如下： {output_json_schema}
    """
    ...


if call_llm_batch.__doc__ is not None:
    call_llm_batch.__doc__ = call_llm_batch.__doc__.format(
        input_json_schema=json.dumps(CallLLMBatchInput.model_json_schema()),
        output_json_schema=json.dumps(CallLLMBatchOutput.model_json_schema()),
    )

CLAUDE_TOOL_MAX_USE = 5
CLAUDE_WEB_SEARCH_TOOL_NAME = "web_search"
CLAUDE_WEB_SEARCH_TOOL_TYPE = "web_search_20250305"
CLAUDE_WEB_FETCH_TOOL_NAME = "web_fetch"
CLAUDE_WEB_FETCH_TOOL_TYPE = "web_fetch_20250910"
CLAUDE_TOOL_DICT: dict[str, dict] = {
    CLAUDE_WEB_SEARCH_TOOL_NAME: {
        "type": CLAUDE_WEB_SEARCH_TOOL_TYPE,
    },
    CLAUDE_WEB_FETCH_TOOL_NAME: {
        "type": CLAUDE_WEB_FETCH_TOOL_TYPE,
    },
}

WEB_FETCH_TOOL_DESCRIPTION = "网页抓取，访问URL并提取网页或文件内容。对于高度动态和比较复杂的网页，抓取工具能够提供更准确的内容。"


class LLMBuiltinTool(BaseModel):
    name: str
    description: str


def get_llm_builtin_tools_info(model_name: str) -> list[LLMBuiltinTool]:
    if "gemini" in model_name:
        return [
            LLMBuiltinTool(
                name=VertexToolName.GOOGLE_SEARCH.value,
                description="google网页搜索，联网搜索最新信息。不支持指定搜索的时间范围和网页快照，只能搜索到网页的当前最新信息。如果查询的是过去某个时间点的历史信息，该工具给出的结果并不可靠。",
            ),
            LLMBuiltinTool(
                name=VertexToolName.URL_CONTEXT.value,
                description=WEB_FETCH_TOOL_DESCRIPTION,
            ),
        ]
    elif "claude" in model_name:
        return [
            LLMBuiltinTool(
                name=CLAUDE_WEB_SEARCH_TOOL_NAME,
                description="网页搜索，联网搜索最新信息",
            ),
            # LLMBuiltinTool(
            #     name=CLAUDE_WEB_FETCH_TOOL_NAME,
            #     description=WEB_FETCH_TOOL_DESCRIPTION,
            # ),
        ]

    else:
        return []


def get_llm_builin_tools(model_name: str, tool_names: list[str] | None) -> list | None:
    if not tool_names:
        return None
    tools = []
    if "gemini" in model_name:
        for tool_name in tool_names:
            tools.append({tool_name: {}})
    elif "claude" in model_name:
        for tool_name in tool_names:
            tool = CLAUDE_TOOL_DICT.get(tool_name)
            if tool is None:
                raise RuntimeError("unsupported tool: {}".format(tool_name))
            tool = deepcopy(tool)
            tool["name"] = tool_name
            tool["max_use"] = CLAUDE_TOOL_MAX_USE
            tools.append(tool)
    else:
        raise RuntimeError("unsupported tool: {}".format(tool_names))
    return tools


async def call_llm_batch_impl(
    model_name: str, batch_input: CallLLMBatchInput
) -> AsyncGenerator[CallLLMOutputWithDebugInfo]:
    for i, input in enumerate(batch_input.inputs):
        raw_errmsg = None
        with contexttimer.Timer() as t:
            try:
                response_format = None
                if input.response_json_schema is not None:
                    response_format = {
                        "type": "json_schema",
                        "json_schema": {
                            "name": "output_schema",
                            "schema": input.response_json_schema,
                            "strict": True,
                        },
                    }
                tools = get_llm_builin_tools(model_name, input.tool_names)

                system_message = {
                    "role": "system",
                    "content": input.system_instructions,
                }
                response = await acompletion(
                    model=model_name,
                    messages=[system_message] + input.messages,
                    response_format=response_format,
                    tools=tools,
                    timeout=90,
                )
                if response is not None:
                    if len(response.choices) > 0:
                        choice = cast(Choices, response.choices[0])
                        output = CallLLMOutput(
                            error=None,
                            id=input.id,
                            text=choice.message.content or "",
                        )

                    else:
                        output = CallLLMOutput(
                            error="model call return a response with no choices",
                            id=input.id,
                            text="",
                        )
                else:
                    output = CallLLMOutput(
                        error="model call return a null response",
                        id=input.id,
                        text="",
                    )
            except Exception as e:
                response = None
                raw_errmsg = "".join(traceback.format_exception(e, limit=3))
                truncated_tb_msg = "".join(raw_errmsg.splitlines()[-10:])
                output = CallLLMOutput(
                    error="model call failed, error：{}".format(truncated_tb_msg),
                    id=input.id,
                    text="",
                )
        yield CallLLMOutputWithDebugInfo(
            output=output,
            raw_errmsg=raw_errmsg,
            time_cost=t.elapsed,
            timestamp=str(datetime.now()),
            response=response,
        )


def summary_batch_output(
    batch_output: CallLLMBatchOutput, output_json_file_path: str
) -> str:
    count = len(batch_output.outputs)
    error_count = 0
    for output in batch_output.outputs:
        if output.error is not None:
            error_count += 1
    if error_count != 0:
        return "Some input process failed, error_count:{}, success_count:{}, process result(including failed) have been saved to {}.".format(
            error_count, count - error_count, output_json_file_path
        )
    return "All inputs were processed successfully."
