using System;
using System.Linq;
using System.Threading;
using FlashCap;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;
using SixLabors.ImageSharp.Processing;

namespace MMirror.Integrations.FaceCamera;

public sealed class FaceDetectionService : IDisposable
{
    private readonly ILogger<FaceDetectionService> logger;
    private readonly TimeProvider timeProvider;

    private readonly Size targetSize;
    private Size videoSize;

    private readonly float threshold;

    private readonly CaptureDeviceDescriptor videoDeviceDescriptor;
    private CaptureDevice? videoDevice;

    private readonly InferenceSession session;
    private readonly DenseTensor<float> tensor;

    private readonly TimeSpan interval;
    private DateTimeOffset nextCheckTime;

    private readonly Lock watchLock = new();
    private Action<bool>? faceDetected;
    public event Action<bool> FaceDetected
    {
        add
        {
            using var scope = watchLock.EnterScope();
            if (faceDetected is null)
            {
                var characteristics =
                    videoDeviceDescriptor
                        .Characteristics.Where(c =>
                            c.PixelFormat is PixelFormats.YUYV
                            && c.Width >= targetSize.Width
                            && c.Height >= targetSize.Height
                        )
                        .MinBy(c => c.Width * c.Height)
                    ?? throw new Exception(
                        $"No viable output format from camera available, options are {string.Join(", ", videoDeviceDescriptor.Characteristics.Select(c => $"{c.PixelFormat} {c.Width}x{c.Height}"))}."
                    );
                videoSize = new(characteristics.Width, characteristics.Height);
                videoDeviceDescriptor
                    .OpenAsync(characteristics, TranscodeFormats.DoNotTranscode, OnPixelBufferArrived)
                    .ContinueWith(t =>
                    {
                        videoDevice = t.Result;
                        videoDevice.StartAsync().ContinueWith(t => logger.LogInformation("Video capture has started"));
                    });
            }
            faceDetected += value;
        }
        remove
        {
            using var scope = watchLock.EnterScope();
            faceDetected -= value;
            if (faceDetected is null)
            {
                if (videoDevice is { } vd)
                {
                    vd.StopAsync()
                        .ContinueWith(_ =>
                        {
                            vd.Dispose();
                            logger.LogInformation("Video capture stopped");
                        });
                }
            }
        }
    }

    public FaceDetectionService(
        ILogger<FaceDetectionService> logger,
        IOptions<CameraOptions> options,
        TimeProvider timeProvider
    )
    {
        this.logger = logger;
        this.timeProvider = timeProvider;

        interval = options.Value.Interval;
        nextCheckTime = timeProvider.GetUtcNow();

        targetSize = new(options.Value.Width, options.Value.Height);
        threshold = options.Value.Threshold;

        var descriptors = new CaptureDevices().GetDescriptors();
        videoDeviceDescriptor =
            descriptors.FirstOrDefault(d => d.Identity as string == options.Value.Name)
            ?? throw new Exception(
                $"No video device named {options.Value.Name} available, options are {string.Join(", ", descriptors.Select(d => d.Identity))}."
            );

        session = new(options.Value.Model);
        tensor = new DenseTensor<float>([1, 3, targetSize.Height, targetSize.Width]);
    }

    private void OnPixelBufferArrived(PixelBufferScope bufferScope)
    {
        if (timeProvider.GetUtcNow() < nextCheckTime)
        {
            return;
        }

        var data = bufferScope.Buffer.ReferImage();
        using (var image = new Image<Rgb24>(videoSize.Width, videoSize.Height))
        {
            var pixels = new ArraySegment<byte>(data.Array!, data.Offset + 54, videoSize.Width * videoSize.Height * 2);
            image.ProcessPixelRows(accessor =>
            {
                for (var y = 0; y < videoSize.Height; y++)
                {
                    var row = accessor.GetRowSpan(y);
                    for (var x = 0; x < videoSize.Width; x += 2)
                    {
                        var y1 = pixels[y * videoSize.Width * 2 + x * 2];
                        var u = pixels[y * videoSize.Width * 2 + x * 2 + 1];
                        var y2 = pixels[y * videoSize.Width * 2 + x * 2 + 2];
                        var v = pixels[y * videoSize.Width * 2 + x * 2 + 3];

                        row[x] = new(
                            (byte)float.Clamp(y1 + 1.402f * (v - 128), 0.0f, 256.0f),
                            (byte)float.Clamp(y1 - 0.344136f * (u - 128) - 0.714136f * (v - 128), 0.0f, 256.0f),
                            (byte)float.Clamp(y1 + 1.772f * (u - 128), 0.0f, 256.0f)
                        );

                        row[x + 1] = new(
                            (byte)float.Clamp(y2 + 1.402f * (v - 128), 0.0f, 256.0f),
                            (byte)float.Clamp(y2 - 0.344136f * (u - 128) - 0.714136f * (v - 128), 0.0f, 256.0f),
                            (byte)float.Clamp(y2 + 1.772f * (u - 128), 0.0f, 256.0f)
                        );
                    }
                }
            });
            image.Mutate(o => o.Resize(options: new() { Mode = ResizeMode.Crop, Size = targetSize }));
            image.ProcessPixelRows(accessor =>
            {
                for (var y = 0; y < image.Height; y++)
                {
                    var row = accessor.GetRowSpan(y);
                    for (var x = 0; x < image.Width; x++)
                    {
                        tensor.SetValue(2 * image.Width * image.Height + y * image.Width + x, row[x].R);
                        tensor.SetValue(1 * image.Width * image.Height + y * image.Width + x, row[x].G);
                        tensor.SetValue(0 * image.Width * image.Height + y * image.Width + x, row[x].B);
                    }
                }
            });
        }

        using (var outputs = session.Run([NamedOnnxValue.CreateFromTensor("input", tensor)]))
        {
            faceDetected?.Invoke(DetectedFace(outputs));
        }

        nextCheckTime = timeProvider.GetUtcNow() + interval;
    }

    private bool DetectedFace(IDisposableReadOnlyCollection<DisposableNamedOnnxValue> outputs) =>
        DetectedFace(outputs, 8) || DetectedFace(outputs, 16) || DetectedFace(outputs, 32);

    private bool DetectedFace(IDisposableReadOnlyCollection<DisposableNamedOnnxValue> outputs, int stride)
    {
        var clsData = outputs.First(o => o.Name == $"cls_{stride}").AsTensor<float>();
        var objData = outputs.First(o => o.Name == $"obj_{stride}").AsTensor<float>();
        return clsData
            .Zip(objData, (cls, obj) => float.Sqrt(float.Clamp(cls, 0.0f, 1.0f) * float.Clamp(obj, 0.0f, 1.0f)))
            .Any(s => s > threshold);
    }

    public void Dispose()
    {
        videoDevice?.Dispose();
        session.Dispose();
    }
}