using System.Net;
using DotNext;
using Common;

namespace Peripherals.JpegClient;

static class JpegAddr
{
    const UInt32 BASE = 0x0000_0000;

    public const UInt32 CAPTURE_RD_CTRL = BASE + 0x0;
    public const UInt32 CAPTURE_WR_CTRL = BASE + 0x1;

    public const UInt32 START_WR_ADDR0 = BASE + 0x2;
    public const UInt32 END_WR_ADDR0 = BASE + 0x3;
    public const UInt32 START_WR_ADDR1 = BASE + 0x4;
    public const UInt32 END_WR_ADDR1 = BASE + 0x5;
    public const UInt32 START_RD_ADDR0 = BASE + 0x6;
    public const UInt32 END_RD_ADDR0 = BASE + 0x7;

    public const UInt32 HDMI_NOT_READY = BASE + 0x8;
    public const UInt32 HDMI_HEIGHT_WIDTH = BASE + 0x9;

    public const UInt32 JPEG_HEIGHT_WIDTH = BASE + 0xA;
    public const UInt32 JPEG_ADD_NEED_FRAME_NUM = BASE + 0xB;
    public const UInt32 JPEG_FRAME_SAVE_NUM = BASE + 0xC;
    public const UInt32 JPEG_FIFO_FRAME_INFO = BASE + 0xD;

    public const UInt32 ADDR_HDMI_WD_START = 0x4000_0000;
    public const UInt32 ADDR_JPEG_START = 0x8000_0000;
    public const UInt32 ADDR_JPEG_END = 0xA000_0000;
}

public class JpegInfo
{
    public UInt32 Width { get; set; }
    public UInt32 Height { get; set; }
    public UInt32 Size { get; set; }

    public JpegInfo(UInt32 width, UInt32 height, UInt32 size)
    {
        Width = width;
        Height = height;
        Size = size;
    }

    public JpegInfo(byte[] data)
    {
        if (data.Length < 8)
            throw new ArgumentException("Invalid data length", nameof(data));

        Width = ((UInt32)(data[5] << 8 + data[6] & 0xF0));
        Height = ((UInt32)((data[6] & 0x0F) << 4 + data[7]));
        Size = Number.BytesToUInt32(data, 0, 4).Value;
    }
}

public enum JpegSampleRate : UInt32
{
    RATE_1_1 = 0b1111_1111_1111_1111_1111_1111_1111_1111,
    RATE_1_2 = 0b1010_1010_1010_1010_1010_1010_1010_1010,
    RATE_1_4 = 0b1000_1000_1000_1000_1000_1000_1000_1000,
    RATE_3_4 = 0b1110_1110_1110_1110_1110_1110_1110_1110,
    RATE_1_8 = 0b1000_0000_1000_0000_1000_0000_1000_0000,
    RATE_3_8 = 0b1001_0010_0100_1001_1001_0010_0100_1001,
    RATE_7_8 = 0b1111_1110_1111_1110_1111_1110_1111_1110,
    RATE_1_16 = 0b1000_0000_0000_0000_1000_0000_0000_0000,
    RATE_3_16 = 0b1000_0100_0010_0000_1000_0100_0010_0000,
    RATE_5_16 = 0b1001_0001_0010_0010_0100_0100_1000_1001,
    RATE_15_16 = 0b1111_1111_1111_1110_1111_1111_1111_1110,
    RATE_1_32 = 0b1000_0000_0000_0000_0000_0000_0000_0000,
    RATE_31_32 = 0b1111_1111_1111_1111_1111_1111_1111_1110,
}

public class Jpeg
{
    private static readonly NLog.Logger logger = NLog.LogManager.GetCurrentClassLogger();

    readonly int timeout = 2000;
    readonly int taskID;
    readonly int port;
    readonly string address;
    private IPEndPoint ep;

    public Jpeg(string address, int port, int taskID, int timeout = 2000)
    {
        if (timeout < 0)
            throw new ArgumentException("Timeout couldn't be negative", nameof(timeout));
        this.address = address;
        this.taskID = taskID;
        this.port = port;
        this.ep = new IPEndPoint(IPAddress.Parse(address), port);
        this.timeout = timeout;
    }

    public async ValueTask<Result<bool>> Init(bool enable = true)
    {
        {
            var ret = await CheckHdmiIsReady();
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to check HDMI ready: {ret.Error}");
                return new(ret.Error);
            }
            if (!ret.Value)
            {
                logger.Error("HDMI not ready");
                return new(false);
            }
        }

        int width = -1, height = -1;
        {
            var ret = await GetHdmiResolution();
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to get HDMI resolution: {ret.Error}");
                return new(ret.Error);
            }
            (width, height) = ret.Value;
        }

        {
            var ret = await ConnectJpeg2Hdmi(width, height);
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to connect JPEG to HDMI: {ret.Error}");
                return new(ret.Error);
            }
            if (!ret.Value)
            {
                logger.Error("Failed to connect JPEG to HDMI");
                return false;
            }
        }

        if (enable)
            return await SetEnable(true);
        else return true;
    }

    public async ValueTask<bool> SetEnable(bool enable)
    {
        if (enable)
        {
            var ret = await UDPClientPool.WriteAddrSeq(
                this.ep,
                this.taskID,
                [JpegAddr.CAPTURE_RD_CTRL, JpegAddr.CAPTURE_WR_CTRL],
                [0b11, 0b01],
                this.timeout
            );
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to set JPEG enable: {ret.Error}");
                return false;
            }
            return ret.Value;
        }
        else
        {
            var ret = await UDPClientPool.WriteAddrSeq(
                this.ep,
                this.taskID,
                [JpegAddr.CAPTURE_RD_CTRL, JpegAddr.CAPTURE_WR_CTRL],
                [0b00, 0b00],
                this.timeout
            );
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to set JPEG disable: {ret.Error}");
                return false;
            }
            return ret.Value;
        }
    }

    public async ValueTask<Result<bool>> CheckHdmiIsReady()
    {
        var ret = await UDPClientPool.ReadAddrWithWait(
            this.ep, this.taskID, JpegAddr.HDMI_NOT_READY, 0b01, 0b01, 100, this.timeout);
        if (!ret.IsSuccessful)
        {
            logger.Error($"Failed to check HDMI status: {ret.Error}");
            return new(ret.Error);
        }
        return ret.Value;
    }

    public async ValueTask<Result<(int, int)>> GetHdmiResolution()
    {
        var ret = await UDPClientPool.ReadAddr(
            this.ep, this.taskID, JpegAddr.HDMI_HEIGHT_WIDTH, 0, this.timeout);
        if (!ret.IsSuccessful)
        {
            logger.Error($"Failed to get HDMI resolution: {ret.Error}");
            return new(ret.Error);
        }

        var data = ret.Value.Options.Data;
        if (data == null || data.Length != 4)
        {
            logger.Error($"Invalid HDMI resolution data length: {data?.Length ?? 0}");
            return new(new Exception("Invalid HDMI resolution data length"));
        }

        var width = data[0] | (data[1] << 8);
        var height = data[2] | (data[3] << 8);
        return new((width, height));
    }

    public async ValueTask<Result<bool>> ConnectJpeg2Hdmi(int width, int height)
    {
        if (width <= 0 || height <= 0)
        {
            logger.Error($"Invalid HDMI resolution: {width}x{height}");
            return new(new ArgumentException("Invalid HDMI resolution"));
        }

        var frameSize = (UInt32)(width * height / 4);

        {
            var ret = await UDPClientPool.WriteAddr(
                this.ep, this.taskID, JpegAddr.START_WR_ADDR0, JpegAddr.ADDR_HDMI_WD_START, this.timeout);
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to set HDMI output start address: {ret.Error}");
                return new(ret.Error);
            }
            if (!ret.Value)
            {
                logger.Error($"Failed to set HDMI output start address");
                return false;
            }
        }

        {
            var ret = await UDPClientPool.WriteAddr(
                this.ep, this.taskID, JpegAddr.END_WR_ADDR0,
                JpegAddr.ADDR_HDMI_WD_START + frameSize, this.timeout);
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to set HDMI output end address: {ret.Error}");
                return new(ret.Error);
            }
            if (!ret.Value)
            {
                logger.Error($"Failed to set HDMI output address");
                return false;
            }
        }

        {
            var ret = await UDPClientPool.WriteAddr(
                this.ep, this.taskID, JpegAddr.START_RD_ADDR0, JpegAddr.ADDR_HDMI_WD_START, this.timeout);
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to set jpeg input start address: {ret.Error}");
                return new(ret.Error);
            }
            if (!ret.Value)
            {
                logger.Error($"Failed to set jpeg input address");
                return false;
            }
        }

        {
            var ret = await UDPClientPool.WriteAddr(
                this.ep, this.taskID, JpegAddr.END_RD_ADDR0,
                JpegAddr.ADDR_HDMI_WD_START + frameSize, this.timeout);
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to set jpeg input end address: {ret.Error}");
                return new(ret.Error);
            }
            if (!ret.Value)
            {
                logger.Error($"Failed to set jpeg input end address");
                return false;
            }
        }

        {
            var ret = await UDPClientPool.WriteAddr(
                this.ep, this.taskID, JpegAddr.START_WR_ADDR1, JpegAddr.ADDR_JPEG_START, this.timeout);
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to set jpeg output start address: {ret.Error}");
                return new(ret.Error);
            }
            if (!ret.Value)
            {
                logger.Error($"Failed to set jpeg output start address");
                return false;
            }
        }

        {
            var ret = await UDPClientPool.WriteAddr(
                this.ep, this.taskID, JpegAddr.END_WR_ADDR1, JpegAddr.ADDR_JPEG_END, this.timeout);
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to set jpeg output end address: {ret.Error}");
                return new(ret.Error);
            }
            if (!ret.Value)
            {
                logger.Error($"Failed to set jpeg output end address");
                return false;
            }
        }

        return true;
    }

    // public async ValueTask<bool> SetSampleRate(uint rate)
    // {
    //     var ret = await UDPClientPool.WriteAddr(
    //         this.ep, this.taskID, JpegAddr.FRAME_SAMPLE_RATE, rate, this.timeout);
    //     if (!ret.IsSuccessful)
    //     {
    //         logger.Error($"Failed to set JPEG sample rate: {ret.Error}");
    //         return false;
    //     }
    //     return ret.Value;
    // }

    // public async ValueTask<bool> SetSampleRate(JpegSampleRate rate)
    // {
    //     return await SetSampleRate((uint)rate);
    // }

    public async ValueTask<uint> GetFrameNumber()
    {
        var ret = await UDPClientPool.ReadAddrByte(
            this.ep, this.taskID, JpegAddr.JPEG_FRAME_SAVE_NUM, this.timeout);
        if (!ret.IsSuccessful)
        {
            logger.Error($"Failed to get JPEG frame number: {ret.Error}");
            return 0;
        }
        return Number.BytesToUInt32(ret.Value.Options.Data ?? Array.Empty<byte>()).Value;
    }

    public async ValueTask<Optional<List<JpegInfo>>> GetFrameInfo(int num)
    {
        var ret = await UDPClientPool.ReadAddr(this.ep, this.taskID, JpegAddr.JPEG_FIFO_FRAME_INFO, num, this.timeout);
        if (!ret.IsSuccessful)
        {
            logger.Error($"Failed to get JPEG frame info: {ret.Error}");
            return new(null);
        }

        var data = ret.Value.Options.Data;
        if (data == null || data.Length == 0)
        {
            logger.Error($"Data is null or empty");
            return new(null);
        }
        if (data.Length != num * 2)
        {
            logger.Error(
                $"Data length should be {num * 2} bytes, instead of {data.Length} bytes");
            return new(null);
        }

        var infos = new List<JpegInfo>();
        for (int i = 0; i < num; i++)
        {
            infos.Add(new JpegInfo(data[i..(i + 1)]));
        }
        return new(infos);
    }

    public async ValueTask<bool> AddFrameNum2Process(uint cnt)
    {
        var ret = await UDPClientPool.WriteAddr(
            this.ep, this.taskID, JpegAddr.JPEG_ADD_NEED_FRAME_NUM, cnt, this.timeout);
        if (!ret.IsSuccessful)
        {
            logger.Error($"Failed to update pointer: {ret.Error}");
            return false;
        }
        return ret.Value;
    }

    public async ValueTask<Result<byte[]?>> GetFrame(uint offset, uint length)
    {
        if (!MsgBus.IsRunning)
        {
            logger.Error("Message bus is not running");
            return new(new Exception("Message bus is not running"));
        }
        MsgBus.UDPServer.ClearUDPData(this.ep.Address.ToString(), this.ep.Port);

        var firstReadLength = (int)(Math.Min(
            length,
            JpegAddr.ADDR_JPEG_END - JpegAddr.ADDR_JPEG_START - offset
        ));
        var secondReadLength = (int)(length - firstReadLength);
        var dataBytes = new byte[length];

        {
            var ret = await UDPClientPool.ReadAddr4Bytes(
                this.ep, this.taskID, JpegAddr.ADDR_JPEG_START + offset, firstReadLength, this.timeout);
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to get JPEG frame data: {ret.Error}");
                return null;
            }
            if (ret.Value.Length != firstReadLength)
            {
                logger.Error($"Data length should be {firstReadLength} bytes, instead of {ret.Value.Length} bytes");
                return null;
            }
            Buffer.BlockCopy(ret.Value, 0, dataBytes, 0, firstReadLength);
        }

        if (secondReadLength > 0)
        {
            var ret = await UDPClientPool.ReadAddr4Bytes(
                this.ep, this.taskID, JpegAddr.ADDR_JPEG_START, secondReadLength, this.timeout);
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to get JPEG frame data: {ret.Error}");
                return null;
            }
            if (ret.Value.Length != secondReadLength)
            {
                logger.Error($"Data length should be {secondReadLength} bytes, instead of {ret.Value.Length} bytes");
                return null;
            }
            Buffer.BlockCopy(ret.Value, 0, dataBytes, firstReadLength, secondReadLength);
        }

        return dataBytes;
    }

    public async ValueTask<List<byte[]>> GetMultiFrames(uint offset, uint[] sizes)
    {
        var frames = new List<byte[]>();
        for (int i = 0; i < sizes.Length; i++)
        {
            var ret = await GetFrame(offset, sizes[i]);
            if (!ret.IsSuccessful)
            {
                logger.Error($"Failed to get JPEG frame {i} data: {ret.Error}");
                continue;
            }
            if (ret.Value == null)
            {
                logger.Error($"Frame {i} data is null");
                continue;
            }
            if (ret.Value.Length != sizes[i])
            {
                logger.Error(
                    $"Frame {i} data length should be {sizes[i]} bytes, instead of {ret.Value.Length} bytes");
                continue;
            }

            frames.Add(ret.Value);
            offset += sizes[i];
        }

        {
            var ret = await AddFrameNum2Process((uint)sizes.Length);
            if (!ret) logger.Error($"Failed to update pointer");
        }

        return frames;
    }

    public async ValueTask<Result<List<byte[]>?>> GetMultiFrames(uint offset)
    {
        if (!MsgBus.IsRunning)
        {
            logger.Error("Message bus is not running");
            return new(new Exception("Message bus is not running"));
        }
        MsgBus.UDPServer.ClearUDPData(this.ep.Address.ToString(), this.ep.Port);

        var frameNum = await GetFrameNumber();
        if (frameNum == 0) return null;

        List<uint>? frameSizes = null;
        {
            var ret = await GetFrameInfo((int)frameNum);
            if (!ret.HasValue || ret.Value.Count == 0)
            {
                logger.Error($"Failed to get frame info");
                return null;
            }
            frameSizes = ret.Value.Select(x => x.Size).ToList();
        }

        var frames = await GetMultiFrames(offset, frameSizes.ToArray());
        if (frames.Count == 0)
        {
            logger.Error($"Failed to get frames");
            return null;
        }

        return frames;
    }
}
