using System.Buffers.Binary;
using System.Net;
using System.Text;

namespace Sunfish.Net;

public enum DnsType : ushort { A = 1, Ptr = 12, Txt = 16, Aaaa = 28, Srv = 33, Any = 255 }

/// <summary>Un registro de recurso DNS (solo los tipos que usa DNS-SD).</summary>
public sealed record DnsRecord(string Name, DnsType Type, uint Ttl)
{
    public string? Target { get; init; }                 // PTR, SRV
    public ushort Port { get; init; }                    // SRV
    public IPAddress? Address { get; init; }             // A, AAAA
    public IReadOnlyList<string> Text { get; init; } = []; // TXT
    public bool CacheFlush { get; init; }
}

/// <summary>
/// Mensaje DNS minimo para mDNS (RFC 6762) y DNS-SD (RFC 6763): preguntas y registros PTR, SRV,
/// TXT, A y AAAA. Al leer se entienden los nombres comprimidos; al escribir no se comprime.
/// </summary>
public sealed class DnsMessage
{
    public ushort Id { get; set; }
    public bool IsResponse { get; set; }
    public List<(string Name, DnsType Type, bool UnicastResponse)> Questions { get; } = new();
    public List<DnsRecord> Answers { get; } = new();
    public List<DnsRecord> Additional { get; } = new();

    public byte[] Encode()
    {
        var b = new List<byte>(512);
        void U16(int v) { b.Add((byte)(v >> 8)); b.Add((byte)v); }
        void U32(uint v) { U16((int)(v >> 16)); U16((int)(v & 0xFFFF)); }
        U16(Id);
        U16(IsResponse ? 0x8400 : 0); // QR + AA
        U16(Questions.Count);
        U16(Answers.Count);
        U16(0);
        U16(Additional.Count);
        foreach (var (name, type, unicast) in Questions)
        {
            WriteName(b, name);
            U16((ushort)type);
            U16(unicast ? 0x8001 : 0x0001);
        }
        foreach (var r in Answers.Concat(Additional))
        {
            WriteName(b, r.Name);
            U16((ushort)r.Type);
            U16(r.CacheFlush ? 0x8001 : 0x0001);
            U32(r.Ttl);
            var data = new List<byte>();
            switch (r.Type)
            {
                case DnsType.Ptr:
                    WriteName(data, r.Target!);
                    break;
                case DnsType.Srv:
                    data.AddRange([0, 0, 0, 0, (byte)(r.Port >> 8), (byte)r.Port]);
                    WriteName(data, r.Target!);
                    break;
                case DnsType.Txt:
                    if (r.Text.Count == 0)
                        data.Add(0);
                    foreach (var item in r.Text)
                    {
                        var bytes = Encoding.UTF8.GetBytes(item);
                        data.Add((byte)Math.Min(255, bytes.Length));
                        data.AddRange(bytes.Take(255));
                    }
                    break;
                case DnsType.A or DnsType.Aaaa:
                    data.AddRange(r.Address!.GetAddressBytes());
                    break;
            }
            U16(data.Count);
            b.AddRange(data);
        }
        return b.ToArray();
    }

    private static void WriteName(List<byte> b, string name)
    {
        foreach (var label in SplitLabels(name))
        {
            var bytes = Encoding.UTF8.GetBytes(label);
            b.Add((byte)Math.Min(63, bytes.Length));
            b.AddRange(bytes.Take(63));
        }
        b.Add(0);
    }

    /// <summary>
    /// "Sunfish (equipo)._http._tcp.local" -> etiquetas. La primera de un nombre de instancia puede
    /// contener puntos, asi que se separa por el tipo de servicio conocido.
    /// </summary>
    internal static List<string> SplitLabels(string name)
    {
        name = name.TrimEnd('.');
        int service = name.IndexOf("._", StringComparison.Ordinal);
        if (service > 0)
        {
            var labels = new List<string> { name[..service] };
            labels.AddRange(name[(service + 1)..].Split('.'));
            return labels;
        }
        return name.Split('.').ToList();
    }

    public static DnsMessage? Decode(ReadOnlySpan<byte> data)
    {
        try
        {
            var m = new DnsMessage
            {
                Id = BinaryPrimitives.ReadUInt16BigEndian(data),
                IsResponse = (data[2] & 0x80) != 0,
            };
            int qd = BinaryPrimitives.ReadUInt16BigEndian(data[4..]);
            int an = BinaryPrimitives.ReadUInt16BigEndian(data[6..]);
            int ns = BinaryPrimitives.ReadUInt16BigEndian(data[8..]);
            int ar = BinaryPrimitives.ReadUInt16BigEndian(data[10..]);
            int pos = 12;
            for (int i = 0; i < qd; i++)
            {
                string name = ReadName(data, ref pos);
                var type = (DnsType)BinaryPrimitives.ReadUInt16BigEndian(data[pos..]);
                int cls = BinaryPrimitives.ReadUInt16BigEndian(data[(pos + 2)..]);
                pos += 4;
                m.Questions.Add((name, type, (cls & 0x8000) != 0));
            }
            for (int i = 0; i < an + ns + ar; i++)
            {
                var record = ReadRecord(data, ref pos);
                if (record is null)
                    continue;
                if (i < an)
                    m.Answers.Add(record);
                else
                    m.Additional.Add(record);
            }
            return m;
        }
        catch (Exception e) when (e is ArgumentOutOfRangeException or IndexOutOfRangeException or InvalidDataException)
        {
            return null;
        }
    }

    private static DnsRecord? ReadRecord(ReadOnlySpan<byte> data, ref int pos)
    {
        string name = ReadName(data, ref pos);
        var type = (DnsType)BinaryPrimitives.ReadUInt16BigEndian(data[pos..]);
        int cls = BinaryPrimitives.ReadUInt16BigEndian(data[(pos + 2)..]);
        uint ttl = BinaryPrimitives.ReadUInt32BigEndian(data[(pos + 4)..]);
        int length = BinaryPrimitives.ReadUInt16BigEndian(data[(pos + 8)..]);
        pos += 10;
        int start = pos;
        int end = pos + length;
        if (end > data.Length)
            throw new InvalidDataException();
        pos = end;
        var r = new DnsRecord(name, type, ttl) { CacheFlush = (cls & 0x8000) != 0 };
        int p = start;
        switch (type)
        {
            case DnsType.Ptr:
                return r with { Target = ReadName(data, ref p) };
            case DnsType.Srv:
                p += 4;
                ushort port = BinaryPrimitives.ReadUInt16BigEndian(data[p..]);
                p += 2;
                return r with { Port = port, Target = ReadName(data, ref p) };
            case DnsType.Txt:
                var items = new List<string>();
                while (p < end)
                {
                    int len = data[p++];
                    if (len > 0)
                        items.Add(Encoding.UTF8.GetString(data.Slice(p, Math.Min(len, end - p))));
                    p += len;
                }
                return r with { Text = items };
            case DnsType.A when length == 4:
                return r with { Address = new IPAddress(data.Slice(start, 4)) };
            case DnsType.Aaaa when length == 16:
                return r with { Address = new IPAddress(data.Slice(start, 16)) };
            default:
                return null;
        }
    }

    private static string ReadName(ReadOnlySpan<byte> data, ref int pos)
    {
        var labels = new List<string>();
        int p = pos;
        bool jumped = false;
        for (int hops = 0; hops < 64; hops++)
        {
            int len = data[p];
            if (len == 0)
            {
                if (!jumped)
                    pos = p + 1;
                return string.Join('.', labels);
            }
            if ((len & 0xC0) == 0xC0)
            {
                int target = ((len & 0x3F) << 8) | data[p + 1];
                if (!jumped)
                    pos = p + 2;
                jumped = true;
                p = target;
                continue;
            }
            labels.Add(Encoding.UTF8.GetString(data.Slice(p + 1, len)));
            p += len + 1;
        }
        throw new InvalidDataException("Name loop");
    }
}