using System.ClientModel;
using Dpz.Core.Entity.Base;
using Dpz.Core.Service.Network.Models;
using ModelContextProtocol.Client;
using OpenAI;
using OpenAI.Chat;
using OpenAI.Models;
using ChatMessage = OpenAI.Chat.ChatMessage;

namespace Dpz.Core.Service.Network;

/// <summary>
/// 基于 OpenAI SDK 的聊天发送和模型查询实现
/// </summary>
public class OpenAiService : IOpenAiService
{
    private const string ConfigurationSectionName = "OpenAI";

    private readonly OpenAiChatOptions _defaults;
    private readonly ChatClient _chatClient;
    private readonly OpenAIModelClient _modelClient;
    private readonly ILogger<OpenAiService> _logger;

    /// <summary>
    /// 初始化 <see cref="OpenAiService"/> 类的新实例
    /// </summary>
    /// <param name="configuration">包含所需 OpenAI 配置节的应用配置</param>
    /// <param name="logger">用于请求诊断的日志记录器</param>
    public OpenAiService(IConfiguration configuration, ILogger<OpenAiService> logger)
    {
        _logger = logger;
        _defaults = LoadOptions(configuration);
        _chatClient = CreateChatClient(_defaults);
        _modelClient = CreateModelClient(_defaults);
    }

    /// <inheritdoc />
    public async Task<ResponseResult<OpenAiChatStreamResult>> SendMessageAsync(
        IReadOnlyList<ChatMessage> messages,
        Func<string, CancellationToken, Task>? onDeltaContentAsync = null,
        Action<OpenAiChatRequestOptions>? options = null,
        CancellationToken cancellationToken = default
    )
    {
        var result = new ResponseResult<OpenAiChatStreamResult>();
        if (messages is not { Count: > 0 })
        {
            return result.WithFail("Messages cannot be empty");
        }

        var requestOptions = CreateRequestOptions();
        options?.Invoke(requestOptions);

        var model = string.IsNullOrWhiteSpace(requestOptions.Model)
            ? _defaults.Model
            : requestOptions.Model;
        var completionOptions = CreateCompletionOptions(requestOptions);
        var chatClient = string.Equals(model, _defaults.Model, StringComparison.Ordinal)
            ? _chatClient
            : CreateChatClient(_defaults, model);

        try
        {
            var contentBuilder = new StringBuilder();
            ChatTokenUsage? usage = null;
            var responseModel = model;

            var updates = chatClient.CompleteChatStreamingAsync(
                messages,
                completionOptions,
                cancellationToken
            );

            await foreach (var update in updates)
            {
                cancellationToken.ThrowIfCancellationRequested();

                if (!string.IsNullOrWhiteSpace(update.Model))
                {
                    responseModel = update.Model;
                }

                if (update.Usage != null)
                {
                    usage = update.Usage;
                }

                // SDK 产出已解析的流式更新；此代码块仅负责拼接文本增量，
                // 并交由调用方实时持久化或转发。
                foreach (var contentPart in update.ContentUpdate)
                {
                    if (string.IsNullOrEmpty(contentPart.Text))
                    {
                        continue;
                    }

                    contentBuilder.Append(contentPart.Text);

                    // 将取消令牌传播到用户代码中，以便调用方能快速停止
                    if (onDeltaContentAsync != null)
                    {
                        await onDeltaContentAsync(contentPart.Text, cancellationToken);
                    }
                }
            }

            return result.WithOk(
                new OpenAiChatStreamResult
                {
                    Model = responseModel,
                    Content = contentBuilder.ToString(),
                    PromptTokens = usage?.InputTokenCount,
                    CompletionTokens = usage?.OutputTokenCount,
                    TotalTokens = usage?.TotalTokenCount,
                }
            );
        }
        catch (OperationCanceledException)
        {
            _logger.LogInformation("OpenAI streaming chat request was cancelled");
            throw;
        }
        catch (Exception e)
        {
            _logger.LogError(e, "OpenAI streaming chat request failed");
            return result.WithFail(e.Message);
        }
    }

    /// <inheritdoc />
    public async Task<ResponseResult<IReadOnlyList<OpenAIModel>>> GetModelsAsync(
        CancellationToken cancellationToken = default
    )
    {
        var result = new ResponseResult<IReadOnlyList<OpenAIModel>>();

        try
        {
            var models = await _modelClient.GetModelsAsync(cancellationToken);
            return result.WithOk(models.Value.ToList());
        }
        catch (OperationCanceledException)
        {
            _logger.LogInformation("OpenAI model list request was cancelled");
            throw;
        }
        catch (Exception e)
        {
            _logger.LogError(e, "OpenAI model list request failed");
            return result.WithFail(e.Message);
        }
    }

    /// <inheritdoc />
    public async Task<ResponseResult<OpenAiChatStreamResult>> SendMessageWithMcpAsync(
        IReadOnlyList<ChatMessage> messages,
        McpClient mcpClient,
        Func<string, CancellationToken, Task>? onDeltaContentAsync = null,
        Action<OpenAiChatRequestOptions>? options = null,
        int maxToolRounds = 8,
        Func<string, string, CancellationToken, Task>? onToolCallAsync = null,
        CancellationToken cancellationToken = default
    )
    {
        var result = new ResponseResult<OpenAiChatStreamResult>();
        if (messages is not { Count: > 0 })
        {
            return result.WithFail("Messages cannot be empty");
        }

        ArgumentNullException.ThrowIfNull(mcpClient);

        var mcpTools = await mcpClient.ListToolsAsync(cancellationToken: cancellationToken);
        if (mcpTools.Count == 0)
        {
            return await SendMessageAsync(
                messages,
                onDeltaContentAsync,
                options,
                cancellationToken
            );
        }

        var requestOptions = CreateRequestOptions();
        options?.Invoke(requestOptions);

        var model = string.IsNullOrWhiteSpace(requestOptions.Model)
            ? _defaults.Model
            : requestOptions.Model;
        var completionOptions = CreateCompletionOptions(requestOptions);
        foreach (var mcpTool in mcpTools)
        {
            completionOptions.Tools.Add(OpenAiMcpToolSupport.ToChatTool(mcpTool));
        }

        var chatClient = string.Equals(model, _defaults.Model, StringComparison.Ordinal)
            ? _chatClient
            : CreateChatClient(_defaults, model);
        var sendMessage = new List<ChatMessage>(messages);
        var allowedRounds = OpenAiMcpToolSupport.ClampMaxToolRounds(maxToolRounds);

        try
        {
            var usage = new ChatUsageAccumulator { ResponseModel = model };
            ChatRoundResult? lastRound = null;
            var preferStreaming = true;

            // 有剩余轮次时允许 function call; 最后一轮改 ToolChoice.None, 不清空 Tools
            // 清空工具声明会让部分兼容端把调用写进文本, 协议侧应继续声明工具但禁止再选
            for (var round = 0; round <= allowedRounds; round++)
            {
                var allowTools = round < allowedRounds;
                completionOptions.ToolChoice = allowTools
                    ? ChatToolChoice.CreateAutoChoice()
                    : ChatToolChoice.CreateNoneChoice();

                var roundResult = await CompleteChatRoundAsync(
                    chatClient,
                    sendMessage,
                    completionOptions,
                    usage,
                    onDeltaContentAsync,
                    preferStreaming,
                    cancellationToken
                );
                lastRound = roundResult;
                preferStreaming = roundResult.WasStreamed;

                if (roundResult.ToolCalls.Count == 0 || !allowTools)
                {
                    break;
                }

                sendMessage.Add(new AssistantChatMessage(roundResult.ToolCalls));
                await ExecuteToolCallsAsync(
                    mcpClient,
                    sendMessage,
                    roundResult.ToolCalls,
                    onToolCallAsync,
                    cancellationToken
                );
            }

            if (lastRound is null)
            {
                return result.WithFail("Model returned no completion");
            }

            if (lastRound.ToolCalls.Count > 0)
            {
                _logger.LogWarning(
                    "MCP chat ended with pending tool calls. FinishReason:{FinishReason} Count:{Count}",
                    lastRound.FinishReason,
                    lastRound.ToolCalls.Count
                );
                return result.WithFail("Model requested tool calls after the last round");
            }

            var content = lastRound.Text;
            if (string.IsNullOrWhiteSpace(content))
            {
                return result.WithFail("Model returned empty content");
            }

            return result.WithOk(
                new OpenAiChatStreamResult
                {
                    Model = usage.ResponseModel,
                    Content = content,
                    PromptTokens = usage.PromptTokens > 0 ? usage.PromptTokens : null,
                    CompletionTokens = usage.CompletionTokens > 0 ? usage.CompletionTokens : null,
                    TotalTokens = usage.TotalTokens > 0 ? usage.TotalTokens : null,
                }
            );
        }
        catch (OperationCanceledException)
        {
            _logger.LogInformation("OpenAI MCP chat request was cancelled");
            throw;
        }
        catch (Exception e)
        {
            _logger.LogError(e, "OpenAI MCP chat request failed");
            return result.WithFail(e.Message);
        }
    }

    private async Task ExecuteToolCallsAsync(
        McpClient mcpClient,
        List<ChatMessage> sendMessage,
        IReadOnlyList<ChatToolCall> toolCalls,
        Func<string, string, CancellationToken, Task>? onToolCallAsync,
        CancellationToken cancellationToken
    )
    {
        foreach (var toolCall in toolCalls)
        {
            cancellationToken.ThrowIfCancellationRequested();
            _logger.LogInformation(
                "MCP tool call {ToolName} {ToolCallId}",
                toolCall.FunctionName,
                toolCall.Id
            );

            var argumentsJson = toolCall.FunctionArguments?.ToString() ?? "";
            if (onToolCallAsync != null)
            {
                await onToolCallAsync(toolCall.FunctionName, argumentsJson, cancellationToken);
            }

            var arguments = OpenAiMcpToolSupport.ParseArguments(toolCall.FunctionArguments);
            string toolText;
            try
            {
                var toolResult = await mcpClient.CallToolAsync(
                    toolCall.FunctionName,
                    arguments,
                    cancellationToken: cancellationToken
                );
                toolText = OpenAiMcpToolSupport.FormatToolResult(toolResult);
            }
            catch (Exception ex) when (ex is not OperationCanceledException)
            {
                _logger.LogWarning(
                    ex,
                    "MCP tool call failed {ToolName} {ToolCallId}",
                    toolCall.FunctionName,
                    toolCall.Id
                );
                toolText = OpenAiMcpToolSupport.Truncate("Error: " + ex.Message);
            }

            sendMessage.Add(new ToolChatMessage(toolCall.Id, toolText));
        }
    }

    /// <summary>
    /// 以流式补全执行一轮对话；流式在首个数据帧之前失败时回退到非流式补全
    /// </summary>
    private async Task<ChatRoundResult> CompleteChatRoundAsync(
        ChatClient chatClient,
        IReadOnlyList<ChatMessage> sendMessage,
        ChatCompletionOptions completionOptions,
        ChatUsageAccumulator usage,
        Func<string, CancellationToken, Task>? onDeltaContentAsync,
        bool preferStreaming,
        CancellationToken cancellationToken
    )
    {
        if (preferStreaming)
        {
            var streamed = await TryCompleteChatStreamingRoundAsync(
                chatClient,
                sendMessage,
                completionOptions,
                usage,
                onDeltaContentAsync,
                cancellationToken
            );
            if (streamed != null)
            {
                return streamed;
            }

            _logger.LogWarning(
                "MCP streaming completion not supported by endpoint, fallback to non-streaming"
            );
        }

        var completion = await CompleteChatAndAccumulateUsageAsync(
            chatClient,
            sendMessage,
            completionOptions,
            usage,
            cancellationToken
        );
        var toolCalls = completion.ToolCalls.ToList();
        var text = GetCompletionText(completion);
        if (text.Length > 0 && toolCalls.Count == 0 && onDeltaContentAsync != null)
        {
            await onDeltaContentAsync(text, cancellationToken);
        }

        return new ChatRoundResult
        {
            Text = text,
            ToolCalls = toolCalls,
            FinishReason = completion.FinishReason,
            WasStreamed = false,
        };
    }

    private async Task<ChatRoundResult?> TryCompleteChatStreamingRoundAsync(
        ChatClient chatClient,
        IReadOnlyList<ChatMessage> sendMessage,
        ChatCompletionOptions completionOptions,
        ChatUsageAccumulator usage,
        Func<string, CancellationToken, Task>? onDeltaContentAsync,
        CancellationToken cancellationToken
    )
    {
        var contentBuilder = new StringBuilder();
        var toolCallAccumulators = new Dictionary<int, StreamingToolCallAccumulator>();
        ChatFinishReason? finishReason = null;
        var receivedAnyUpdate = false;

        await using var enumerator = chatClient
            .CompleteChatStreamingAsync(sendMessage, completionOptions, cancellationToken)
            .GetAsyncEnumerator(cancellationToken);

        try
        {
            while (await enumerator.MoveNextAsync())
            {
                receivedAnyUpdate = true;
                var update = enumerator.Current;
                cancellationToken.ThrowIfCancellationRequested();

                if (!string.IsNullOrWhiteSpace(update.Model))
                {
                    usage.ResponseModel = update.Model;
                }

                if (update.Usage != null)
                {
                    usage.PromptTokens += update.Usage.InputTokenCount;
                    usage.CompletionTokens += update.Usage.OutputTokenCount;
                    usage.TotalTokens += update.Usage.TotalTokenCount;
                }

                if (update.FinishReason != null)
                {
                    finishReason = update.FinishReason;
                }

                foreach (var contentPart in update.ContentUpdate)
                {
                    if (string.IsNullOrEmpty(contentPart.Text))
                    {
                        continue;
                    }

                    contentBuilder.Append(contentPart.Text);

                    // 将取消令牌传播到用户代码中，以便调用方能快速停止
                    if (onDeltaContentAsync != null)
                    {
                        await onDeltaContentAsync(contentPart.Text, cancellationToken);
                    }
                }

                foreach (var toolUpdate in update.ToolCallUpdates)
                {
                    if (!toolCallAccumulators.TryGetValue(toolUpdate.Index, out var accumulator))
                    {
                        accumulator = new StreamingToolCallAccumulator { Index = toolUpdate.Index };
                        toolCallAccumulators.Add(toolUpdate.Index, accumulator);
                    }

                    if (!string.IsNullOrEmpty(toolUpdate.ToolCallId))
                    {
                        accumulator.ToolCallId = toolUpdate.ToolCallId;
                    }

                    if (!string.IsNullOrEmpty(toolUpdate.FunctionName))
                    {
                        accumulator.FunctionName = toolUpdate.FunctionName;
                    }

                    if (toolUpdate.FunctionArgumentsUpdate != null)
                    {
                        accumulator.Arguments.Append(toolUpdate.FunctionArgumentsUpdate);
                    }
                }
            }
        }
        catch (OperationCanceledException)
        {
            throw;
        }
        catch (Exception e) when (!receivedAnyUpdate)
        {
            // 首个数据帧之前失败（常见于服务端不支持流式工具调用），交由调用方回退非流式
            _logger.LogWarning(
                e,
                "MCP streaming completion failed before first update, retry with non-streaming"
            );
            return null;
        }

        var toolCalls = new List<ChatToolCall>(toolCallAccumulators.Count);
        foreach (var accumulator in toolCallAccumulators.Values)
        {
            if (string.IsNullOrEmpty(accumulator.FunctionName))
            {
                continue;
            }

            toolCalls.Add(
                ChatToolCall.CreateFunctionToolCall(
                    accumulator.ToolCallId,
                    accumulator.FunctionName,
                    BinaryData.FromString(accumulator.Arguments.ToString())
                )
            );
        }

        return new ChatRoundResult
        {
            Text = contentBuilder.ToString(),
            ToolCalls = toolCalls,
            FinishReason = finishReason,
        };
    }

    private static async Task<ChatCompletion> CompleteChatAndAccumulateUsageAsync(
        ChatClient chatClient,
        IReadOnlyList<ChatMessage> sendMessage,
        ChatCompletionOptions completionOptions,
        ChatUsageAccumulator usage,
        CancellationToken cancellationToken
    )
    {
        var completion = (
            await chatClient.CompleteChatAsync(sendMessage, completionOptions, cancellationToken)
        ).Value;

        if (!string.IsNullOrWhiteSpace(completion.Model))
        {
            usage.ResponseModel = completion.Model;
        }

        if (completion.Usage != null)
        {
            usage.PromptTokens += completion.Usage.InputTokenCount;
            usage.CompletionTokens += completion.Usage.OutputTokenCount;
            usage.TotalTokens += completion.Usage.TotalTokenCount;
        }

        return completion;
    }

    private sealed class ChatUsageAccumulator
    {
        public string ResponseModel { get; set; } = "";

        public int PromptTokens { get; set; }

        public int CompletionTokens { get; set; }

        public int TotalTokens { get; set; }
    }

    private sealed class ChatRoundResult
    {
        public string Text { get; set; } = "";

        public List<ChatToolCall> ToolCalls { get; set; } = [];

        public ChatFinishReason? FinishReason { get; set; }

        /// <summary>
        /// 是否由流式补全产生（false 表示发生过非流式回退）
        /// </summary>
        public bool WasStreamed { get; set; } = true;
    }

    private sealed class StreamingToolCallAccumulator
    {
        public int Index { get; init; }

        public string ToolCallId { get; set; } = "";

        public string FunctionName { get; set; } = "";

        public StringBuilder Arguments { get; } = new();
    }

    private static string GetCompletionText(ChatCompletion? completion)
    {
        if (completion is null || completion.Content.Count == 0)
        {
            return "";
        }

        var builder = new StringBuilder();
        foreach (var part in completion.Content)
        {
            if (!string.IsNullOrEmpty(part.Text))
            {
                builder.Append(part.Text);
            }
        }

        return builder.ToString();
    }

    private static OpenAiChatOptions LoadOptions(IConfiguration configuration)
    {
        var section = configuration.GetSection(ConfigurationSectionName);
        if (!section.Exists())
        {
            throw new InvalidConfigurationException(
                $"{ConfigurationSectionName} configuration section is missing"
            );
        }

        var options = section.Get<OpenAiChatOptions>() ?? new OpenAiChatOptions();
        if (string.IsNullOrWhiteSpace(options.Endpoint))
        {
            throw new InvalidConfigurationException("OpenAI:Endpoint configuration is missing");
        }

        if (string.IsNullOrWhiteSpace(options.ApiKey))
        {
            throw new InvalidConfigurationException("OpenAI:ApiKey configuration is missing");
        }

        if (string.IsNullOrWhiteSpace(options.Model))
        {
            throw new InvalidConfigurationException("OpenAI:Model configuration is missing");
        }

        return options;
    }

    private static ChatClient CreateChatClient(OpenAiChatOptions options)
    {
        return CreateChatClient(options, options.Model);
    }

    private static ChatClient CreateChatClient(OpenAiChatOptions options, string model)
    {
        var credential = new ApiKeyCredential(options.ApiKey);
        var clientOptions = new OpenAIClientOptions { Endpoint = new Uri(options.Endpoint) };
        return new ChatClient(model, credential: credential, clientOptions);
    }

    private static OpenAIModelClient CreateModelClient(OpenAiChatOptions options)
    {
        var credential = new ApiKeyCredential(options.ApiKey);
        var clientOptions = new OpenAIClientOptions { Endpoint = new Uri(options.Endpoint) };
        return new OpenAIModelClient(credential, clientOptions);
    }

    private OpenAiChatRequestOptions CreateRequestOptions()
    {
        return new OpenAiChatRequestOptions
        {
            Model = _defaults.Model,
            MaxOutputTokenCount = _defaults.MaxOutputTokenCount,
            Temperature = _defaults.Temperature,
            TopP = _defaults.TopP,
            FrequencyPenalty = _defaults.FrequencyPenalty,
        };
    }

    private static ChatCompletionOptions CreateCompletionOptions(OpenAiChatRequestOptions options)
    {
        return new ChatCompletionOptions
        {
            MaxOutputTokenCount = options.MaxOutputTokenCount,
            Temperature = options.Temperature,
            TopP = options.TopP,
            FrequencyPenalty = options.FrequencyPenalty,
            ResponseFormat = options.ResponseFormat,
        };
    }
}