using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;
using Betalgo.Ranul.OpenAI.ObjectModels.ResponseModels;
using Microsoft.AI.Foundry.Local;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;

namespace Reviewer.Cli.FoundryLocal;

internal sealed class FoundryLocalChatClient : IChatClient, IAsyncDisposable
{
    private readonly ILoggerFactory loggerFactory;
    private readonly ILogger<FoundryLocalChatClient> logger;
    private readonly IServiceProvider serviceProvider;

    private readonly CancellationTokenSource initializationCancellationTokenSource;
    private readonly Task<IModel> modelTask;
    private readonly Task<OpenAIChatClient> chatClientTask;

    public FoundryLocalChatClient(
        ILoggerFactory loggerFactory,
        ILogger<FoundryLocalChatClient> logger,
        IOptions<FoundryLocalOptions> options,
        IServiceProvider serviceProvider
    )
    {
        this.loggerFactory = loggerFactory;
        this.logger = logger;
        this.serviceProvider = serviceProvider;

        initializationCancellationTokenSource = new();
        modelTask = InitializeModel(options.Value, initializationCancellationTokenSource.Token);
        chatClientTask = InitializeChatClient(initializationCancellationTokenSource.Token);
    }

    private async Task<IModel> InitializeModel(FoundryLocalOptions options, CancellationToken cancellationToken)
    {
        await FoundryLocalManager.CreateAsync(
            new() { AppName = "ReviewCli" },
            loggerFactory.CreateLogger("FoundryLocal")
        );

        var result = await FoundryLocalManager.Instance.DownloadAndRegisterEpsAsync(cancellationToken);
        if (!result.Success)
        {
            logger.LogUnsuccessfulExecutionProviderRegistration(result.Status);
        }

        var catalog = await FoundryLocalManager.Instance.GetCatalogAsync(cancellationToken);
        var model =
            await catalog.GetModelAsync(options.Model, cancellationToken)
            ?? throw new ModelNotFoundException(options.Model);

        if (!await model.IsCachedAsync())
        {
            await model.DownloadAsync(logger.LogModelDownloadProgress, cancellationToken);
        }
        await model.LoadAsync(cancellationToken);
        return model;
    }

    private async Task<OpenAIChatClient> InitializeChatClient(CancellationToken cancellationToken)
    {
        var model = await modelTask;
        return await model.GetChatClientAsync(cancellationToken);
    }

    public async Task<ChatResponse> GetResponseAsync(
        IEnumerable<ChatMessage> messages,
        ChatOptions? options = null,
        CancellationToken cancellationToken = default
    )
    {
        var chatClient = await chatClientTask.WaitAsync(cancellationToken);
        return MapChatResponse(
            await chatClient.CompleteChatAsync(messages.Select(MapChatMessage), CancellationToken.None)
        );
        var mappedMessages = messages.Select(MapChatMessage).ToList();
        var promptChars = mappedMessages.Sum(m => m.Content?.Length ?? 0);
        logger.LogInformation(
            "Starting chat completion. Messages: {MessageCount}, prompt length: ~{PromptChars} chars. Caller token cancelled: {CallerCancelled}",
            mappedMessages.Count,
            promptChars,
            cancellationToken.IsCancellationRequested
        );

        try
        {
            // Pass CancellationToken.None on purpose: this completion should run to completion even if
            // the workflow's token is cancelled. That way, if this throws "cancelled", we know the
            // cancellation came from inside Foundry Local Core rather than from our token.
            return MapChatResponse(await chatClient.CompleteChatAsync(mappedMessages, CancellationToken.None));
        }
        catch (Exception ex)
        {
            logger.LogError(
                ex,
                "CompleteChatAsync failed ({ExceptionType}). Caller token cancelled: {CallerCancelled}, init token cancelled: {InitCancelled}",
                ex.GetType().FullName,
                cancellationToken.IsCancellationRequested,
                initializationCancellationTokenSource.IsCancellationRequested
            );
            throw;
        }
    }

    public IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
        IEnumerable<ChatMessage> messages,
        ChatOptions? options = null,
        CancellationToken cancellationToken = default
    )
    {
        // cancellationToken.Register(() =>
        // {
        //     logger.LogWarning("Outer token cancelled: {StackTrace}", new StackTrace(true));
        // });
        // initializationCancellationTokenSource.Token.Register(() =>
        // {
        //     logger.LogWarning("Inner token cancelled: {StackTrace}", new StackTrace(true));
        // });
        // var chatClient = await chatClientTask;
        // var stream = chatClient.CompleteChatStreamingAsync(messages.Select(MapChatMessage), cancellationToken);
        // await foreach (var message in stream)
        // {
        //     yield return MapChatResponseUpdate(message);
        // }
        throw new NotImplementedException();
    }

    private static Betalgo.Ranul.OpenAI.ObjectModels.RequestModels.ChatMessage MapChatMessage(
        ChatMessage chatMessage
    ) => new(role: chatMessage.Role.Value, content: chatMessage.Text ?? "", name: chatMessage.AuthorName);

    private static ChatResponse MapChatResponse(ChatCompletionCreateResponse chatCompletionCreateResponse) =>
        new(
            new ChatMessage(
                ChatRole.Assistant,
                chatCompletionCreateResponse
                    .Choices.Select(c => c?.Message?.ContentCalculated as string ?? c?.Message?.Content)
                    .OfType<string>()
                    .First()
            )
        )
        {
            RawRepresentation = chatCompletionCreateResponse,
        };

    private static ChatResponseUpdate MapChatResponseUpdate(
        ChatCompletionCreateResponse chatCompletionCreateResponse
    ) =>
        new(
            ChatRole.Assistant,
            chatCompletionCreateResponse
                .Choices.Select(c => c.Message?.ContentCalculated as string ?? c.Message?.Content)
                .OfType<string>()
                .First()
        );

    public object? GetService(Type serviceType, object? serviceKey = null) =>
        serviceKey is null
            ? serviceProvider.GetService(serviceType)
            : serviceProvider.GetKeyedService(serviceType, serviceKey);

    public void Dispose() => DisposeAsync().AsTask().GetAwaiter().GetResult();

    public async ValueTask DisposeAsync()
    {
        if (modelTask.IsCompletedSuccessfully)
        {
            var model = await modelTask;
            await model.UnloadAsync();
        }
        else
        {
            logger.LogWarning("Disposing foundry local client");
            await initializationCancellationTokenSource.CancelAsync();
        }
    }
}

file sealed class ModelNotFoundException(string modelAlias) : Exception($"Model \"{modelAlias}\" not found.");