// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

using System.Collections.Concurrent;
using System.Text.Json;
using Microsoft.Agents.A365.Observability.Runtime.Tracing.Contracts;
using Microsoft.Agents.A365.Observability.Runtime.Tracing.Scopes;
using Microsoft.Extensions.AI;
using Microsoft.Teams.Apps;
using Microsoft.Teams.Apps.Api.Clients;
using Microsoft.Teams.Apps.Handlers;
using Microsoft.Teams.Apps.Schema;
using Microsoft.Teams.Apps.Schema.Entities;
using Microsoft.Teams.Apps.State;

namespace ObservabilityBot;

public class ObservabilityBotApp : TeamsBotApplication
{
    private readonly IChatClient _chatClient;
    private readonly ChatOptions _chatOptions;
    private readonly ConcurrentDictionary<string, List<ChatMessage>> _chatHistories = new();
    private readonly string _deploymentName;

    public ObservabilityBotApp(
        ApiClient teamsApiClient,
        IHttpContextAccessor httpContextAccessor,
        ILogger<ObservabilityBotApp> logger,
        IChatClient chatClient,
        ChatOptions chatOptions,
        TeamsBotApplicationOptions? teamsOptions = null,
        TurnStateLoader? turnStateLoader = null)
        : base(teamsApiClient, httpContextAccessor, logger, teamsOptions, turnStateLoader)
    {
        _chatClient = chatClient;
        _chatOptions = chatOptions;
        _deploymentName = Environment.GetEnvironmentVariable("AZURE_OPENAI_DEPLOYMENT") ?? "unknown";

        this.OnMessage(HandleMessageAsync);
    }

    private async Task HandleMessageAsync(Context<MessageActivity> context, CancellationToken ct)
    {
        ArgumentNullException.ThrowIfNull(context.Activity);
        ArgumentNullException.ThrowIfNull(context.Activity.Conversation);
        ArgumentNullException.ThrowIfNull(context.Activity.Conversation.Id);

        await context.Typing(ct);

        string conversationId = context.Activity.Conversation.Id;
        List<ChatMessage> history = context.State.UserState?.Get<List<ChatMessage>>() ?? [];

        lock (history)
        {
            history.Add(new ChatMessage(ChatRole.User, context.Activity.Text));
        }

        // Build Agent365 scope contracts from the turn context.
        TeamsConversationAccount? recipient = context.Activity.Recipient;
        var agentDetails = new AgentDetails(
            agentId: recipient?.AgenticAppId ?? recipient?.Id,
            agentName: recipient?.Name,
            agenticUserId: recipient?.AgenticUserId,
            agentBlueprintId: recipient?.AgenticAppBlueprintId,
            tenantId: recipient?.TenantId);

        var request = new Request(
            content: context.Activity.Text,
            conversationId: conversationId,
            channel: new Channel(context.Activity.ChannelId));

        // === InvokeAgentScope: wraps the entire agent turn ===
        // Opened here (not by the SDK) so we can reliably record both input and output.
        // The SDK has already set cert-required baggage on Baggage.Current via TeamsBaggageBuilder,
        // so this span inherits tenant.id, agent.id, user.id etc. automatically.
        var invokeAgentScopeDetails = new InvokeAgentScopeDetails(context.Activity.ServiceUrl);
        using var invokeScope = InvokeAgentScope.Start(request, invokeAgentScopeDetails, agentDetails);

        try
        {
            // === InferenceScope: wraps the LLM + tool-call loop ===
            var inferenceDetails = new InferenceCallDetails(
                InferenceOperationType.Chat,
                model: _deploymentName,
                providerName: "AzureOpenAI");

            List<ChatMessage> snapshot;
            lock (history) { snapshot = [.. history]; }

            ChatResponse chatResponse;
            using (var inferenceScope = InferenceScope.Start(request, inferenceDetails, agentDetails))
            {
                chatResponse = await _chatClient.GetResponseAsync(snapshot, _chatOptions, ct);

                if (chatResponse.Usage is { } usage)
                {
                    if (usage.InputTokenCount is { } inputTokens)
                        inferenceScope.RecordInputTokens((int)inputTokens);
                    if (usage.OutputTokenCount is { } outputTokens)
                        inferenceScope.RecordOutputTokens((int)outputTokens);
                }

                string finishReason = chatResponse.FinishReason?.Value ?? "stop";
                inferenceScope.RecordFinishReasons([finishReason]);
            }

            lock (history)
            {
                history.AddRange(chatResponse.Messages);
            }

            // === ExecuteToolScope: record each tool invocation ===
            var toolCalls = chatResponse.Messages
                .SelectMany(m => m.Contents.OfType<FunctionCallContent>())
                .GroupBy(fc => fc.CallId ?? fc.Name ?? "")
                .ToDictionary(g => g.Key, g => g.First());

            foreach (FunctionResultContent? funcResult in chatResponse.Messages
                .SelectMany(m => m.Contents.OfType<FunctionResultContent>()))
            {
                toolCalls.TryGetValue(funcResult.CallId ?? "", out FunctionCallContent? matchingCall);

                var toolDetails = new ToolCallDetails(
                    toolName: matchingCall?.Name ?? "unknown",
                    arguments: matchingCall?.Arguments is { } args ? JsonSerializer.Serialize(args) : null,
                    toolCallId: funcResult.CallId);

                using var toolScope = ExecuteToolScope.Start(request, toolDetails, agentDetails);
                if (funcResult.Result is not null)
                {
                    toolScope.RecordResponse(funcResult.Result.ToString()!);
                }
            }

            // Extract citations from tool results.
            var citations = chatResponse.Messages
                .SelectMany(m => m.Contents.OfType<FunctionResultContent>())
                .Where(frc => frc.Result is not null)
                .SelectMany(frc =>
                {
                    try
                    {
                        JsonElement json = JsonSerializer.Deserialize<JsonElement>(frc.Result!.ToString()!);
                        if (json.TryGetProperty("structuredContent", out JsonElement sc) &&
                            sc.TryGetProperty("results", out JsonElement results))
                        {
                            return results.EnumerateArray()
                                .Where(r => r.TryGetProperty("contentUrl", out _))
                                .Select(r => (
                                    Title: r.GetProperty("title").GetString() ?? "",
                                    Url: r.GetProperty("contentUrl").GetString() ?? "",
                                    Content: r.TryGetProperty("content", out JsonElement c) ? c.GetString() ?? "" : ""
                                ));
                        }
                    }
                    catch (JsonException) { }
                    return [];
                })
                .DistinctBy(c => c.Url)
                .Take(5).ToList();

            string responseText = chatResponse.Text;

            for (int i = 0; i < citations.Count; i++)
            {
                responseText += $"[{i + 1}] ";
            }

            // Record output on the top-level invoke_agent span before it closes.
            invokeScope.RecordOutputMessages([responseText]);

            TeamsActivityBuilder builder = TeamsActivity.CreateBuilder()
                .WithText(responseText, TextFormats.Markdown)
                .AddMention(context.Activity?.From!)
                .AddAIGenerated();

            for (int i = 0; i < citations.Count; i++)
            {
                (string? Title, string? Url, string? Content) = citations[i];
                string abstract_ = Content.Length > 160 ? Content[..157] + "..." : Content;
                builder.AddCitation(i + 1, new CitationAppearance() { Name = Title, Url = new Uri(Url), Abstract = abstract_, Icon = CitationIcon.Text });
            }

            await context.Send(builder.Build(), ct);
        }
        catch (Exception ex)
        {
            invokeScope.RecordError(ex);
            throw;
        }
        finally
        {
            context.State.UserState?.Set(history);
        }
    }
}
