using System.Runtime.CompilerServices;
using NewLife.AI.Clients;
using NewLife.AI.Models;
using NewLife.Log;

namespace NewLife.AI.Tools;

/// <summary>工具对话客户端中间件。注入多个 <see cref="IToolProvider"/> 的工具定义，并自动处理多轮工具调用回路</summary>
/// <remarks>
/// 工作流（非流式 / 流式统一）：
/// <list type="number">
/// <item>请求前，聚合所有 <see cref="Providers"/> 的工具定义与 <c>ChatOptions.Tools</c></item>
/// <item>调用内层客户端获取响应</item>
/// <item>若响应含 <c>tool_calls</c>，按工具名路由到对应 Provider 执行 <see cref="ExecuteToolAsync"/></item>
/// <item>循环重新调用模型，直到无更多工具调用（最多 <see cref="MaxIterations"/> 轮）</item>
/// </list>
/// 使用方式：
/// <code>
/// var client = provider.CreateClient(providerOptions)
///     .AsBuilder()
///     .UseTools(registry, mcpProvider)  // 多个 IToolProvider 按工具名路由
///     .Build();
/// </code>
/// </remarks>
public class ToolChatClient : DelegatingChatClient, ILogFeature, ITracerFeature
{
    #region 属性

    /// <summary>工具提供者列表（按工具名直接路由执行工具调用）</summary>
    public IReadOnlyList<IToolProvider> Providers { get; }

    private Int32 _maxIterations = 10;
    /// <summary>最大工具调用循环次数，防止无限递归。默认 10；设为 0 或负数时自动回退为 10</summary>
    public Int32 MaxIterations { get => _maxIterations; set => _maxIterations = value > 0 ? value : 10; }

    /// <summary>工具结果最大字符数。超过此长度时自动截断并追加省略提示，0表示不限制</summary>
    public Int32 MaxResultLength { get; set; }

    /// <summary>工具审批提供者。设置后在每次工具执行前请求审批，未设置时直接执行</summary>
    public IToolApprovalProvider? ApprovalProvider { get; set; }
    #endregion

    #region 构造

    /// <summary>初始化工具对话客户端中间件</summary>
    /// <param name="innerClient">内层客户端</param>
    /// <param name="providers">工具提供者列表（按工具名路由；未找到则抛 <see cref="InvalidOperationException"/>）</param>
    public ToolChatClient(IChatClient innerClient, params IToolProvider[] providers) : base(innerClient)
    {
        Providers = (providers ?? []).ToList().AsReadOnly();
    }

    #endregion

    #region 方法

    /// <summary>非流式对话完成。注入工具定义并自动处理工具调用回路</summary>
    /// <param name="request">内部对话请求</param>
    /// <param name="cancellationToken">取消令牌</param>
    public override async Task<IChatResponse> GetResponseAsync(IChatRequest request, CancellationToken cancellationToken = default)
    {
        if (request == null) throw new ArgumentNullException(nameof(request));

        var (mergedTools, toolMap) = GetMergedTools(request);
        if (mergedTools.Count == 0)
            return await InnerClient.GetResponseAsync(request, cancellationToken).ConfigureAwait(false);

        // 合并工具定义到选项（不修改调用方的原始选项）
        var workOptions = MergeToolOptions(request, mergedTools);
        var workMessages = request.Messages.ToList();

        IChatResponse response;
        var iterations = 0;
        UsageDetails? accumulatedUsage = null;

        while (true)
        {
            cancellationToken.ThrowIfCancellationRequested();
            response = await InnerClient.GetResponseAsync(ChatRequest.Create(workMessages, workOptions), cancellationToken).ConfigureAwait(false);

            // 累加每轮 LLM 调用的 Token 用量（N 次工具调用 = N+1 次 LLM 调用，每轮都有独立 Usage）
            if (response.Usage != null) accumulatedUsage = accumulatedUsage?.Add(response.Usage) ?? response.Usage;

            // 从第一个 Choice 中获取工具调用
            var assistantMessage = response.Messages?.FirstOrDefault()?.Message;
            var toolCalls = assistantMessage?.ToolCalls;
            if (toolCalls == null || toolCalls.Count == 0) break;
            if (++iterations > MaxIterations) break;

            // 追加 assistant 消息（含工具调用）
            // DeepSeek 思考模式要求：有工具调用时必须将 reasoning_content 一并回传，否则 API 返回 400
            workMessages.Add(new ChatMessage
            {
                Role = "assistant",
                Content = assistantMessage?.Content,
                ReasoningContent = assistantMessage?.ReasoningContent,
                ToolCalls = toolCalls.Select(tc => new ToolCall { Id = tc.Id, Type = tc.Type, Function = tc.Function }).ToList(),
            });

            // 依次执行所有工具调用
            foreach (var tc in toolCalls)
            {
                if (tc.Function == null) continue;
                var result = await ExecuteToolAsync(tc.Function.Name, tc.Function.Arguments, toolMap, cancellationToken).ConfigureAwait(false);
                workMessages.Add(new ChatMessage { Role = "tool", ToolCallId = tc.Id, Content = result });
            }
        }

        // 将所有轮次的 Token 用量累加值写回最终 response，供上层（如 InvokeLlmDirectAsync）使用
        if (accumulatedUsage != null) response.Usage = accumulatedUsage;

        return response;
    }

    /// <summary>流式对话完成。注入工具定义，流式执行多轮工具调用回路，对外透明</summary>
    /// <param name="request">内部对话请求</param>
    /// <param name="cancellationToken">取消令牌</param>
    public override async IAsyncEnumerable<IChatResponse> GetStreamingResponseAsync(
        IChatRequest request,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        if (request == null) throw new ArgumentNullException(nameof(request));

        var (mergedTools, toolMap) = GetMergedTools(request);
        if (mergedTools.Count == 0)
        {
            await foreach (var chunk in InnerClient.GetStreamingResponseAsync(request, cancellationToken).ConfigureAwait(false))
                yield return chunk;
            yield break;
        }

        var workOptions = MergeToolOptions(request, mergedTools);
        var workMessages = request.Messages.ToList();

        UsageDetails? accumulatedUsage = null;

        for (var iteration = 0; iteration < MaxIterations; iteration++)
        {
            cancellationToken.ThrowIfCancellationRequested();

            var toolCallCollector = new List<ToolCall>();
            String? finishReason = null;
            var assistantContent = (String?)null;
            var assistantReasoningContent = (String?)null;
            UsageDetails? iterUsage = null;

            await foreach (var chunk in InnerClient.GetStreamingResponseAsync(ChatRequest.Create(workMessages, workOptions, stream: true), cancellationToken).ConfigureAwait(false))
            {
                // 轮次内合并 chunk Usage（各协议差异由 MergeChunkUsage 虚拟方法处理）
                if (chunk.Usage != null)
                    iterUsage = MergeChunkUsage(iterUsage, chunk.Usage);

                var choice = chunk.Messages?.FirstOrDefault();
                if (choice != null)
                {
                    finishReason = choice.FinishReason?.ToApiString() ?? finishReason;
                    var delta = choice.Delta;
                    if (delta != null)
                    {
                        // 累积正文内容（供追加 assistant 消息）
                        var text = delta.Content as String;
                        if (!String.IsNullOrEmpty(text))
                            assistantContent = (assistantContent ?? String.Empty) + text;

                        // 累积思维链内容（DeepSeek 思考模式要求：有工具调用时必须将 reasoning_content 一并回传）
                        if (!String.IsNullOrEmpty(delta.ReasoningContent))
                            assistantReasoningContent = (assistantReasoningContent ?? String.Empty) + delta.ReasoningContent;

                        // 合并流式 tool_calls 增量
                        if (delta.ToolCalls != null)
                        {
                            foreach (var tc in delta.ToolCalls)
                                MergeToolCallDelta(toolCallCollector, tc);
                        }
                    }
                }

                // 始终透传原始 chunk，不做任何抑制
                yield return chunk;

                // 尽早原则：多轮场景下（有历史轮 accumulatedUsage），每个含 Usage 的 chunk 后
                // 立即追加一个运行时累计总量 chunk，让消费方随时能获取到正确的跨轮累计值
                if (chunk.Usage != null && accumulatedUsage != null)
                    yield return new ChatResponse { Usage = accumulatedUsage.Add(iterUsage!) };
            }

            // 跨轮 Token 累加：将本轮 Usage 加到全局累加值
            if (iterUsage != null)
            {
                DefaultSpan.Current?.AppendTag($"Tokens: {iterUsage.InputTokens}+{iterUsage.OutputTokens}={iterUsage.TotalTokens} finishReason: {finishReason}");

                accumulatedUsage = accumulatedUsage?.Add(iterUsage) ?? iterUsage;
            }

            var isToolRound = finishReason.EqualIgnoreCase("tool_calls") ||
                              (toolCallCollector.Count > 0 && String.IsNullOrEmpty(finishReason));

            if (!isToolRound || toolCallCollector.Count == 0)
            {
                // 兜底：最终轮无 Usage chunk 但存在历史轮（极少见），补发累计总量
                if (iterUsage == null && accumulatedUsage != null)
                    yield return new ChatResponse { Usage = accumulatedUsage };
                yield break;
            }

            // 追加 assistant 消息（含工具调用）
            workMessages.Add(new ChatMessage
            {
                Role = "assistant",
                Content = assistantContent,
                ReasoningContent = assistantReasoningContent,
                ToolCalls = toolCallCollector.ToList(),
            });

            foreach (var tc in toolCallCollector)
            {
                if (tc.Function == null) continue;

                // 通知调用方：工具调用开始
                yield return new ChatResponse
                {
                    ToolCallEvents = [new ToolCallEventInfo("start", tc.Id, tc.Function.Name, tc.Function.Arguments)]
                };

                // 在 try/catch 中执行工具，收集结果（yield 不能出现在 try/catch 内）
                ToolCallEventInfo resultEvent;
                var span = Tracer?.NewSpan($"ai:tool:{tc.Function.Name}", tc.Function.Arguments);
                try
                {
                    var toolResult = await ExecuteToolAsync(tc.Function.Name, tc.Function.Arguments, toolMap, cancellationToken).ConfigureAwait(false);
                    workMessages.Add(new ChatMessage { Role = "tool", ToolCallId = tc.Id, Content = toolResult });
                    if (span != null && toolResult != null)
                        span.AppendTag(toolResult, toolResult.Length);

                    resultEvent = new ToolCallEventInfo("done", tc.Id, tc.Function.Name, toolResult);
                }
                catch (Exception ex)
                {
                    span?.SetError(ex, null);
                    // 将错误作为工具结果反馈给模型，让模型有机会修正，不中断流
                    workMessages.Add(new ChatMessage { Role = "tool", ToolCallId = tc.Id, Content = $"Error: {ex.Message}" });
                    resultEvent = new ToolCallEventInfo("error", tc.Id, tc.Function.Name, ex.Message);
                }
                finally
                {
                    span?.Dispose();
                }

                // 通知调用方：工具调用结果
                yield return new ChatResponse { ToolCallEvents = [resultEvent] };
            }
            // 继续下一轮（下一轮流的 chunk 透传给调用方）
        }
        // 超过最大轮次，静默退出（调用方已收到全部 chunk）
    }

    #endregion

    #region 辅助

    /// <summary>按工具名路由到对应 Provider 执行工具调用。未找到则抛 <see cref="InvalidOperationException"/></summary>
    /// <param name="toolName">工具名称</param>
    /// <param name="argumentsJson">参数 JSON 字符串（模型原文）</param>
    /// <param name="toolMap">工具名到 Provider 的路由字典</param>
    /// <param name="cancellationToken">取消令牌</param>
    private async Task<String> ExecuteToolAsync(String toolName, String? argumentsJson, Dictionary<String, IToolProvider> toolMap, CancellationToken cancellationToken)
    {
        if (!toolMap.TryGetValue(toolName, out var provider))
            throw new InvalidOperationException($"Tool not found: '{toolName}', searched {Providers.Count} providers");

        // 审批拦截：设置了 ApprovalProvider 时在执行前请求用户确认
        if (ApprovalProvider != null)
        {
            var approval = await ApprovalProvider.RequestApprovalAsync(toolName, argumentsJson, cancellationToken).ConfigureAwait(false);
            if (!approval.Approved)
                return $"工具调用被用户拒绝: {toolName}";
        }

        var result = await provider.CallToolAsync(toolName, argumentsJson, cancellationToken).ConfigureAwait(false);

        // 结果超长时截断并追加省略提示
        if (MaxResultLength > 0 && result != null && result.Length > MaxResultLength)
        {
            result = result.Substring(0, MaxResultLength) + $"\n\n[... 内容已截断，原始长度 {result.Length} 字符，仅保留前 {MaxResultLength} 字符]";
        }

        return result!;
    }

    /// <summary>合并流式 tool_call 增量到收集列表。OpenAI 流式协议中 tool_calls 分块到达</summary>
    private static void MergeToolCallDelta(List<ToolCall> collector, ToolCall delta)
    {
        if (delta == null) return;

        ToolCall? existing = null;
        if (!String.IsNullOrEmpty(delta.Id))
            existing = collector.FirstOrDefault(t => t.Id == delta.Id);
        else if (delta.Index != null)
            existing = collector.FirstOrDefault(t => t.Index == delta.Index);
        else if (collector.Count > 0)
            existing = collector[collector.Count - 1];  // 兜底取最后一个（单工具调用时常见）

        if (existing == null && !String.IsNullOrEmpty(delta.Id))
        {
            collector.Add(new ToolCall
            {
                Index = delta.Index,
                Id = delta.Id,
                Type = delta.Type,
                Function = new FunctionCall
                {
                    Name = delta.Function?.Name ?? String.Empty,
                    Arguments = delta.Function?.Arguments ?? String.Empty,
                },
            });
            return;
        }

        if (existing?.Function != null && delta.Function != null)
        {
            if (!String.IsNullOrEmpty(delta.Function.Name))
                existing.Function.Name += delta.Function.Name;
            if (!String.IsNullOrEmpty(delta.Function.Arguments))
                existing.Function.Arguments += delta.Function.Arguments;
        }
    }

    /// <summary>聚合所有提供者的工具定义，合并 options.Tools，同时建立工具名到 Provider 的路由字典</summary>
    private (List<ChatTool> tools, Dictionary<String, IToolProvider> toolMap) GetMergedTools(IChatRequest? options)
    {
        var tools = new List<ChatTool>();
        var toolMap = new Dictionary<String, IToolProvider>(StringComparer.OrdinalIgnoreCase);
        foreach (var provider in Providers)
        {
            foreach (var t in provider.GetTools())
            {
                tools.Add(t);
                var name = t.Function?.Name;
                if (!name.IsNullOrEmpty())
                    toolMap[name] = provider;
            }
        }
        if (options?.Tools != null)
        {
            foreach (var t in options.Tools)
                tools.Add(t);
        }
        return (tools, toolMap);
    }

    /// <summary>克隆 ChatOptions 并注入合并后的工具列表（不修改调用方的原始选项）</summary>
    private static ChatOptions MergeToolOptions(IChatRequest? request, List<ChatTool> mergedTools)
        => new()
        {
            Model = request?.Model,
            Temperature = request?.Temperature,
            TopP = request?.TopP,
            TopK = request?.TopK,
            MaxTokens = request?.MaxTokens,
            Stop = request?.Stop,
            PresencePenalty = request?.PresencePenalty,
            FrequencyPenalty = request?.FrequencyPenalty,
            Tools = mergedTools,
            ToolChoice = request?.ToolChoice ?? "auto",
            User = request?.User,
            EnableThinking = request?.EnableThinking,
            ResponseFormat = request?.ResponseFormat,
            ParallelToolCalls = request?.ParallelToolCalls,
            UserId = request?.UserId,
            ConversationId = request?.ConversationId,
            Items = request?.Items ?? new Dictionary<String, Object?>(),
        };

    #endregion

    #region 日志
    /// <summary>日志</summary>
    public ILog Log { get; set; } = Logger.Null;

    /// <summary>追踪器</summary>
    public ITracer? Tracer { get; set; }
    #endregion
}