using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization;

namespace Compiler;

[JsonConverter(typeof(NubTypeJsonConverter))]
public abstract class NubType
{
    public abstract override string ToString();

    [Obsolete("Use IsAssignableTo instead of ==", error: true)]
    public static bool operator ==(NubType? a, NubType? b) => throw new InvalidOperationException("Use IsAssignableTo");

    [Obsolete("Use IsAssignableTo instead of ==", error: true)]
    public static bool operator !=(NubType? a, NubType? b) => throw new InvalidOperationException("Use IsAssignableTo");

    public bool IsAssignableTo(NubType target)
    {
        return (this, target) switch
        {
            (NubTypeEnumVariant variant, NubTypeEnum targetEnum) => ReferenceEquals(variant.EnumType, targetEnum),
            _ => ReferenceEquals(this, target),
        };
    }
}

public class NubTypeVoid : NubType
{
    public static readonly NubTypeVoid Instance = new();

    private NubTypeVoid()
    {
    }

    public override string ToString() => "void";
}

public class NubTypeUInt : NubType
{
    private static readonly Dictionary<int, NubTypeUInt> Cache = new();

    public static NubTypeUInt Get(int width)
    {
        if (!Cache.TryGetValue(width, out var type))
            Cache[width] = type = new NubTypeUInt(width);

        return type;
    }

    public int Width { get; }

    private NubTypeUInt(int width)
    {
        Width = width;
    }

    public override string ToString() => $"u{Width}";
}

public class NubTypeSInt : NubType
{
    private static readonly Dictionary<int, NubTypeSInt> Cache = new();

    public static NubTypeSInt Get(int width)
    {
        if (!Cache.TryGetValue(width, out var type))
            Cache[width] = type = new NubTypeSInt(width);

        return type;
    }

    public int Width { get; }

    private NubTypeSInt(int width)
    {
        Width = width;
    }

    public override string ToString() => $"i{Width}";
}

public class NubTypeBool : NubType
{
    public static readonly NubTypeBool Instance = new();

    private NubTypeBool()
    {
    }

    public override string ToString() => "bool";
}

public class NubTypeString : NubType
{
    public static readonly NubTypeString Instance = new();

    private NubTypeString()
    {
    }

    public override string ToString() => "string";
}

public class NubTypeChar : NubType
{
    public static readonly NubTypeChar Instance = new();

    private NubTypeChar()
    {
    }

    public override string ToString() => "char";
}

public class NubTypeStruct : NubType
{
    private static readonly Dictionary<(string Module, string Name), NubTypeStruct> Cache = new();

    public static NubTypeStruct Get(string module, string name)
    {
        if (!Cache.TryGetValue((module, name), out var structType))
            Cache[(module, name)] = structType = new NubTypeStruct(module, name);

        return structType;
    }

    private NubTypeStruct(string module, string name)
    {
        Module = module;
        Name = name;
    }

    public string Module { get; }
    public string Name { get; }

    public override string ToString() => $"{Module}::{Name}";
}

public class NubTypeAnonymousStruct : NubType
{
    private static readonly Dictionary<Signature, NubTypeAnonymousStruct> Cache = new();

    public static NubTypeAnonymousStruct Get(List<Field> fields)
    {
        var sig = new Signature(fields);

        if (!Cache.TryGetValue(sig, out var func))
            Cache[sig] = func = new NubTypeAnonymousStruct(fields);

        return func;
    }

    private NubTypeAnonymousStruct(IReadOnlyList<Field> fields)
    {
        Fields = fields;
    }

    public IReadOnlyList<Field> Fields { get; }

    public override string ToString() => $"{{ {string.Join(", ", Fields.Select(x => $"{x.Name}: {x.Type}"))} }}";

    public class Field(string name, NubType type)
    {
        public string Name { get; } = name;
        public NubType Type { get; } = type;
    }

    private record Signature(IReadOnlyList<Field> Fields);
}

public class NubTypeEnum : NubType
{
    private static readonly Dictionary<(string Module, string Name), NubTypeEnum> Cache = new();

    public static NubTypeEnum Get(string module, string name)
    {
        if (!Cache.TryGetValue((module, name), out var enumType))
            Cache[(module, name)] = enumType = new NubTypeEnum(module, name);

        return enumType;
    }

    private NubTypeEnum(string module, string name)
    {
        Module = module;
        Name = name;
    }

    public string Module { get; }
    public string Name { get; }

    public override string ToString() => $"{Module}::{Name}";
}

public class NubTypeEnumVariant : NubType
{
    private static readonly Dictionary<(NubTypeEnum EnumType, string Variant), NubTypeEnumVariant> Cache = new();

    public static NubTypeEnumVariant Get(NubTypeEnum enumType, string variant)
    {
        if (!Cache.TryGetValue((enumType, variant), out var variantType))
            Cache[(enumType, variant)] = variantType = new NubTypeEnumVariant(enumType, variant);

        return variantType;
    }

    private NubTypeEnumVariant(NubTypeEnum enumType, string variant)
    {
        EnumType = enumType;
        Variant = variant;
    }

    public NubTypeEnum EnumType { get; }
    public string Variant { get; }

    public override string ToString() => $"{EnumType}.{Variant}";
}

public class NubTypePointer : NubType
{
    private static readonly Dictionary<NubType, NubTypePointer> Cache = new();

    public static NubTypePointer Get(NubType to)
    {
        if (!Cache.TryGetValue(to, out var ptr))
            Cache[to] = ptr = new NubTypePointer(to);

        return ptr;
    }

    public NubType To { get; }

    private NubTypePointer(NubType to)
    {
        To = to;
    }

    public override string ToString() => $"^{To}";
}

public class NubTypeFunc : NubType
{
    private static readonly Dictionary<Signature, NubTypeFunc> Cache = new();

    public static NubTypeFunc Get(List<NubType> parameters, NubType returnType)
    {
        var sig = new Signature(parameters, returnType);

        if (!Cache.TryGetValue(sig, out var func))
            Cache[sig] = func = new NubTypeFunc(parameters, returnType);

        return func;
    }

    public IReadOnlyList<NubType> Parameters { get; }
    public NubType ReturnType { get; }

    private NubTypeFunc(List<NubType> parameters, NubType returnType)
    {
        Parameters = parameters;
        ReturnType = returnType;
    }

    public override string ToString() => $"func({string.Join(' ', Parameters)}): {ReturnType}";

    private record Signature(IReadOnlyList<NubType> Parameters, NubType ReturnType);
}

public class NubTypeArray : NubType
{
    private static readonly Dictionary<NubType, NubTypeArray> Cache = new();

    public static NubTypeArray Get(NubType to)
    {
        if (!Cache.TryGetValue(to, out var ptr))
            Cache[to] = ptr = new NubTypeArray(to);

        return ptr;
    }

    public NubType ElementType { get; }

    private NubTypeArray(NubType elementType)
    {
        ElementType = elementType;
    }

    public override string ToString() => $"[]{ElementType}";
}

public class TypeEncoder
{
    public static string Encode(NubType type)
    {
        return new TypeEncoder().EncodeRoot(type);
    }

    private TypeEncoder()
    {
    }

    private string EncodeRoot(NubType type)
    {
        var sb = new StringBuilder();
        EncodeType(sb, type);
        return sb.ToString();
    }

    private void EncodeType(StringBuilder sb, NubType type)
    {
        switch (type)
        {
            case NubTypeVoid:
                sb.Append('V');
                break;

            case NubTypeBool:
                sb.Append('B');
                break;

            case NubTypeUInt u:
                sb.Append($"U(");
                sb.Append(u.Width);
                sb.Append(')');
                break;

            case NubTypeSInt s:
                sb.Append($"I(");
                sb.Append(s.Width);
                sb.Append(')');
                break;

            case NubTypeString:
                sb.Append('S');
                break;

            case NubTypeChar:
                sb.Append('C');
                break;

            case NubTypePointer p:
                sb.Append("P(");
                EncodeType(sb, p.To);
                sb.Append(')');
                break;

            case NubTypeStruct st:
                sb.Append("TN(");
                sb.Append(st.Module);
                sb.Append(':');
                sb.Append(st.Name);
                sb.Append(')');
                break;

            case NubTypeEnum e:
                sb.Append("EN(");
                sb.Append(e.Module);
                sb.Append(':');
                sb.Append(e.Name);
                sb.Append(')');
                break;

            case NubTypeEnumVariant ev:
                sb.Append("EV(");
                sb.Append(ev.EnumType.Module);
                sb.Append(':');
                sb.Append(ev.EnumType.Name);
                sb.Append(':');
                sb.Append(ev.Variant);
                sb.Append(')');
                break;

            case NubTypeFunc fn:
                sb.Append("F(");
                for (int i = 0; i < fn.Parameters.Count; i++)
                {
                    EncodeType(sb, fn.Parameters[i]);
                }
                EncodeType(sb, fn.ReturnType);
                sb.Append(')');
                break;

            case NubTypeAnonymousStruct s:
                sb.Append("TA(");
                foreach (var field in s.Fields)
                {
                    sb.Append(field.Name);
                    sb.Append(':');
                    EncodeType(sb, field.Type);
                }
                sb.Append(')');
                break;

            case NubTypeArray a:
                sb.Append("A(");
                EncodeType(sb, a.ElementType);
                sb.Append(')');
                break;

            default:
                throw new NotSupportedException(type.GetType().Name);
        }
    }

    private class Definition(int index)
    {
        public int Index { get; } = index;
        public string? Encoded { get; set; }
    }
}

public class TypeDecoder
{
    public static NubType Decode(string encoded)
    {
        return new TypeDecoder(encoded).DecodeType();
    }

    private TypeDecoder(string encoded)
    {
        this.encoded = encoded;
    }

    private readonly string encoded;
    private int index = 0;

    private NubType DecodeType()
    {
        var start = Consume();
        return start switch
        {
            'V' => NubTypeVoid.Instance,
            'B' => NubTypeBool.Instance,
            'U' => DecodeUInt(),
            'I' => DecodeSInt(),
            'S' => NubTypeString.Instance,
            'C' => NubTypeChar.Instance,
            'P' => DecodePointer(),
            'F' => DecodeFunc(),
            'T' => DecodeStruct(),
            'E' => DecodeEnum(),
            'A' => DecodeArray(),
            _ => throw new Exception($"'{start}' is not a valid start to a type")
        };
    }

    private NubTypeUInt DecodeUInt()
    {
        Expect('(');
        var width = ExpectInt();
        Expect(')');
        return NubTypeUInt.Get(width);
    }

    private NubTypeSInt DecodeSInt()
    {
        Expect('(');
        var width = ExpectInt();
        Expect(')');
        return NubTypeSInt.Get(width);
    }

    private NubTypePointer DecodePointer()
    {
        Expect('(');
        var to = DecodeType();
        Expect(')');
        return NubTypePointer.Get(to);
    }

    private NubTypeFunc DecodeFunc()
    {
        var types = new List<NubType>();

        Expect('(');
        while (!TryExpect(')'))
        {
            types.Add(DecodeType());
        }

        return NubTypeFunc.Get(types.Take(types.Count - 1).ToList(), types.Last());
    }

    private NubType DecodeStruct()
    {

        if (TryExpect('A'))
        {
            var sb = new StringBuilder();
            var fields = new List<NubTypeAnonymousStruct.Field>();

            Expect('(');
            while (!TryExpect(')'))
            {
                while (!TryExpect(':'))
                {
                    sb.Append(Consume());
                }

                var name = sb.ToString();
                sb.Clear();

                var type = DecodeType();

                fields.Add(new NubTypeAnonymousStruct.Field(name, type));
            }

            return NubTypeAnonymousStruct.Get(fields);
        }

        if (TryExpect('N'))
        {
            var sb = new StringBuilder();

            Expect('(');
            while (!TryExpect(':'))
                sb.Append(Consume());

            var module = sb.ToString();
            sb.Clear();

            while (!TryExpect(')'))
                sb.Append(Consume());

            var name = sb.ToString();

            return NubTypeStruct.Get(module, name);
        }

        throw new Exception("Expected 'A' or 'N'");
    }

    private NubType DecodeEnum()
    {
        var sb = new StringBuilder();

        if (TryExpect('V'))
        {
            Expect('(');
            while (!TryExpect(':'))
                sb.Append(Consume());

            var module = sb.ToString();
            sb.Clear();

            while (!TryExpect(':'))
                sb.Append(Consume());

            var name = sb.ToString();

            while (!TryExpect(')'))
                sb.Append(Consume());

            var variant = sb.ToString();

            return NubTypeEnumVariant.Get(NubTypeEnum.Get(module, name), variant);
        }
        else if (TryExpect('N'))
        {
            Expect('(');
            while (!TryExpect(':'))
                sb.Append(Consume());

            var module = sb.ToString();
            sb.Clear();

            while (!TryExpect(')'))
                sb.Append(Consume());

            var name = sb.ToString();

            return NubTypeEnum.Get(module, name);
        }

        throw new Exception($"Expected 'V' or 'N'");
    }

    private NubTypeArray DecodeArray()
    {
        Expect('(');
        var elementType = DecodeType();
        Expect(')');
        return NubTypeArray.Get(elementType);
    }

    private bool TryPeek(out char c)
    {
        if (index >= encoded.Length)
        {
            c = '\0';
            return false;
        }

        c = encoded[index];
        return true;
    }

    private bool TryConsume(out char c)
    {
        if (index >= encoded.Length)
        {
            c = '\0';
            return false;
        }

        c = encoded[index];
        index += 1;

        return true;
    }

    private char Consume()
    {
        if (!TryConsume(out var c))
            throw new Exception("Unexpected end of string");

        return c;
    }

    private bool TryExpect(char c)
    {
        if (index >= encoded.Length)
            return false;

        if (encoded[index] != c)
            return false;

        Consume();
        return true;
    }

    private void Expect(char c)
    {
        if (!TryExpect(c))
            throw new Exception($"Expected '{c}'");
    }

    private int ExpectInt()
    {
        var buf = string.Empty;

        while (TryPeek(out var c))
        {
            if (!char.IsDigit(c))
                break;

            buf += Consume();
        }

        return int.Parse(buf);
    }
}

public class NubTypeJsonConverter : JsonConverter<NubType>
{
    public override NubType Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
    {
        return TypeDecoder.Decode(reader.GetString()!);
    }

    public override void Write(Utf8JsonWriter writer, NubType value, JsonSerializerOptions options)
    {
        writer.WriteStringValue(TypeEncoder.Encode(value));
    }
}

public static class Hashing
{
    public static ulong Fnv1a64(string text)
    {
        const ulong offset = 14695981039346656037UL;
        const ulong prime = 1099511628211UL;

        ulong hash = offset;
        foreach (var c in Encoding.UTF8.GetBytes(text))
        {
            hash ^= c;
            hash *= prime;
        }

        return hash;
    }
}

public static class NameMangler
{
    public static string Mangle(string module, string name, NubType type)
    {
        var canonical = TypeEncoder.Encode(type);
        var hash = Hashing.Fnv1a64(canonical);

        return $"nub_{Sanitize(module)}_{Sanitize(name)}_{hash:x16}";
    }

    private static string Sanitize(string s)
    {
        var sb = new StringBuilder(s.Length);
        foreach (var c in s)
        {
            if (char.IsLetterOrDigit(c))
                sb.Append(c);
            else
                sb.Append('_');
        }
        return sb.ToString();
    }
}