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();
}

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 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() => $"struct {Module}::{Name}";
}

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() => $"enum {Module}::{Name}";
}

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 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 NubTypePointer p:
                sb.Append("P(");
                EncodeType(sb, p.To);
                sb.Append(')');
                break;

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

            case NubTypeEnum st:
                sb.Append("E(");
                sb.Append(st.Module);
                sb.Append(':');
                sb.Append(st.Name);
                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;

            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,
            'P' => DecodePointer(),
            'F' => DecodeFunc(),
            'T' => DecodeStruct(),
            'E' => DecodeEnum(),
            _ => 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 NubTypeStruct DecodeStruct()
    {
        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);
    }

    private NubTypeEnum DecodeEnum()
    {
        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 NubTypeEnum.Get(module, name);
    }

    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();
    }
}