# coding=utf-8
"""
FrontendTool 执行器。

前端工具的执行逻辑：
1. 发起工具调用，生成 trace id
2. 将工具调用信息写入数据库（ToolCallRecord），等待前端处理
3. 异步轮询数据库，通过 trace id 查询前端返回结果
4. 超时后返回默认结果（前端状态未知，需向用户确认）

使用方式：
- 运行时创建 FrontendToolRuntime 实例
- 调用 stream/a_stream 启动流式执行
- 发送 FrontendToolPendingEvent 通知前端有新工具待执行
- 轮询等待前端通过 API 回写结果到 ToolCallRecord
"""

from __future__ import annotations

import asyncio
import os
import time
import uuid
from typing import Any, AsyncIterator, Callable, Dict, Iterator, Optional

from loguru import logger

from turbo_agent_core.schema.enums import JSON
from turbo_agent_core.schema.events import (
    BaseEvent,
    RunLifecycleCreatedEvent,
    RunLifecycleCreatedPayload,
    RunLifecycleCompletedEvent,
    RunLifecycleCompletedPayload,
    RunLifecycleFailedEvent,
    RunLifecycleFailedPayload,
    ContentAgentResultStartEvent,
    ContentAgentResultStartPayload,
    ContentAgentResultEndEvent,
    ContentAgentResultEndPayload,
    ExecutorMetadata,
    UserInfo,
)
from turbo_agent_core.schema.agents import FrontendTool
from turbo_agent_core.schema.refs import ToolCallRecordRef
from turbo_agent_runtime.utils.executor_identity import build_executor_id


# 前端工具配置（从环境变量读取默认值）
FRONTEND_TOOL_POLL_INTERVAL_MS = int(os.getenv("FRONTEND_TOOL_POLL_INTERVAL_MS", "500"))
FRONTEND_TOOL_TIMEOUT_MS = int(os.getenv("FRONTEND_TOOL_TIMEOUT_MS", "30000"))


class FrontendToolRuntime(FrontendTool):
    """前端工具运行时：在浏览器/客户端环境中执行的工具。
    
    执行流程：
    1. 创建 ToolCallRecord，状态为 pending
    2. 发送事件通知前端有待执行工具
    3. 轮询数据库等待前端回写结果
    4. 超时或获取结果后返回
    
    配置（环境变量）：
    - FRONTEND_TOOL_POLL_INTERVAL_MS: 轮询间隔（毫秒），默认 500ms
    - FRONTEND_TOOL_TIMEOUT_MS: 超时时间（毫秒），默认 30000ms (30s)
    """

    def _resolve_user_metadata(self, **kwargs) -> 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 _get_poll_config(self, **kwargs) -> tuple[int, int]:
        """获取轮询配置。
        
        Returns:
            (poll_interval_ms, timeout_ms)
        """
        poll_interval = kwargs.get("poll_interval_ms", FRONTEND_TOOL_POLL_INTERVAL_MS)
        timeout = kwargs.get("timeout_ms", self.timeout_ms or FRONTEND_TOOL_TIMEOUT_MS)
        return poll_interval, timeout

    async def _wait_for_frontend_result(
        self,
        trace_id: str,
        record_id: str,
        poll_interval_ms: int,
        timeout_ms: int,
        record_fetcher: Optional[Callable[[str], Optional[ToolCallRecordRef]]] = None,
    ) -> Dict[str, Any]:
        """等待前端返回结果。
        
        Args:
            trace_id: 执行跟踪 ID
            record_id: ToolCallRecord ID
            poll_interval_ms: 轮询间隔（毫秒）
            timeout_ms: 超时时间（毫秒）
            record_fetcher: 记录查询函数（用于从数据库获取最新状态）
            
        Returns:
            结果字典，包含 status 和 result
        """
        start_time = time.time()
        poll_interval_sec = poll_interval_ms / 1000.0
        timeout_sec = timeout_ms / 1000.0
        
        while True:
            # 检查是否超时
            elapsed = time.time() - start_time
            if elapsed > timeout_sec:
                logger.warning(
                    f"FrontendTool timeout: trace_id={trace_id}, "
                    f"record_id={record_id}, elapsed={elapsed:.2f}s"
                )
                return {
                    "status": "timeout",
                    "result": {
                        "error": "前端工具执行超时",
                        "message": "前端未在指定时间内返回结果，请向用户确认执行状态",
                        "trace_id": trace_id,
                        "record_id": record_id,
                    },
                }
            
            # 如果提供了查询函数，查询最新状态
            if record_fetcher:
                try:
                    record = record_fetcher(record_id)
                    if record:
                        # 检查记录状态
                        record_status = getattr(record, "status", None)
                        if record_status == "completed":
                            return {
                                "status": "success",
                                "result": getattr(record, "result", {}),
                            }
                        elif record_status == "failed":
                            return {
                                "status": "error",
                                "result": {
                                    "error": "前端工具执行失败",
                                    "message": getattr(record, "error_message", "未知错误"),
                                },
                            }
                except Exception as e:
                    logger.warning(f"Failed to fetch ToolCallRecord: {e}")
            
            # 等待下一轮轮询
            await asyncio.sleep(poll_interval_sec)

    def run(self, input: JSON, **kwargs) -> JSON:
        """同步执行（前端工具不支持真正的同步执行，返回错误）。
        
        前端工具必须在异步环境中执行，因为需要轮询等待前端返回。
        """
        return {
            "error": "FrontendTool does not support synchronous execution",
            "tool_id": self.id,
            "message": "Please use a_run or a_stream instead",
        }

    async def a_run(self, input: JSON, **kwargs) -> JSON:
        """异步执行前端工具。
        
        流程：
        1. 验证输入
        2. 创建 ToolCallRecord
        3. 等待前端返回结果（轮询）
        4. 返回结果或超时
        """
        try:
            # 验证输入
            data = input if isinstance(input, dict) else {"value": input}
            valid_input = self.validate_input(data)
            
            # 生成 trace_id 和 record_id
            trace_id = str(uuid.uuid4())
            record_id = f"frontend_{self.id}_{int(time.time() * 1000)}"
            
            # 获取轮询配置
            poll_interval_ms, timeout_ms = self._get_poll_config(**kwargs)
            
            # 获取记录查询函数（可选）
            record_fetcher = kwargs.get("record_fetcher")
            
            # 创建 ToolCallRecord（通过回调函数，因为 runtime 不直接操作数据库）
            record_creator = kwargs.get("record_creator")
            if record_creator:
                record_creator(
                    ToolCallRecordRef(
                        id=record_id,
                        tool_id=self.id,
                        tool_name=self.name,
                        input_data=valid_input,
                        status="pending",
                        trace_id=trace_id,
                    )
                )
            
            # 发送通知前端的事件（通过外部回调）
            frontend_notifier = kwargs.get("frontend_notifier")
            if frontend_notifier:
                frontend_notifier(
                    tool_id=self.id,
                    tool_name=self.name,
                    trace_id=trace_id,
                    record_id=record_id,
                    input_data=valid_input,
                )
            
            # 等待前端返回结果
            wait_result = await self._wait_for_frontend_result(
                trace_id=trace_id,
                record_id=record_id,
                poll_interval_ms=poll_interval_ms,
                timeout_ms=timeout_ms,
                record_fetcher=record_fetcher,
            )
            
            # 处理结果
            if wait_result["status"] == "success":
                result = wait_result["result"]
                shaped = self.validate_output(result)
                return {
                    "tool_id": self.id,
                    "version_id": getattr(self, "version_id", None),
                    "output": shaped,
                    "raw": result,
                    "trace_id": trace_id,
                    "record_id": record_id,
                }
            elif wait_result["status"] == "timeout":
                # 超时返回特殊标记
                return {
                    "tool_id": self.id,
                    "version_id": getattr(self, "version_id", None),
                    "output": wait_result["result"],
                    "raw": wait_result["result"],
                    "trace_id": trace_id,
                    "record_id": record_id,
                    "status": "timeout",
                    "requires_user_confirmation": True,
                }
            else:
                # 错误
                return {
                    "error": wait_result["result"].get("error", "前端工具执行失败"),
                    "tool_id": self.id,
                    "trace_id": trace_id,
                    "record_id": record_id,
                }
                
        except Exception as e:
            logger.exception(f"FrontendTool 异步执行失败: {e}")
            return {"error": str(e), "tool_id": self.id}

    def stream(self, input: JSON, **kwargs) -> Iterator[BaseEvent]:
        """同步流式执行（前端工具不支持真正的同步流式）。
        
        返回单一错误事件。
        """
        trace_id = str(uuid.uuid4())
        run_id = f"run_{self.id}_{int(time.time() * 1000)}"
        executor_type = self.run_type
        executor_id = build_executor_id(self, executor_type)
        executor_path = [executor_id]

        yield RunLifecycleFailedEvent(
            trace_id=trace_id,
            trace_path=[],
            run_id=run_id,
            run_path=[run_id],
            executor_id=executor_id,
            executor_type=executor_type,
            executor_path=executor_path,
            payload=RunLifecycleFailedPayload(
                error={
                    "code": "NotSupported",
                    "message": "FrontendTool does not support synchronous stream",
                }
            ),
        )

    async def a_stream(self, input: JSON, **kwargs) -> AsyncIterator[BaseEvent]:
        """异步流式执行前端工具。
        
        流程：
        1. 发送生命周期开始事件
        2. 创建 ToolCallRecord 并通知前端
        3. 轮询等待结果（期间可发送进度事件）
        4. 发送结果事件和生命周期结束事件
        """
        trace_id = str(uuid.uuid4())
        run_id = f"run_{self.id}_{int(time.time() * 1000)}"
        executor_type = self.run_type
        executor_id = build_executor_id(self, executor_type)
        executor_path = [executor_id]

        # 获取轮询配置
        poll_interval_ms, timeout_ms = self._get_poll_config(**kwargs)
        
        # 获取回调函数
        record_creator = kwargs.get("record_creator")
        record_fetcher = kwargs.get("record_fetcher")
        frontend_notifier = kwargs.get("frontend_notifier")

        # 发送生命周期开始事件
        yield RunLifecycleCreatedEvent(
            trace_id=trace_id,
            trace_path=[],
            run_id=run_id,
            run_path=[run_id],
            executor_id=executor_id,
            executor_type=executor_type,
            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=self.version_id,
                version=self.version_tag,
            ),
            user_metadata=self._resolve_user_metadata(**kwargs),
            payload=RunLifecycleCreatedPayload(input_data=input),
        )

        try:
            # 验证输入
            data = input if isinstance(input, dict) else {"value": input}
            valid_input = self.validate_input(data)
            
            # 生成 record_id
            record_id = f"frontend_{self.id}_{int(time.time() * 1000)}"

            # 发送结果开始事件
            yield ContentAgentResultStartEvent(
                trace_id=trace_id,
                run_id=run_id,
                executor_id=executor_id,
                executor_type=executor_type,
                executor_path=executor_path,
                payload=ContentAgentResultStartPayload(mode="json"),
            )

            # 创建 ToolCallRecord
            if record_creator:
                record_creator(
                    ToolCallRecordRef(
                        id=record_id,
                        tool_id=self.id,
                        tool_name=self.name,
                        input_data=valid_input,
                        status="pending",
                        trace_id=trace_id,
                    )
                )

            # 通知前端
            if frontend_notifier:
                frontend_notifier(
                    tool_id=self.id,
                    tool_name=self.name,
                    trace_id=trace_id,
                    record_id=record_id,
                    input_data=valid_input,
                )

            # 等待前端返回结果
            wait_result = await self._wait_for_frontend_result(
                trace_id=trace_id,
                record_id=record_id,
                poll_interval_ms=poll_interval_ms,
                timeout_ms=timeout_ms,
                record_fetcher=record_fetcher,
            )

            # 处理结果
            if wait_result["status"] == "success":
                result = wait_result["result"]
                shaped = self.validate_output(result)
                
                yield ContentAgentResultEndEvent(
                    trace_id=trace_id,
                    run_id=run_id,
                    executor_id=executor_id,
                    executor_path=executor_path,
                    executor_type=executor_type,
                    payload=ContentAgentResultEndPayload(
                        status="success",
                        full_result=shaped,
                    ),
                )
                yield RunLifecycleCompletedEvent(
                    trace_id=trace_id,
                    trace_path=[],
                    run_id=run_id,
                    run_path=[run_id],
                    executor_id=executor_id,
                    executor_path=executor_path,
                    executor_type=executor_type,
                    payload=RunLifecycleCompletedPayload(
                        output=shaped,
                        usage={"total_tokens": 0},
                    ),
                )
                
            elif wait_result["status"] == "timeout":
                # 超时：返回默认结果，提示需要用户确认
                timeout_result = wait_result["result"]
                
                yield ContentAgentResultEndEvent(
                    trace_id=trace_id,
                    run_id=run_id,
                    executor_id=executor_id,
                    executor_path=executor_path,
                    executor_type=executor_type,
                    payload=ContentAgentResultEndPayload(
                        status="timeout",
                        full_result=timeout_result,
                    ),
                )
                # 超时视为完成，但状态特殊
                yield RunLifecycleCompletedEvent(
                    trace_id=trace_id,
                    trace_path=[],
                    run_id=run_id,
                    run_path=[run_id],
                    executor_id=executor_id,
                    executor_path=executor_path,
                    executor_type=executor_type,
                    payload=RunLifecycleCompletedPayload(
                        output=timeout_result,
                        usage={"total_tokens": 0},
                    ),
                )
            else:
                # 错误
                error_result = wait_result["result"]
                yield ContentAgentResultEndEvent(
                    trace_id=trace_id,
                    run_id=run_id,
                    executor_id=executor_id,
                    executor_path=executor_path,
                    executor_type=executor_type,
                    payload=ContentAgentResultEndPayload(
                        status="error",
                        full_result=error_result,
                    ),
                )
                yield RunLifecycleFailedEvent(
                    trace_id=trace_id,
                    trace_path=[],
                    run_id=run_id,
                    run_path=[run_id],
                    executor_id=executor_id,
                    executor_path=executor_path,
                    executor_type=executor_type,
                    payload=RunLifecycleFailedPayload(
                        error={
                            "code": "FrontendToolError",
                            "message": error_result.get("message", "前端工具执行失败"),
                        }
                    ),
                )

        except Exception as e:
            logger.exception(f"FrontendTool 异步流式执行失败: {e}")
            yield ContentAgentResultEndEvent(
                trace_id=trace_id,
                run_id=run_id,
                executor_id=executor_id,
                executor_path=executor_path,
                executor_type=executor_type,
                payload=ContentAgentResultEndPayload(
                    status="error",
                    full_result=str(e),
                ),
            )
            yield RunLifecycleFailedEvent(
                trace_id=trace_id,
                trace_path=[],
                run_id=run_id,
                run_path=[run_id],
                executor_id=executor_id,
                executor_path=executor_path,
                executor_type=executor_type,
                payload=RunLifecycleFailedPayload(
                    error={"code": "RuntimeError", "message": str(e)}
                ),
            )
