using System.Diagnostics;
using System.Text;

namespace Compiler;

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<Signature, NubTypeStruct> Cache = new();

    public static NubTypeStruct Get(bool packed, List<Field> fields)
    {
        var sig = new Signature(packed, fields);

        if (!Cache.TryGetValue(sig, out var sturctType))
            Cache[sig] = sturctType = new NubTypeStruct(packed, fields);

        return sturctType;
    }

    private NubTypeStruct(bool packed, List<Field> fields)
    {
        Packed = packed;
        this.fields = fields;
    }

    private NubTypeStruct(bool packed)
    {
        Packed = packed;
    }

    public bool Packed { get; }

    private IReadOnlyList<Field>? fields;
    public IReadOnlyList<Field> Fields
    {
        get
        {
            if (fields == null)
                throw new InvalidOperationException("Fields has not been set");

            return fields;
        }
    }

    public static NubTypeStruct CreateWithoutFields(bool packed)
    {
        return new NubTypeStruct(packed);
    }

    public void SetFields(List<Field> fields)
    {
        if (this.fields != null)
            throw new InvalidOperationException("Fields can only be set once");

        this.fields = fields;
        Cache[new Signature(Packed, this.fields)] = this;
    }

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

    public record Field(string Name, NubType Type);
    private record Signature(bool Packed, IReadOnlyList<Field> Fields);
}

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

    public string MangledName(string name, string module) => SymbolNameGen.Exported(name, module, this);

    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 Dictionary<NubTypeStruct, Definition> structDefinitions = new();

    private string EncodeRoot(NubType type)
    {
        var sb = new StringBuilder();

        EncodeType(sb, type);

        foreach (var definition in structDefinitions.Values.OrderBy(x => x.Index))
        {
            Debug.Assert(definition.Encoded != null);
            sb.Insert(0, definition.Encoded);
        }

        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(GetOrCreateStructDefinition(st));
                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 string GetOrCreateStructDefinition(NubTypeStruct st)
    {
        if (!structDefinitions.TryGetValue(st, out var definition))
        {
            definition = new Definition(structDefinitions.Count);
            structDefinitions[st] = definition;

            var sb = new StringBuilder();

            sb.Append("D(");
            sb.Append(st.Packed ? '1' : '0');
            sb.Append('{');
            for (var i = 0; i < st.Fields.Count; i++)
            {
                var field = st.Fields[i];
                sb.Append(field.Name);
                sb.Append(':');
                EncodeType(sb, field.Type);
            }
            sb.Append('}');
            sb.Append(')');

            var encoded = sb.ToString();
            definition.Encoded = encoded;
        }

        return $"T({definition.Index})";
    }

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

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

    private readonly string encoded;
    private int index;
    private readonly List<NubTypeStruct> structDefinitions = new();

    private NubType DecodeRoot()
    {
        // First pass: Collect all declarations
        while (TryExpect('D'))
        {
            Expect('(');
            var packedChar = Consume();
            if (packedChar is not '1' and not '0')
                throw new Exception("Expected '0' or '1' for struct packing");

            Expect('{');
            while (TryPeek(out var c))
            {
                Consume();
                if (c == '}')
                    break;
            }

            Expect(')');

            structDefinitions.Add(NubTypeStruct.CreateWithoutFields(packedChar == '1'));
        }

        index = 0;
        var defIndex = 0;

        // Second pass: Set field types
        while (TryExpect('D'))
        {
            var fields = new List<NubTypeStruct.Field>();

            Consume();
            Consume();
            Consume();

            while (!TryExpect('}'))
            {
                var sb = new StringBuilder();

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

                var type = DecodeType();

                fields.Add(new NubTypeStruct.Field(sb.ToString(), type));
            }
            Expect(')');

            structDefinitions[defIndex].SetFields(fields);

            defIndex += 1;
        }

        return DecodeType();
    }

    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(),
            _ => 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()
    {
        Expect('(');
        var index = ExpectInt();
        Expect(')');
        return structDefinitions[index];
    }

    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 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 SymbolNameGen
{
    public static string Exported(string module, string function, NubType type)
    {
        var canonical = TypeEncoder.Encode(type);
        var hash = Hashing.Fnv1a64(canonical);

        return $"nub_{Sanitize(module)}_{Sanitize(function)}_{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();
    }
}