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
{
    public string Name { get; }
    public string Module { get; }
    public bool Packed { get; }

    private IReadOnlyList<Field>? _resolvedFields;

    public IReadOnlyList<Field> Fields => _resolvedFields ?? throw new InvalidOperationException();

    public NubTypeStruct(string module, string name, bool packed)
    {
        Module = module;
        Name = name;
        Packed = packed;
    }

    public void ResolveFields(IReadOnlyList<Field> fields)
    {
        if (_resolvedFields != null)
            throw new InvalidOperationException($"{Name} already resolved");

        _resolvedFields = fields;
    }

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

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

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 readonly record struct Signature(IReadOnlyList<NubType> Parameters, NubType ReturnType);
}

static class TypeEncoder
{
    public static string Encode(NubType type)
    {
        var sb = new StringBuilder();
        Encode(type, sb);
        return sb.ToString();
    }

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

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

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

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

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

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

            case NubTypeStruct st:
                sb.Append("T(");
                sb.Append(st.Module).Append("::").Append(st.Name);
                sb.Append(',').Append(st.Packed ? '1' : '0');
                sb.Append('{');
                for (int i = 0; i < st.Fields.Count; i++)
                {
                    var field = st.Fields[i];
                    sb.Append(field.Name).Append(':');
                    Encode(field.Type, sb);
                    if (i < st.Fields.Count - 1)
                        sb.Append(',');
                }
                sb.Append("})");
                break;

            case NubTypeFunc fn:
                sb.Append("F(");
                foreach (var parameter in fn.Parameters)
                {
                    sb.Append(Encode(parameter));
                    sb.Append(',');
                }
                sb.Append(Encode(fn.ReturnType));
                sb.Append(')');
                break;

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

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

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

    private readonly string encoded;
    private int pos;

    private NubType Decode()
    {
        return Parse();
    }

    private NubType Parse()
    {
        if (pos >= encoded.Length)
            throw new InvalidOperationException("Unexpected end of string");

        char c = encoded[pos++];
        return c switch
        {
            'V' => NubTypeVoid.Instance,
            'B' => NubTypeBool.Instance,
            'S' => NubTypeString.Instance,
            'U' => ParseUInt(),
            'I' => ParseSInt(),
            'P' => ParsePointer(),
            'T' => ParseStruct(),
            'F' => ParseFunc(),
            _ => throw new NotSupportedException($"Unknown type code '{c}' at position {pos - 1}")
        };
    }

    private NubTypeUInt ParseUInt()
    {
        ExpectChar('(');
        int width = ReadNumber();
        ExpectChar(')');
        return NubTypeUInt.Get(width);
    }

    private NubTypeSInt ParseSInt()
    {
        ExpectChar('(');
        int width = ReadNumber();
        ExpectChar(')');
        return NubTypeSInt.Get(width);
    }

    private NubTypePointer ParsePointer()
    {
        ExpectChar('(');
        var to = Parse();
        ExpectChar(')');
        return NubTypePointer.Get(to);
    }

    private NubTypeStruct ParseStruct()
    {
        ExpectChar('(');
        int start = pos;
        while (pos < encoded.Length && encoded[pos] != ',' && encoded[pos] != '{') pos++;
        var fullName = encoded[start..pos];
        var parts = fullName.Split("::");
        if (parts.Length != 2)
            throw new InvalidOperationException($"Invalid struct name: {fullName}");

        string module = parts[0], name = parts[1];

        bool packed = false;
        if (encoded[pos] == ',')
        {
            pos += 1;
            packed = encoded[pos += 1] == '1';
        }

        var st = new NubTypeStruct(module, name, packed);

        ExpectChar('{');

        var fields = new List<NubTypeStruct.Field>();
        while (encoded[pos] != '}')
        {
            int nameStart = pos;
            while (encoded[pos] != ':') pos += 1;
            string fieldName = encoded[nameStart..pos];
            pos += 1;

            var fieldType = Parse();
            fields.Add(new NubTypeStruct.Field(fieldName, fieldType));

            if (encoded[pos] == ',') pos += 1;
        }

        ExpectChar('}');
        ExpectChar(')');

        st.ResolveFields(fields);
        return st;
    }

    private NubTypeFunc ParseFunc()
    {
        ExpectChar('(');
        var parameters = new List<NubType>();
        while (true)
        {
            if (encoded[pos] == ')')
            {
                pos++;
                break;
            }

            var param = Parse();
            parameters.Add(param);

            if (encoded[pos] == ',')
            {
                pos += 1;
            }
            else if (encoded[pos] == ')')
            {
                pos += 1;
                break;
            }
            else
            {
                throw new InvalidOperationException($"Unexpected char '{encoded[pos]}' in function type at {pos}");
            }
        }

        if (parameters.Count == 0)
            throw new InvalidOperationException("Function must have a return type");

        var returnType = parameters[^1];
        var paramTypes = parameters.Take(parameters.Count - 1).ToList();

        return NubTypeFunc.Get(paramTypes, returnType);
    }

    private void ExpectChar(char expected)
    {
        if (pos >= encoded.Length || encoded[pos] != expected)
            throw new InvalidOperationException($"Expected '{expected}' at position {pos}");

        pos += 1;
    }

    private int ReadNumber()
    {
        int start = pos;
        while (pos < encoded.Length && char.IsDigit(encoded[pos])) pos += 1;
        if (start == pos) throw new InvalidOperationException($"Expected number at position {start}");
        return int.Parse(encoded[start..pos]);
    }
}

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

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