using NubLang.Diagnostics;
using NubLang.Syntax;

namespace NubLang.Ast;

public class TypeResolver
{
    private readonly Dictionary<string, Module> _modules;
    private readonly Dictionary<(string Module, string Name), NubType> _typeCache = new();
    private readonly HashSet<(string Module, string Name)> _resolvingTypes = [];

    public TypeResolver(Dictionary<string, Module> modules)
    {
        _modules = modules;
    }

    public NubType ResolveType(TypeSyntax type, string currentModule)
    {
        return type switch
        {
            ArrayTypeSyntax arr => new NubArrayType(ResolveType(arr.BaseType, currentModule)),
            BoolTypeSyntax => new NubBoolType(),
            IntTypeSyntax i => new NubIntType(i.Signed, i.Width),
            FloatTypeSyntax f => new NubFloatType(f.Width),
            FuncTypeSyntax func => new NubFuncType(func.Parameters.Select(x => ResolveType(x, currentModule)).ToList(), ResolveType(func.ReturnType, currentModule)),
            SliceTypeSyntax slice => new NubSliceType(ResolveType(slice.BaseType, currentModule)),
            ConstArrayTypeSyntax arr => new NubConstArrayType(ResolveType(arr.BaseType, currentModule), arr.Size),
            PointerTypeSyntax ptr => new NubPointerType(ResolveType(ptr.BaseType, currentModule)),
            StringTypeSyntax => new NubStringType(),
            CustomTypeSyntax c => ResolveCustomType(c, currentModule),
            VoidTypeSyntax => new NubVoidType(),
            _ => throw new NotSupportedException($"Unknown type syntax: {type}")
        };
    }

    private NubType ResolveCustomType(CustomTypeSyntax customType, string currentModule)
    {
        var module = _modules[customType.ModuleToken?.Value ?? currentModule];

        var enumDef = module.Enums(true).FirstOrDefault(x => x.NameToken.Value == customType.NameToken.Value);
        if (enumDef != null)
        {
            return enumDef.Type != null ? ResolveType(enumDef.Type, currentModule) : new NubIntType(false, 64);
        }

        var structDef = module.Structs(true).FirstOrDefault(x => x.NameToken.Value == customType.NameToken.Value);
        if (structDef != null)
        {
            var key = (customType.ModuleToken?.Value ?? currentModule, customType.NameToken.Value);

            if (_typeCache.TryGetValue(key, out var cachedType))
            {
                return cachedType;
            }

            if (!_resolvingTypes.Add(key))
            {
                var placeholder = new NubStructType(customType.ModuleToken?.Value ?? currentModule, customType.NameToken.Value, structDef.Packed, []);
                _typeCache[key] = placeholder;
                return placeholder;
            }

            try
            {
                var result = new NubStructType(customType.ModuleToken?.Value ?? currentModule, structDef.NameToken.Value, structDef.Packed, []);
                _typeCache[key] = result;

                var fields = structDef.Fields
                    .Select(x => new NubStructFieldType(x.NameToken.Value, ResolveType(x.Type, currentModule), x.Value != null))
                    .ToList();

                result.Fields.AddRange(fields);
                return result;
            }
            finally
            {
                _resolvingTypes.Remove(key);
            }
        }

        throw new TypeResolverException(Diagnostic
            .Error($"Type {customType.NameToken.Value} not found in module {customType.ModuleToken?.Value ?? currentModule}")
            .At(customType)
            .Build());
    }
}

public class TypeResolverException : Exception
{
    public Diagnostic Diagnostic { get; }

    public TypeResolverException(Diagnostic diagnostic) : base(diagnostic.Message)
    {
        Diagnostic = diagnostic;
    }
}