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 ILogger<FoundryLocalChatClient> logger;
    private readonly IServiceProvider serviceProvider;

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

    public FoundryLocalChatClient(
        ILogger<FoundryLocalChatClient> logger,
        IOptions<FoundryLocalOptions> options,
        IServiceProvider serviceProvider
    )
    {
        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" }, logger);

        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
    ) =>
        MapChatResponse(
            await (await chatClientTask).CompleteChatAsync(messages.Select(MapChatMessage), cancellationToken)
        );

    public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
        IEnumerable<ChatMessage> messages,
        ChatOptions? options = null,
        [EnumeratorCancellation] CancellationToken cancellationToken = default
    )
    {
        var chatClient = await chatClientTask;
        var stream = chatClient.CompleteChatStreamingAsync(messages.Select(MapChatMessage), cancellationToken);
        await foreach (var message in stream)
        {
            yield return MapChatResponseUpdate(message);
        }
    }

    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
        {
            await initializationCancellationTokenSource.CancelAsync();
        }
    }
}

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