using System.Text;
using Azaion.Common.Database;
using Azaion.Common.DTO.Config;
using Azaion.CommonSecurity;
using Azaion.CommonSecurity.DTO.Commands;
using MessagePack;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using NetMQ;
using NetMQ.Sockets;
using Newtonsoft.Json;

namespace Azaion.Common.Services;

public interface IInferenceService
{
    Task RunInference(List<string> mediaPaths, Func<AnnotationImage, Task> processAnnotation, CancellationToken ct = default);
}

public class PythonInferenceService(ILogger<PythonInferenceService> logger, IOptions<AIRecognitionConfig> aiConfigOptions) : IInferenceService
{
    public async Task RunInference(List<string> mediaPaths, Func<AnnotationImage, Task> processAnnotation, CancellationToken ct = default)
    {
        using var dealer = new DealerSocket();
        var clientId = Guid.NewGuid();
        dealer.Options.Identity = Encoding.UTF8.GetBytes(clientId.ToString("N"));
        dealer.Connect($"tcp://{SecurityConstants.ZMQ_HOST}:{SecurityConstants.ZMQ_PORT}");

        var data = MessagePackSerializer.Serialize(aiConfigOptions.Value);
        var filename = JsonConvert.SerializeObject(mediaPaths);
        dealer.SendFrame(MessagePackSerializer.Serialize(new RemoteCommand(CommandType.Inference, filename, data)));

        while (!ct.IsCancellationRequested)
        {
            try
            {
                var annotationStream = dealer.Get<AnnotationImage>(bytes => bytes.Length == 4 && Encoding.UTF8.GetString(bytes) == "DONE", ct: ct);
                if (annotationStream == null)
                    break;

                await processAnnotation(annotationStream);
            }
            catch (Exception e)
            {
                logger.LogError(e, e.Message);
                break;
            }
        }
    }
}