using NubLang.Syntax.Node;

namespace NubLang;

public sealed class DefinitionTable
{
    private readonly List<TopLevelNode> _topLevelNodes;

    public DefinitionTable(IEnumerable<SyntaxTree> syntaxTrees)
    {
        _topLevelNodes = syntaxTrees.SelectMany(x => x.TopLevelNodes).ToList();
    }

    public IEnumerable<LocalFuncNode> LookupLocalFunc(string @namespace, string name)
    {
        return _topLevelNodes
            .OfType<LocalFuncNode>()
            .Where(x => x.Namespace == @namespace && x.Name == name);
    }

    public IEnumerable<ExternFuncNode> LookupExternFunc(string @namespace, string name)
    {
        return _topLevelNodes
            .OfType<ExternFuncNode>()
            .Where(x => x.Namespace == @namespace && x.Name == name);
    }

    public IEnumerable<StructNode> LookupStruct(string @namespace, string name)
    {
        return _topLevelNodes
            .OfType<StructNode>()
            .Where(x => x.Namespace == @namespace && x.Name == name);
    }

    public IEnumerable<StructFieldNode> LookupStructField(StructNode structNode, string field)
    {
        return structNode.Fields.Where(x => x.Name == field);
    }

    public IEnumerable<TraitFuncImplNode> LookupTraitFuncImpl(NubType forType, string name)
    {
        return _topLevelNodes
            .OfType<TraitImplNode>()
            .Where(x => x.ForType == forType)
            .SelectMany(x => x.Functions)
            .Where(x => x.Name == name);
    }

    public IEnumerable<TraitNode> LookupTrait(string @namespace, string name)
    {
        return _topLevelNodes
            .OfType<TraitNode>()
            .Where(x => x.Namespace == @namespace && x.Name == name);
    }

    public IEnumerable<TraitFuncNode> LookupTraitFunc(TraitNode trait, string name)
    {
        return trait.Functions.Where(x => x.Name == name);
    }
}

public sealed class BoundDefinitionTable
{
    private readonly List<BoundTopLevelNode> _topLevelNodes;

    public BoundDefinitionTable(IEnumerable<BoundSyntaxTree> syntaxTrees)
    {
        _topLevelNodes = syntaxTrees.SelectMany(x => x.TopLevelNodes).ToList();
    }

    public BoundLocalFuncNode LookupLocalFunc(string @namespace, string name)
    {
        return _topLevelNodes
            .OfType<BoundLocalFuncNode>()
            .First(x => x.Namespace == @namespace && x.Name == name);
    }

    public BoundExternFuncNode LookupExternFunc(string @namespace, string name)
    {
        return _topLevelNodes
            .OfType<BoundExternFuncNode>()
            .First(x => x.Namespace == @namespace && x.Name == name);
    }

    public BoundStructNode LookupStruct(string @namespace, string name)
    {
        return _topLevelNodes
            .OfType<BoundStructNode>()
            .First(x => x.Namespace == @namespace && x.Name == name);
    }

    public BoundStructFieldNode LookupStructField(BoundStructNode structNode, string field)
    {
        return structNode.Fields.First(x => x.Name == field);
    }

    public IEnumerable<BoundTraitImplNode> LookupTraitImpls(BoundNubType itemType)
    {
        return _topLevelNodes
            .OfType<BoundTraitImplNode>()
            .Where(x => x.ForType == itemType);
    }

    public BoundTraitFuncImplNode LookupTraitFuncImpl(BoundNubType forType, string name)
    {
        return _topLevelNodes
            .OfType<BoundTraitImplNode>()
            .Where(x => x.ForType == forType)
            .SelectMany(x => x.Functions)
            .First(x => x.Name == name);
    }

    public BoundTraitNode LookupTrait(string @namespace, string name)
    {
        return _topLevelNodes
            .OfType<BoundTraitNode>()
            .First(x => x.Namespace == @namespace && x.Name == name);
    }

    public BoundTraitFuncNode LookupTraitFunc(BoundTraitNode trait, string name)
    {
        return trait.Functions.First(x => x.Name == name);
    }

    public IEnumerable<BoundStructNode> GetStructs()
    {
        return _topLevelNodes.OfType<BoundStructNode>();
    }

    public IEnumerable<BoundTraitNode> GetTraits()
    {
        return _topLevelNodes.OfType<BoundTraitNode>();
    }
}