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

namespace Azaion.Common.Services;

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

public class PythonInferenceService(ILogger<PythonInferenceService> logger, IOptions<AIRecognitionConfig> aiConfigOptions) : IInferenceService
{
    public async Task RunInference(string mediaPath, Func<AnnotationImage, CancellationToken, 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);
        dealer.SendFrame(MessagePackSerializer.Serialize(new RemoteCommand(CommandType.Inference, mediaPath, data)));

        while (true)
        {
            byte[] bytes = [];
            try
            {
                var annotationStream = dealer.Get<AnnotationImage>(out bytes);
                if (annotationStream == null)
                {
                    if (bytes.Length == 4 && Encoding.UTF8.GetString(bytes) == "DONE")
                        break;
                    continue;
                }

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