from __future__ import annotations

from typing import AsyncIterator, Iterator, List, Optional, Union, Dict, Any
import uuid
import json
from loguru import logger
import time
from turbo_agent_core.schema.agents import Tool,LLMModel
from turbo_agent_core.schema.events import (
    BaseEvent,
    RunLifecycleCreatedEvent,
    RunLifecycleCreatedPayload,
    RunLifecycleCompletedEvent,
    RunLifecycleCompletedPayload,
    RunLifecycleFailedEvent,
    RunLifecycleFailedPayload,
    ContentTextStartEvent,
    ContentTextStartPayload,
    ContentTextDeltaEvent,
    ContentTextDeltaPayload,
    ContentTextEndEvent,
    ContentTextEndPayload,
    ContentReasoningStartEvent,
    ContentReasoningStartPayload,
    ContentReasoningDeltaEvent,
    ContentReasoningDeltaPayload,
    ContentReasoningEndEvent,
    ContentReasoningEndPayload,
    ContentActionStartEvent,
    ContentActionStartPayload,
    ContentActionDeltaEvent,
    ContentActionDeltaPayload,
    ContentActionEndEvent,
    ContentActionEndPayload,
    ExecutorMetadata,
    UserInfo,
)
from turbo_agent_core.schema.states import Conversation, Message, MessageRole
from turbo_agent_core.utils.json_stream import JsonStreamParser
from turbo_agent_runtime.utils.message_converter import (
    convert_state_message_to_openai_message,
    convert_openai_message_to_state_message
)
from turbo_agent_runtime.utils.executor_identity import build_executor_id
from turbo_agent_runtime.utils.tool_convert import convert_tools_to_openai_functions

import litellm

LLMInput = Union[Conversation, List[Message]]
class LLMModelRuntime(LLMModel):
    _litellm_debug_enabled = False

    @classmethod
    def _ensure_litellm_debug_enabled(cls) -> None:
        if cls._litellm_debug_enabled:
            return
        try:
            turn_on_debug = getattr(litellm, "_turn_on_debug", None)
            if callable(turn_on_debug):
                turn_on_debug()
            if hasattr(litellm, "set_verbose"):
                litellm.set_verbose = True
            cls._litellm_debug_enabled = True
            logger.warning("LiteLLM Debug 已开启，后续将输出底层调试日志")
        except Exception as e:
            logger.warning("开启 LiteLLM Debug 失败，将继续按普通模式执行: {}", e)

    def _resolve_user_metadata(self, **kwargs: Any) -> UserInfo:
        raw = kwargs.get("user_metadata") or kwargs.get("user_info")
        if isinstance(raw, UserInfo):
            return raw
        if isinstance(raw, dict):
            try:
                return UserInfo.model_validate(raw)
            except Exception:
                pass

        user_id = kwargs.get("user_id")
        username = kwargs.get("username")
        return UserInfo(id=str(user_id or "local"), username=str(username or "local"))

    def _dump_for_event(self, obj: Any) -> Any:
        if obj is None:
            return None
        if hasattr(obj, "model_dump"):
            try:
                return obj.model_dump()
            except Exception:
                return obj
        if isinstance(obj, list):
            return [self._dump_for_event(x) for x in obj]
        if isinstance(obj, dict):
            return {k: self._dump_for_event(v) for k, v in obj.items()}
        return obj

    def _convert_tools(self, tools: List[Tool]) -> List[Dict[str, Any]]:
        return convert_tools_to_openai_functions(tools)

    def _parse_input(self, input: LLMInput, leaf_message_id: Optional[str] = None) -> List[Dict[str, Any]]:
        core_messages: List[Message] = []
        if isinstance(input, Conversation):
            core_messages = input.get_history(leaf_message_id)
        elif isinstance(input, list):
            core_messages = input
        else:
            raise ValueError(f"Invalid input type: {type(input)}")
            
        openai_msgs = []
        for msg in core_messages:
            converted = convert_state_message_to_openai_message(msg)
            if converted:
                openai_msgs.extend(converted)
        return openai_msgs

    def _get_model_config(self) -> Dict[str, Any]:
        if not self.instances:
            return {}
        instance = self.instances[0]
        endpoint = instance.endpoint
        
        base_url = getattr(endpoint, "baseURL", None) or getattr(endpoint, "url", None)
        api_key = getattr(endpoint, "accessKey", None) or getattr(endpoint, "accessKeyOrToken", None) or "dummy"
        
        params = self.defaultParameters

        model_id = instance.request_model_id or self.id
        instance_request_params = getattr(instance, "request_params", None) or {}
        # is_managed_litellm = bool(instance_request_params.get("turbo_managed_litellm"))
        # if not is_managed_litellm:
        #     # 非托管模型沿用 provider/model_id 的既有规则。
        if hasattr(instance, "provider") and instance.provider:
            provider_val = instance.provider.value if hasattr(instance.provider, "value") else str(instance.provider)
            model_id = f"{provider_val}/{model_id}"
        else:
            model_id = f"openai/{model_id}"

        model_params = {
            "model": model_id,
            "temperature": params.temperature,
            "max_tokens": params.max_tokens,
            "reasoning_effort": "low", 
            "drop_params": True
        }
        
        if params.top_p is not None:
            model_params["top_p"] = params.top_p
        if params.presence_penalty is not None:
            model_params["presence_penalty"] = params.presence_penalty
        if params.frequency_penalty is not None:
            model_params["frequency_penalty"] = params.frequency_penalty
        if params.logit_bias is not None:
            model_params["logit_bias"] = params.logit_bias
        if params.stop is not None:
            model_params["stop"] = params.stop
        if params.user is not None:
            model_params["user"] = params.user
        if params.extra:
            model_params.update(params.extra)
            
        return {
            "base_url": base_url,
            "api_key": api_key,
            "model_params": model_params
        }

    def _populate_final_usage(self, combined_response: Any, final_usage: Dict[str, Any]) -> None:
        cost = None
        logger.debug("已禁用 LiteLLM 费用计算，跳过 completion_cost 以避免额外远程请求。")

        usage = getattr(combined_response, "usage", None)
        if not usage:
            return

        final_usage["prompt_tokens"] = getattr(usage, "prompt_tokens", None)
        final_usage["completion_tokens"] = getattr(usage, "completion_tokens", None)
        final_usage["total_tokens"] = getattr(usage, "total_tokens", None)
        if cost is not None:
            final_usage["cost"] = cost

        details = getattr(usage, "completion_tokens_details", None)
        if not details:
            return

        if hasattr(details, "reasoning_tokens"):
            final_usage["reasoning_tokens"] = details.reasoning_tokens
        elif isinstance(details, dict):
            final_usage["reasoning_tokens"] = details.get("reasoning_tokens")

    def _format_output(self, message: Any) -> Message:
        msg = convert_openai_message_to_state_message(message)
        if msg:
            if not msg.id:
                msg.id = str(uuid.uuid4())
            return msg
        return Message(
            id=str(uuid.uuid4()),
            role=MessageRole.assistant,
            content=""
        )

    def _stub_response(self, messages: List[Dict[str, Any]]) -> str:
        return "I am a stub response because no model instance is configured."

    def run(self, conversation: LLMInput, tools: Optional[List[Tool]] = None, leaf_message_id: Optional[str] = None, **kwargs) -> Message:
        """
        同步执行 LLM 模型调用。
        """
        messages = self._parse_input(conversation, leaf_message_id)
        config = self._get_model_config()
        
        if config:
            model_params = config["model_params"]
            
            if tools:
                openai_tools = self._convert_tools(tools)
                if openai_tools:
                    model_params["tools"] = openai_tools
            try:
                # litellm handles api_key and base_url. model name is in model_params['model']
                resp = litellm.completion(
                    messages=messages, 
                    base_url=config.get("base_url"), 
                    api_key=config.get("api_key"),
                    **model_params
                )
                return self._format_output(resp.choices[0].message)
            except Exception as e:
                logger.error(f"LLM invoke failed: {e}")
                raise e
        
        return Message(
            id=str(uuid.uuid4()),
            role=MessageRole.assistant,
            content=self._stub_response(messages)
        )

    def _update_model_cost_map(self, base_url: str, api_key: str, model_id: str):
        """
        尝试从远端服务获取模型价格信息并注册到 litellm.model_cost
        用于解决 Model not mapped 警告。
        
        策略：
        1. 假设目标服务是 LiteLLM Server。
        2.请求 /v1/model/info (不带参数)。
        3. 遍历返回列表，匹配 model_name。
        """
        try:
            import requests
            
            # 构造请求 URL
            # 假设 base_url 类似于 http://host:port/v1 或 http://host:port
            info_url = base_url.rstrip("/")
            if info_url.endswith("/v1"):
                info_url = f"{info_url}/model/info"
            else:
                info_url = f"{info_url}/v1/model/info"
            
            # 预处理：移除 provider 前缀以便匹配远程列表
            # model_id 此时是 "provider/actual_model_id" 格式，例如 "openai/Qwen/Qwen2.5-7B"
            # 对于 litellm 来说，前缀是必须带的，所以可以直接 split 处理，取后半部分作为远程匹配的名称
            search_model_name = model_id
            if "/" in model_id:
                _, search_model_name = model_id.split("/", 1)

            logger.debug(f"[费用同步] 正在尝试从 {info_url} 获取模型 {search_model_name} (原始ID: {model_id}) 的价格信息...")
            
            # 发起请求 (不带 query params)
            resp = requests.get(info_url, headers={"Authorization": f"Bearer {api_key}"}, timeout=2)
            
            if resp.status_code == 200:
                data = resp.json()
                
                if "data" in data and isinstance(data["data"], list):
                    found = False
                    for item in data["data"]:
                        remote_name = item.get("model_name")

                        # 匹配 model_name (使用去除 provider 后的名称)
                        if remote_name == search_model_name:
                            info = item.get("model_info", {})
                            logger.debug(f"[费用同步] 模型信息接口返回数据: {info}")
                            # 提取价格信息
                            input_cost = info.get("input_cost_per_token", 0)
                            output_cost = info.get("output_cost_per_token", 0)
                            max_tokens = info.get("max_tokens", 16000)
                            
                            # 注册到 litellm 本地映射
                            # 注意：我们注册的 key 必须是当前使用的 model_id (带 provider 的完整串)，
                            # 这样 litellm 在计算本次调用的 cost 时才能查找到。
                            litellm.register_model({
                                search_model_name: {
                                    "input_cost_per_token": input_cost,
                                    "output_cost_per_token": output_cost,
                                    "max_tokens": max_tokens,
                                    "litellm_provider": "openai" # 标记为 OpenAI 兼容协议
                                }
                            })
                            logger.debug(f"[费用同步] 成功注册模型 {search_model_name} 的价格 (匹配到远程模型: {remote_name}): 输入={input_cost}, 输出={output_cost}")
                            found = True
                            break
                    
                    if not found:
                         logger.warning(f"[费用同步] 接口返回了模型列表，但未找到匹配的模型: {search_model_name}")
                else:
                    logger.warning(f"[费用同步] 接口返回格式异常: data 字段不存在或不是列表")
            else:
                logger.warning(f"[费用同步] 请求失败，Web状态码: {resp.status_code}")

        except Exception as e:
            # 静默失败，不影响主流程
            logger.warning(f"[费用同步] 获取远程模型价格信息时发生异常: {e}")

    def stream(self, conversation: LLMInput, tools: Optional[List[Tool]] = None, leaf_message_id: Optional[str] = None, **kwargs) -> Iterator[BaseEvent]:
        # ... logic ...
        start_time = time.time()
        trace_id = str(uuid.uuid4())
        run_id = f"run_{self.id}_{time.time()*1000}"
        executor_type = self.run_type
        executor_id = build_executor_id(self, executor_type)
        executor_path = [executor_id]
        yield RunLifecycleCreatedEvent(
            trace_id=trace_id,
            trace_path=[],
            run_id=run_id,
            run_path=[run_id],
            executor_type=executor_type,
            executor_id=executor_id,
            executor_path=executor_path,
            executor_metadata=ExecutorMetadata(
                id=self.id,
                name=self.name,
                belong_to_project_path=self.belong_to_project_path,
                name_id=self.name_id,
                description=self.description,
                avatar_uri=self.avatar_uri,
                run_type=self.run_type,
                version_id=None,
                version=None
            ),
            user_metadata=self._resolve_user_metadata(**kwargs),
            payload=RunLifecycleCreatedPayload(
                input_data={
                    "conversation": self._dump_for_event(conversation),
                    "leaf_message_id": leaf_message_id,
                    "tools": self._dump_for_event(tools),
                }
            ),
        )
        
        try:
            messages = self._parse_input(conversation, leaf_message_id)
            config = self._get_model_config()
            self._ensure_litellm_debug_enabled()
            
            if config:
                # 尝试注册模型价格 (One-off)
                if config.get("base_url"):
                     self._update_model_cost_map(config.get("base_url"), config.get("api_key"), config["model_params"].get("model"))

                model_params = config["model_params"]
                if tools:
                    openai_tools = self._convert_tools(tools)
                    if openai_tools:
                        model_params["tools"] = openai_tools
                
                model_params["stream"] = True
                # 开启流式 Usage 返回 (支持 OpenAI/LitellmServer 等标准)
                model_params["stream_options"] = {"include_usage": True}
                
                yield ContentTextStartEvent(
                    trace_id=trace_id, run_id=run_id,
                    executor_type=executor_type,
                    executor_id=executor_id, executor_path=executor_path,
                    payload=ContentTextStartPayload(format="plain")
                )
                
                accumulated_content = []
                accumulated_reasoning = []
                tool_call_states = {}
                all_chunks = []
                # 初始化 Usage，确保字段存在
                final_usage = {
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0
                }

                # Use litellm.completion with stream=True
                response = litellm.completion(
                    messages=messages, 
                    api_key=config.get("api_key"),
                    base_url=config.get("base_url"),
                    **model_params
                )

                for chunk in response:
                    all_chunks.append(chunk)
                    # 尝试捕获 Usage 信息 (通常在最后一个 Chunk)
                    logger.debug(f"Received chunk: {chunk}")
                    if hasattr(chunk, "usage") and chunk.usage:
                        # 严格映射标准字段
                        p_t = getattr(chunk.usage, "prompt_tokens", 0)
                        c_t = getattr(chunk.usage, "completion_tokens", 0)
                        t_t = getattr(chunk.usage, "total_tokens", 0)
                        if t_t == 0 and p_t > 0:
                            t_t = p_t + c_t
                        final_usage = {
                           "prompt_tokens": p_t,
                           "completion_tokens": c_t,
                           "total_tokens": t_t
                        }

                    if not chunk.choices or len(chunk.choices) == 0:
                        continue
                        
                    delta = chunk.choices[0].delta
                    
                    # Reasoning
                    if hasattr(delta, "reasoning_content") and delta.reasoning_content:
                        if not accumulated_reasoning:
                            yield ContentReasoningStartEvent(
                                trace_id=trace_id, run_id=run_id,
                                executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                                payload=ContentReasoningStartPayload(format="plain")
                            )
                        accumulated_reasoning.append(delta.reasoning_content)
                        yield ContentReasoningDeltaEvent(
                            trace_id=trace_id, run_id=run_id,
                            executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                            payload=ContentReasoningDeltaPayload(delta=delta.reasoning_content)
                        )

                    # Content
                    part = delta.content
                    if part:
                        accumulated_content.append(part)
                        yield ContentTextDeltaEvent(
                            trace_id=trace_id, run_id=run_id,
                            executor_type=executor_type,
                            executor_id=executor_id, executor_path=executor_path,
                            payload=ContentTextDeltaPayload(delta=str(part))
                        )
                    
                    # Tool Calls
                    if delta.tool_calls:
                        for tc in delta.tool_calls:
                            idx = tc.index
                            if idx not in tool_call_states:
                                tool_call_states[idx] = {
                                    "id": "", 
                                    "name": "", 
                                    "buffer": "", 
                                    "parser": JsonStreamParser(), 
                                    "has_started": False
                                }
                            
                            state = tool_call_states[idx]
                            
                            if tc.id:
                                state["id"] = tc.id
                            if tc.function and tc.function.name:
                                state["name"] += tc.function.name
                            
                            if not state["has_started"] and state["id"] and state["name"]:
                                state["has_started"] = True
                                yield ContentActionStartEvent(
                                    trace_id=trace_id, run_id=run_id,
                                    executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                                    action_id=state["id"],
                                    payload=ContentActionStartPayload(
                                        name=state["name"], 
                                        call_type="TOOL",
                                    )
                                )

                            if tc.function and tc.function.arguments:
                                fragment = tc.function.arguments
                                state["buffer"] += fragment
                                
                                if state["has_started"]:
                                    deltas = state["parser"].feed(fragment)
                                    for path, char in deltas:
                                        yield ContentActionDeltaEvent(
                                            trace_id=trace_id, run_id=run_id,
                                            executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                                            action_id=state["id"],
                                            payload=ContentActionDeltaPayload(
                                                part="args",
                                                delta=char,
                                                key_path=list(path)
                                            )
                                        )

                # 使用 stream_chunk_builder 重组完整响应
                combined_response = litellm.stream_chunk_builder(all_chunks, messages=messages)
                combined_msg = combined_response.choices[0].message

                self._populate_final_usage(combined_response, final_usage)

                full_text = combined_msg.content or ""
                yield ContentTextEndEvent(
                    trace_id=trace_id, run_id=run_id,
                    executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                    payload=ContentTextEndPayload(full_text=str(full_text))
                )
                
                # 尝试提取 Reasoning (如果 builder 支持)
                full_reasoning = getattr(combined_msg, "reasoning_content", "") or "".join(accumulated_reasoning)
                if full_reasoning:
                    yield ContentReasoningEndEvent(
                        trace_id=trace_id, run_id=run_id,
                        executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                        payload=ContentReasoningEndPayload(full_text=full_reasoning)
                    )
                
                # 处理 Tool Calls
                tool_calls = combined_msg.tool_calls or []
                final_tool_calls_data = [] 

                for tc in tool_calls:
                    # 尝试解析参数
                    args_dict =tc.function.arguments
                    if tc.function and tc.function.arguments:
                        try:
                            args_dict = json.loads(tc.function.arguments)
                        except:
                            pass
                        
                    yield ContentActionEndEvent(
                        trace_id=trace_id, run_id=run_id,
                        executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                        action_id=tc.id,
                        payload=ContentActionEndPayload(arguments=args_dict)
                    )
                    # 兼容性处理 so that _format_output or downstream works
                    if hasattr(tc, "model_dump"):
                        final_tool_calls_data.append(tc.model_dump())
                    else:
                        final_tool_calls_data.append(tc)

                # 构造最终输出
                # 注意：litellm 1.x stream_chunk_builder 返回的 message 可能直接可用
                # 这里为了保险，手动构造一个 dict 传给 _format_output，就像之前做的那样
                # 或者直接利用 combined_msg 对象 (如果 _format_output 支持)
                # 之前代码: output = self._format_output(final_msg_dict)
                # final_msg_dict 是 dict。convert_openai_message_to_state_message 支持 dict 或 object
                output = self._format_output(combined_msg)
                
                # Standardize usage
                standard_usage = {
                    "promptToken": final_usage.get("prompt_tokens"),
                    "completionToken": final_usage.get("completion_tokens"),
                    "reasoningToken": final_usage.get("reasoning_tokens"),
                    "tokenCost": final_usage.get("total_tokens"),
                    "moneyCost": final_usage.get("cost"),
                    "timeCost": time.time() - start_time
                }

                yield RunLifecycleCompletedEvent(
                    trace_id=trace_id,
                    trace_path=[],
                    run_id=run_id,
                    run_path=[run_id],
                    executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                    payload=RunLifecycleCompletedPayload(output=output, usage=standard_usage),
                )

            else:
                stub = self._stub_response(messages)
                yield ContentTextStartEvent(
                    trace_id=trace_id, run_id=run_id,
                    executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                    payload=ContentTextStartPayload(format="plain")
                )
                yield ContentTextDeltaEvent(
                    trace_id=trace_id, run_id=run_id,
                    executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                    payload=ContentTextDeltaPayload(delta=stub)
                )
                yield ContentTextEndEvent(
                    trace_id=trace_id, run_id=run_id,
                    executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                    payload=ContentTextEndPayload(full_text=stub)
                )
                yield RunLifecycleCompletedEvent(
                    trace_id=trace_id,
                    trace_path=[],
                    run_id=run_id,
                    run_path=[run_id],
                    executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                    payload=RunLifecycleCompletedPayload(output=None, usage={
                        "promptToken": 0,
                        "completionToken": 0,
                        "tokenCost": 0,
                        "moneyCost": 0.0,
                        "timeCost": time.time() - start_time
                    }),
                )
                
        except Exception as e:
            logger.exception(e)
            yield RunLifecycleFailedEvent(
                trace_id=trace_id,
                trace_path=[],
                run_id=run_id,
                run_path=[run_id],
                executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                payload=RunLifecycleFailedPayload(error={"code": "RuntimeError", "message": str(e)}),
            )

    async def a_stream(self, conversation: LLMInput, tools: Optional[List[Tool]] = None, leaf_message_id: Optional[str] = None, **kwargs) -> AsyncIterator[BaseEvent]:
        start_time = time.time()
        trace_id = str(uuid.uuid4())
        run_id = f"run_{self.id}_{time.time()*1000}"
        executor_type = self.run_type
        executor_id = build_executor_id(self, executor_type)
        executor_path = [executor_id]
        yield RunLifecycleCreatedEvent(
            trace_id=trace_id,
            trace_path=[],
            run_id=run_id,
            run_path=[run_id],
            executor_type=executor_type,
            executor_id=executor_id,
            executor_path=executor_path,
            executor_metadata=ExecutorMetadata(
                id=self.id,
                name=self.name,
                belong_to_project_path=self.belong_to_project_path,
                name_id=self.name_id,
                description=self.description,
                avatar_uri=self.avatar_uri,
                run_type=self.run_type,
                version_id=None,
                version=None
            ),
            user_metadata=self._resolve_user_metadata(**kwargs),
            payload=RunLifecycleCreatedPayload(
                input_data={
                    "conversation": self._dump_for_event(conversation),
                    "leaf_message_id": leaf_message_id,
                    "tools": self._dump_for_event(tools),
                }
            ),
        )
        
        try:
            messages = self._parse_input(conversation, leaf_message_id)
            logger.debug(f"Parsed messages for a_stream: {messages}")
            config = self._get_model_config()
            
            if config:
                # 尝试注册模型价格 (One-off)
                if config.get("base_url"):
                     self._update_model_cost_map(config.get("base_url"), config.get("api_key"), config["model_params"].get("model"))

                model_params = config["model_params"]
                if tools:
                    openai_tools = self._convert_tools(tools)
                    if openai_tools:
                        model_params["tools"] = openai_tools
                
                model_params["stream"] = True
                # 开启流式 Usage 返回
                model_params["stream_options"] = {"include_usage": True}
                
                yield ContentTextStartEvent(
                    trace_id=trace_id, run_id=run_id,
                    executor_type=executor_type,
                    executor_id=executor_id, executor_path=executor_path,
                    payload=ContentTextStartPayload(format="plain")
                )
                
                accumulated_content = []
                accumulated_reasoning = []
                tool_call_states = {}
                all_chunks = []
                # 初始化 Usage，确保字段存在
                final_usage = {
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0
                }
                logger.debug(f"Starting async LLM completion stream...")
                logger.debug(
                    "LiteLLM 调试请求参数: model={}, base_url={}, has_api_key={}, tools_count={}, stream_options={}",
                    model_params.get("model"),
                    config.get("base_url"),
                    bool(config.get("api_key")),
                    len(model_params.get("tools") or []),
                    model_params.get("stream_options"),
                )
                response = await litellm.acompletion(
                    messages=messages, 
                    api_key=config.get("api_key"),
                    base_url=config.get("base_url"),
                    **model_params
                )
                logger.debug(f"Async LLM completion stream started. {response}")

                async for chunk in response:
                    all_chunks.append(chunk)
                    logger.debug(f"LLM Received chunk: {chunk}")
                    # 尝试捕获 Usage 信息 (通常在最后一个 Chunk)
                    if hasattr(chunk, "usage") and chunk.usage:
                        # 严格映射标准字段
                        p_t = getattr(chunk.usage, "prompt_tokens", 0)
                        c_t = getattr(chunk.usage, "completion_tokens", 0)
                        t_t = getattr(chunk.usage, "total_tokens", 0)
                        if t_t == 0 and p_t > 0:
                            t_t = p_t + c_t
                        final_usage = {
                           "prompt_tokens": p_t,
                           "completion_tokens": c_t,
                           "total_tokens": t_t
                        }

                    if not chunk.choices or len(chunk.choices) == 0:
                        continue

                    delta = chunk.choices[0].delta
                    
                    # Reasoning
                    if hasattr(delta, "reasoning_content") and delta.reasoning_content:
                        if not accumulated_reasoning:
                            yield ContentReasoningStartEvent(
                                trace_id=trace_id, run_id=run_id,
                                executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                                payload=ContentReasoningStartPayload(format="plain")
                            )
                        accumulated_reasoning.append(delta.reasoning_content)
                        yield ContentReasoningDeltaEvent(
                            trace_id=trace_id, run_id=run_id,
                            executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                            payload=ContentReasoningDeltaPayload(delta=delta.reasoning_content)
                        )

                    # Content
                    part = delta.content
                    if part:
                        accumulated_content.append(part)
                        yield ContentTextDeltaEvent(
                            trace_id=trace_id, run_id=run_id,
                            executor_type=executor_type,
                            executor_id=executor_id, executor_path=executor_path,
                            payload=ContentTextDeltaPayload(delta=str(part))
                        )
                    
                    # Tool Calls
                    if delta.tool_calls:
                        for tc in delta.tool_calls:
                            idx = tc.index
                            if idx not in tool_call_states:
                                tool_call_states[idx] = {
                                    "id": "", 
                                    "name": "", 
                                    "buffer": "", 
                                    "parser": JsonStreamParser(), 
                                    "has_started": False
                                }
                            
                            state = tool_call_states[idx]
                            
                            if tc.id:
                                state["id"] = tc.id
                            if tc.function and tc.function.name:
                                state["name"] += tc.function.name
                            
                            if not state["has_started"] and state["id"] and state["name"]:
                                state["has_started"] = True
                                yield ContentActionStartEvent(
                                    trace_id=trace_id, run_id=run_id,
                                    executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                                    action_id=state["id"],
                                    payload=ContentActionStartPayload(
                                        name=state["name"], 
                                        call_type="TOOL",
                                    )
                                )

                            if tc.function and tc.function.arguments:
                                fragment = tc.function.arguments
                                state["buffer"] += fragment
                                
                                if state["has_started"]:
                                    deltas = state["parser"].feed(fragment)
                                    for path, char in deltas:
                                        yield ContentActionDeltaEvent(
                                            trace_id=trace_id, run_id=run_id,
                                            executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                                            action_id=state["id"],
                                            payload=ContentActionDeltaPayload(
                                                part="args",
                                                delta=char,
                                                key_path=list(path)
                                            )
                                        )

                # 使用 stream_chunk_builder 重组完整响应
                combined_response = litellm.stream_chunk_builder(all_chunks, messages=messages)
                logger.debug(f"Combined response stream_chunk_builder: {combined_response}")
                combined_msg = combined_response.choices[0].message

                self._populate_final_usage(combined_response, final_usage)

                full_text = combined_msg.content or ""
                yield ContentTextEndEvent(
                    trace_id=trace_id, run_id=run_id,
                    executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                    payload=ContentTextEndPayload(full_text=str(full_text))
                )
                
                # 尝试提取 Reasoning (如果 builder 支持)
                full_reasoning = getattr(combined_msg, "reasoning_content", "") or "".join(accumulated_reasoning)
                if full_reasoning:
                    yield ContentReasoningEndEvent(
                        trace_id=trace_id, run_id=run_id,
                        executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                        payload=ContentReasoningEndPayload(full_text=full_reasoning)
                    )
                
                # 处理 Tool Calls
                tool_calls = combined_msg.tool_calls or []
                final_tool_calls_data = [] 

                for tc in tool_calls:
                    # 尝试解析参数
                    args_dict = {}
                    if tc.function and tc.function.arguments:
                        try:
                            args_dict = json.loads(tc.function.arguments)
                        except:
                            pass
                        
                    yield ContentActionEndEvent(
                        trace_id=trace_id, run_id=run_id,
                        executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                        action_id=tc.id,
                        payload=ContentActionEndPayload(arguments=args_dict)
                    )
                    # 兼容性处理 so that _format_output or downstream works
                    if hasattr(tc, "model_dump"):
                        final_tool_calls_data.append(tc.model_dump())
                    else:
                        final_tool_calls_data.append(tc)

                output = self._format_output(combined_msg)
                
                # Standardize usage
                standard_usage = {
                    "promptToken": final_usage.get("prompt_tokens"),
                    "completionToken": final_usage.get("completion_tokens"),
                    "reasoningToken": final_usage.get("reasoning_tokens"),
                    "tokenCost": final_usage.get("total_tokens"),
                    "moneyCost": final_usage.get("cost"),
                    "timeCost": time.time() - start_time
                }

                logger.debug(f"Calculated final_usage: {standard_usage}")
                yield RunLifecycleCompletedEvent(
                    trace_id=trace_id,
                    trace_path=[],
                    run_id=run_id,
                    run_path=[run_id],
                    executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                    payload=RunLifecycleCompletedPayload(output=output, usage=standard_usage),
                )
                

            else:
                stub = self._stub_response(messages)
                yield ContentTextStartEvent(
                    trace_id=trace_id, run_id=run_id,
                    executor_type=executor_type,
                    executor_id=executor_id, executor_path=executor_path,
                    payload=ContentTextStartPayload(format="plain")
                )
                yield ContentTextDeltaEvent(
                    trace_id=trace_id, run_id=run_id,
                    executor_type=executor_type,
                    executor_id=executor_id, executor_path=executor_path,
                    payload=ContentTextDeltaPayload(delta=stub)
                )
                yield ContentTextEndEvent(
                    trace_id=trace_id, run_id=run_id,
                    executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                    payload=ContentTextEndPayload(full_text=stub)
                )
                yield RunLifecycleCompletedEvent(
                    trace_id=trace_id,
                    trace_path=[],
                    run_id=run_id,
                    run_path=[run_id],
                    executor_type=executor_type,
                    executor_id=executor_id,
                    executor_path=executor_path,
                    payload=RunLifecycleCompletedPayload(output=None, usage={"total_tokens": 0}),
                )
                
        except Exception as e:
            logger.exception(e)
            yield RunLifecycleFailedEvent(
                trace_id=trace_id,
                trace_path=[],
                run_id=run_id,
                run_path=[run_id],
                executor_type=executor_type,
                executor_id=executor_id,
                executor_path=executor_path,
                payload=RunLifecycleFailedPayload(error={"code": "RuntimeError", "message": str(e)}),
            )
        logger.debug(f"Async a_stream completed for trace_id: {trace_id}")

    async def a_run(self, conversation: LLMInput, tools: Optional[List[Tool]] = None, leaf_message_id: Optional[str] = None, **kwargs) -> Message:
        """
        异步执行 LLM 模型调用。
        """
        messages = self._parse_input(conversation, leaf_message_id)
        config = self._get_model_config()
        
        if config:
            model_params = config["model_params"]
            self._ensure_litellm_debug_enabled()
            
            if tools:
                openai_tools = self._convert_tools(tools)
                if openai_tools:
                    model_params["tools"] = openai_tools
            try:
                logger.warning(
                    "LiteLLM 调试请求参数(a_run): model={}, base_url={}, has_api_key={}, tools_count={}",
                    model_params.get("model"),
                    config.get("base_url"),
                    bool(config.get("api_key")),
                    len(model_params.get("tools") or []),
                )
                resp = await litellm.acompletion(
                    messages=messages, 
                    api_key=config.get("api_key"),
                    base_url=config.get("base_url"),
                    **model_params
                )
                return self._format_output(resp.choices[0].message)
            except Exception as e:
                logger.error(f"LLM invoke failed: {e}")
                raise e
        
        return Message(
            id=str(uuid.uuid4()),
            role=MessageRole.assistant,
            content=self._stub_response(messages)
        )


