"""
使用装饰器实现全栈函数调用追踪
解决LangGraph无法捕获子函数内部状态的问题
"""

import functools
import inspect
import asyncio
import contextvars
from typing import Callable, Any, Dict, Optional, List
from datetime import datetime
from collections import deque

# 使用 contextvars 实现请求级别的隔离，避免多个并发请求互相干扰
_event_callbacks: contextvars.ContextVar[List[Callable]] = contextvars.ContextVar('event_callbacks', default=None)
_current_stream_handler: contextvars.ContextVar[Optional[Any]] = contextvars.ContextVar('current_stream_handler',
                                                                                        default=None)


def register_event_callback(callback: Callable[[Dict[str, Any]], None]):
    """
    注册事件回调函数（使用 contextvars 实现请求隔离）
    """
    callbacks = _event_callbacks.get()
    if callbacks is None:
        callbacks = []
        _event_callbacks.set(callbacks)
    callbacks.append(callback)


def unregister_event_callback(callback: Callable[[Dict[str, Any]], None]):
    """
    移除事件回调函数（使用 contextvars 实现请求隔离）
    """
    callbacks = _event_callbacks.get()
    if callbacks and callback in callbacks:
        callbacks.remove(callback)


def clear_event_callbacks():
    """
    清空所有事件回调（使用 contextvars 实现请求隔离）
    """
    callbacks = _event_callbacks.get()
    if callbacks:
        callbacks.clear()


def get_all_callbacks():
    """
    获取所有注册的回调函数（使用 contextvars 实现请求隔离）
    """
    callbacks = _event_callbacks.get()
    return callbacks.copy() if callbacks else []


def extract_function_description(func: Callable) -> str:
    """
    提取函数的描述（docstring的第一行或前100个字符）
    """
    if func.__doc__:
        doc = func.__doc__.strip()
        # 取第一行或前100个字符
        first_line = doc.split('\n')[0].strip()
        return first_line[:100] if len(first_line) > 100 else first_line
    return ""


def _extract_details_from_result(result: Any) -> str:
    """
    从函数返回值中提取 details，顺序固定为：先 details 再中文。
    1) 先看返回的 dict 里是否有「details 类」字段（__details__、_details、details 或键名含 "details"），
       且其值为非空，则用该值；
    2) 若没有或为空，再走原本的中文兜底：从返回结构中找一段包含中文的内容作为 details。
    """
    default = "完成"
    if not result:
        return default

    # 第一步：优先用 details 类字段（仅当有值且非空时采用）
    if isinstance(result, dict):
        for key in ("__details__", "_details", "details"):
            if key in result and result[key] is not None:
                s = str(result[key]).strip()
                if s:
                    return str(result[key])
        for k, v in result.items():
            if "details" in k.lower() and v is not None:
                s = str(v).strip()
                if s:
                    return str(v)

    # 第二步：保留原本的中文兜底
    def _contains_chinese(text: str) -> bool:
        if not text:
            return False
        return any("\u4e00" <= c <= "\u9fff" for c in str(text))

    def _find_chinese(obj: Any, max_length: int = 200) -> Optional[str]:
        if isinstance(obj, str):
            return (obj[:max_length] + "...") if _contains_chinese(obj) and len(obj) > max_length else (
                obj if _contains_chinese(obj) else None)
        if isinstance(obj, dict):
            for key, value in obj.items():
                if "details" in key.lower():
                    continue
                found = _find_chinese(value, max_length)
                if found:
                    return found
        if isinstance(obj, list):
            for item in obj:
                found = _find_chinese(item, max_length)
                if found:
                    return found
        return None

    chinese = _find_chinese(result)
    return chinese if chinese else default


def track_function_calls(func: Callable) -> Callable:
    """
    装饰器：追踪函数调用过程，支持同步和异步函数
    
    Args:
        func: 被装饰的函数
    """

    @functools.wraps(func)
    def sync_wrapper(*args, **kwargs):
        # 获取函数签名信息
        sig = inspect.signature(func)
        bound_args = sig.bind(*args, **kwargs)
        bound_args.apply_defaults()

        func_name = func.__name__
        module_name = func.__module__

        # 尝试从函数参数中获取 state
        state = None
        if args and isinstance(args[0], dict):
            state = args[0]
        elif 'state' in bound_args.arguments:
            state = bound_args.arguments.get('state')

        func_description = extract_function_description(func)

        # 准备函数参数信息
        params = {k: v for k, v in bound_args.arguments.items()}

        try:
            # 执行原函数
            result = func(*args, **kwargs)

            # 构建 details：优先用返回值中的 details 类字段（__details__、_details、details 或键名含 details）
            details = _extract_details_from_result(result)

            # 合并函数执行事件（只在执行完成后发送一次）
            execution_event = {
                "type": "function_execution",
                "title": func_description,
                "details": details,
                "result": result,  # 添加函数返回结果
                "timestamp": datetime.now().isoformat()
            }

            # 调用所有注册的回调（从上下文变量获取）
            callbacks = _event_callbacks.get() or []
            for callback in callbacks:
                try:
                    callback(execution_event)
                except Exception as e:
                    print(f"回调执行出错: {e}")

            return result

        except Exception as e:
            # 发送函数执行失败事件
            error_event = {
                "type": "function_error",
                "title": func_description,
                "details": f"错误: {str(e)}",
                "timestamp": datetime.now().isoformat()
            }

            callbacks = _event_callbacks.get() or []
            for callback in callbacks:
                try:
                    callback(error_event)
                except Exception as ex:
                    print(f"回调执行出错: {ex}")

            raise

    @functools.wraps(func)
    async def async_wrapper(*args, **kwargs):
        # 获取函数签名信息
        sig = inspect.signature(func)
        bound_args = sig.bind(*args, **kwargs)
        bound_args.apply_defaults()

        func_name = func.__name__
        module_name = func.__module__

        # 尝试从函数参数中获取 state
        state = None
        if args and isinstance(args[0], dict):
            state = args[0]
        elif 'state' in bound_args.arguments:
            state = bound_args.arguments.get('state')

        func_description = extract_function_description(func)

        # 准备函数参数信息
        params = {k: v for k, v in bound_args.arguments.items()}

        try:
            # 执行原函数
            result = await func(*args, **kwargs)

            # 构建 details：优先用返回值中的 details 类字段（__details__、_details、details 或键名含 details）
            details = _extract_details_from_result(result)

            # 合并函数执行事件（只在执行完成后发送一次）
            execution_event = {
                "type": "function_execution",
                "title": func_description,
                "details": details,
                "result": result,  # 添加函数返回结果
                "timestamp": datetime.now().isoformat()
            }

            # 调用所有注册的回调（从上下文变量获取）
            callbacks = _event_callbacks.get() or []
            for callback in callbacks:
                try:
                    callback(execution_event)
                except Exception as e:
                    print(f"回调执行出错: {e}")

            return result

        except Exception as e:
            # 发送函数执行失败事件
            error_event = {
                "type": "function_error",
                "title": func_description,
                "details": f"错误: {str(e)}",
                "timestamp": datetime.now().isoformat()
            }

            callbacks = _event_callbacks.get() or []
            for callback in callbacks:
                try:
                    callback(error_event)
                except Exception as ex:
                    print(f"回调执行出错: {ex}")

            raise
    
    # 根据函数是否为协程函数返回对应的装饰器
    if asyncio.iscoroutinefunction(func):
        return async_wrapper
    else:
        return sync_wrapper


# 事件处理器
def console_event_handler(event: Dict[str, Any]):
    """
    控制台事件处理器
    """
    event_type = event["type"]
    title = event.get("title", "")
    details = event.get("details", "")

    if event_type == "function_execution":
        print(f"✅ 函数执行: {title}")
        if details:
            print(f"   详情: {details}")

    elif event_type == "function_error":
        print(f"❌ 函数错误: {title}")
        if details:
            print(f"   详情: {details}")


class StreamEventHandler:
    """
    流式事件处理器，将装饰器事件放入队列供流式输出使用
    """
    def __init__(self, queue: asyncio.Queue):
        self.queue = queue
    
    def __call__(self, event: Dict[str, Any]):
        """
        处理事件，将事件放入队列
        注意：这是同步方法，但队列是异步的，需要特殊处理
        """
        try:
            # 使用线程安全的方式将事件放入队列
            # 如果队列已满，使用 put_nowait 并忽略错误
            try:
                self.queue.put_nowait(event)
            except asyncio.QueueFull:
                # 队列已满，跳过此事件（不应该发生，因为队列大小足够）
                pass
        except Exception as e:
            # 忽略错误，避免影响主流程
            pass
    
    def send_stream_content(self, content: str, title: str = ""):
        """
        发送流式内容事件
        
        Args:
            content: 流式内容片段
            title: 函数描述/标题（可选）
        """
        try:
            event = {
                "type": "stream_content",
                "content": content,
                "title": title,
                "timestamp": datetime.now().isoformat()
            }
            self.queue.put_nowait(event)
        except Exception as e:
            # 忽略错误，避免影响主流程
            pass


def create_stream_event_handler() -> tuple[StreamEventHandler, asyncio.Queue]:
    """
    创建流式事件处理器和对应的队列（使用 contextvars 实现请求隔离）
    
    Returns:
        (handler, queue): 事件处理器和事件队列
    """
    queue = asyncio.Queue(maxsize=1000)  # 设置足够大的队列大小
    handler = StreamEventHandler(queue)
    _current_stream_handler.set(handler)  # 保存到上下文变量
    return handler, queue


def get_current_stream_handler() -> Optional[StreamEventHandler]:
    """
    获取当前活动的流式事件处理器（使用 contextvars 实现请求隔离）
    
    Returns:
        当前流式事件处理器，如果没有则返回 None
    """
    return _current_stream_handler.get()


def clear_current_stream_handler():
    """
    清除当前活动的流式事件处理器（使用 contextvars 实现请求隔离）
    """
    _current_stream_handler.set(None)
