using System.Buffers.Binary;
using System.IO.Compression;
using System.Text.Json;
using Fleet.Core;

namespace Monitor.Protocol;

/// <summary>
/// RawChunk 的编解码（Probe、Agent、Backend、离线工具共用同一份代码）。
/// 布局：<c>"MNRC" | u8 container_version | u32 crc32c(压缩数据) | Brotli( u32 header_len | header JSON | 各列数据 )</c>。
/// 列数据：可空列先写有效位图，再写有效值；plain 为无符号 varint（有符号类型用 zigzag），delta 为与上一个有效值之差的 zigzag varint。
/// </summary>
public static class RawChunkCodec
{
    private const int PrefixBytes = 4 + 1 + 4;

    private static ReadOnlySpan<byte> Magic => "MNRC"u8;

    /// <summary>压缩之前的正文：<c>u32 header_len | header JSON | 各列数据</c>。布局变化时黄金字节测试会失败，提醒升级版本。</summary>
    public static byte[] EncodeBody(RawChunk chunk)
    {
        RawChunkHeader header = chunk.Header;
        header.ContainerVersion = ProtocolVersion.Container;
        header.Streams = [];
        foreach (RawStream stream in chunk.Streams)
        {
            var streamHeader = new StreamHeader { Name = stream.Name, Count = stream.Count };
            foreach (RawColumn column in stream.Columns)
            {
                streamHeader.Columns.Add(new ColumnHeader
                {
                    Name = column.Name,
                    Type = TypeName(column.Type),
                    Enc = column.Enc == ColumnEnc.Delta ? "delta" : "plain",
                    Nullable = column.Nullable,
                    Role = RoleName(column.Role),
                });
            }

            header.Streams.Add(streamHeader);
        }

        byte[] headerJson = JsonSerializer.SerializeToUtf8Bytes(header, MonitorJsonContext.Default.RawChunkHeader);

        using var body = new MemoryStream();
        Span<byte> length = stackalloc byte[4];
        BinaryPrimitives.WriteUInt32LittleEndian(length, (uint)headerJson.Length);
        body.Write(length);
        body.Write(headerJson);
        foreach (RawStream stream in chunk.Streams)
        {
            foreach (RawColumn column in stream.Columns)
            {
                WriteColumn(body, column);
            }
        }

        return body.ToArray();
    }

    public static byte[] Encode(RawChunk chunk)
    {
        byte[] plain = EncodeBody(chunk);
        using var compressed = new MemoryStream();
        using (var brotli = new BrotliStream(compressed, CompressionLevel.Optimal, leaveOpen: true))
        {
            brotli.Write(plain);
        }

        byte[] payload = compressed.ToArray();
        byte[] result = new byte[PrefixBytes + payload.Length];
        Magic.CopyTo(result);
        result[4] = ProtocolVersion.Container;
        BinaryPrimitives.WriteUInt32LittleEndian(result.AsSpan(5), Crc32c.Compute(payload));
        payload.CopyTo(result, PrefixBytes);
        return result;
    }

    public static Result<RawChunk> Decode(ReadOnlySpan<byte> data)
    {
        if (data.Length < PrefixBytes || !data[..4].SequenceEqual(Magic))
        {
            return Corrupt("bad magic");
        }

        if (data[4] != ProtocolVersion.Container)
        {
            return Corrupt($"unsupported container version {data[4]}");
        }

        ReadOnlySpan<byte> payload = data[PrefixBytes..];
        if (BinaryPrimitives.ReadUInt32LittleEndian(data[5..]) != Crc32c.Compute(payload))
        {
            return Corrupt("crc mismatch");
        }

        byte[]? plain = Inflate(payload);
        if (plain is null)
        {
            return Corrupt("invalid or oversized compressed data");
        }

        if (plain.Length < 4)
        {
            return Corrupt("truncated header length");
        }

        uint headerLength = BinaryPrimitives.ReadUInt32LittleEndian(plain);
        if (headerLength > plain.Length - 4)
        {
            return Corrupt("header length exceeds data");
        }

        RawChunkHeader? header;
        try
        {
            header = JsonSerializer.Deserialize(plain.AsSpan(4, (int)headerLength), MonitorJsonContext.Default.RawChunkHeader);
        }
        catch (JsonException)
        {
            return Corrupt("header is not valid json");
        }

        if (header is null || header.Streams.Count > ProtocolLimits.MaxStreamsPerChunk)
        {
            return Corrupt("header is missing or has too many streams");
        }

        ReadOnlySpan<byte> columns = plain.AsSpan(4 + (int)headerLength);
        var streams = new List<RawStream>(header.Streams.Count);
        foreach (StreamHeader streamHeader in header.Streams)
        {
            if (streamHeader.Count is < 0 or > ProtocolLimits.MaxRecordsPerChunk || streamHeader.Columns.Count > ProtocolLimits.MaxColumnsPerStream)
            {
                return Corrupt($"stream {streamHeader.Name} has an invalid size");
            }

            var cols = new List<RawColumn>(streamHeader.Columns.Count);
            foreach (ColumnHeader columnHeader in streamHeader.Columns)
            {
                if (!TryParseType(columnHeader.Type, out ColumnType type) || !TryParseEnc(columnHeader.Enc, out ColumnEnc enc))
                {
                    return Corrupt($"column {columnHeader.Name} has an unknown type or encoding");
                }

                RawColumn? column = ReadColumn(ref columns, columnHeader, type, enc, streamHeader.Count);
                if (column is null)
                {
                    return Corrupt($"column {streamHeader.Name}.{columnHeader.Name} is truncated or out of range");
                }

                cols.Add(column);
            }

            streams.Add(new RawStream(streamHeader.Name, streamHeader.Count, cols));
        }

        return columns.Length == 0 ? Result<RawChunk>.Ok(new RawChunk(header, streams)) : Corrupt("trailing bytes");
    }

    private static Result<RawChunk> Corrupt(string message) =>
        Result<RawChunk>.Fail(MonitorErrors.RawCorrupt, message);

    private static byte[]? Inflate(ReadOnlySpan<byte> payload)
    {
        using var input = new MemoryStream(payload.ToArray());
        using var brotli = new BrotliStream(input, CompressionMode.Decompress);
        using var output = new MemoryStream();
        byte[] buffer = new byte[16 * 1024];
        try
        {
            int read;
            while ((read = brotli.Read(buffer, 0, buffer.Length)) > 0)
            {
                output.Write(buffer, 0, read);
                if (output.Length > ProtocolLimits.MaxDecodedChunkBytes)
                {
                    return null;
                }
            }
        }
        catch (InvalidDataException)
        {
            return null;
        }

        return output.ToArray();
    }

    private static void WriteColumn(MemoryStream body, RawColumn column)
    {
        int count = column.Count;
        if (column.Nullable)
        {
            byte[] bitmap = new byte[(count + 7) / 8];
            for (int i = 0; i < count; i++)
            {
                if (column.IsValid(i))
                {
                    bitmap[i >> 3] |= (byte)(1 << (i & 7));
                }
            }

            body.Write(bitmap);
        }

        bool signed = column.Type is ColumnType.I32 or ColumnType.I64;
        long previous = 0;
        for (int i = 0; i < count; i++)
        {
            if (!column.IsValid(i))
            {
                continue;
            }

            long value = column.Values[i];
            if (column.Enc == ColumnEnc.Delta)
            {
                VarInt.Write(body, VarInt.ZigZag(value - previous));
                previous = value;
            }
            else
            {
                VarInt.Write(body, signed ? VarInt.ZigZag(value) : (ulong)value);
            }
        }
    }

    private static RawColumn? ReadColumn(ref ReadOnlySpan<byte> data, ColumnHeader header, ColumnType type, ColumnEnc enc, int count)
    {
        bool[]? valid = null;
        if (header.Nullable)
        {
            int bitmapBytes = (count + 7) / 8;
            if (data.Length < bitmapBytes)
            {
                return null;
            }

            valid = new bool[count];
            for (int i = 0; i < count; i++)
            {
                valid[i] = (data[i >> 3] & (1 << (i & 7))) != 0;
            }

            data = data[bitmapBytes..];
        }

        bool signed = type is ColumnType.I32 or ColumnType.I64;
        long[] values = new long[count];
        long previous = 0;
        for (int i = 0; i < count; i++)
        {
            if (valid is not null && !valid[i])
            {
                continue;
            }

            if (!VarInt.TryRead(data, out ulong raw, out int consumed))
            {
                return null;
            }

            data = data[consumed..];
            long value;
            if (enc == ColumnEnc.Delta)
            {
                value = previous + VarInt.UnZigZag(raw);
                previous = value;
            }
            else
            {
                value = signed ? VarInt.UnZigZag(raw) : unchecked((long)raw);
            }

            if (!InRange(type, value))
            {
                return null;
            }

            values[i] = value;
        }

        return new RawColumn(header.Name, type, enc, header.Nullable, ParseRole(header.Role), values, valid);
    }

    private static bool InRange(ColumnType type, long value) => type switch
    {
        ColumnType.U8 => value is >= 0 and <= byte.MaxValue,
        ColumnType.U16 => value is >= 0 and <= ushort.MaxValue,
        ColumnType.U32 => value is >= 0 and <= uint.MaxValue,
        ColumnType.I32 => value is >= int.MinValue and <= int.MaxValue,
        _ => true,
    };

    public static string TypeName(ColumnType type) => type switch
    {
        ColumnType.U8 => "u8",
        ColumnType.U16 => "u16",
        ColumnType.U32 => "u32",
        ColumnType.U64 => "u64",
        ColumnType.I32 => "i32",
        _ => "i64",
    };

    private static bool TryParseType(string text, out ColumnType type)
    {
        switch (text)
        {
            case "u8": type = ColumnType.U8; return true;
            case "u16": type = ColumnType.U16; return true;
            case "u32": type = ColumnType.U32; return true;
            case "u64": type = ColumnType.U64; return true;
            case "i32": type = ColumnType.I32; return true;
            case "i64": type = ColumnType.I64; return true;
            default: type = default; return false;
        }
    }

    private static bool TryParseEnc(string text, out ColumnEnc enc)
    {
        switch (text)
        {
            case "plain": enc = ColumnEnc.Plain; return true;
            case "delta": enc = ColumnEnc.Delta; return true;
            default: enc = default; return false;
        }
    }

    private static string? RoleName(ColumnRole role) => role switch
    {
        ColumnRole.TsProbe => "ts_probe",
        ColumnRole.TsAgent => "ts_agent",
        ColumnRole.TsLocal => "ts_local",
        _ => null,
    };

    private static ColumnRole ParseRole(string? text) => text switch
    {
        "ts_probe" => ColumnRole.TsProbe,
        "ts_agent" => ColumnRole.TsAgent,
        "ts_local" => ColumnRole.TsLocal,
        _ => ColumnRole.None,
    };
}