using NubLang.Ast;
using NubLang.Diagnostics;
using NubLang.Modules;
using NubLang.Syntax;
using OmniSharp.Extensions.LanguageServer.Protocol;

namespace NubLang.LSP;

public class WorkspaceManager(DiagnosticsPublisher diagnosticsPublisher)
{
    private readonly Dictionary<string, SyntaxTree> _syntaxTrees = new();
    private readonly Dictionary<string, List<TopLevelNode>> _compilationUnits = new();
    private ModuleRepository? _repository;

    public void Init(string rootPath)
    {
        var files = Directory.GetFiles(rootPath, "*.nub", SearchOption.AllDirectories);
        foreach (var path in files)
        {
            var text = File.ReadAllText(path);
            var tokenizer = new Tokenizer();

            tokenizer.Tokenize(path, text);
            diagnosticsPublisher.Publish(path, tokenizer.Diagnostics);

            var parser = new Parser();
            var parseResult = parser.Parse(tokenizer.Tokens);
            diagnosticsPublisher.Publish(path, parser.Diagnostics);

            _syntaxTrees[path] = parseResult;
        }

        ModuleRepository repository;

        foreach (var (fsPath, syntaxTree) in _syntaxTrees)
        {
            try
            {
                repository = ModuleRepository.Create(_syntaxTrees.Select(x => x.Value).ToList());
            }
            catch (CompileException e)
            {
                return;
            }

            var typeChecker = new TypeChecker();
            var result = typeChecker.Check(syntaxTree, repository);
            diagnosticsPublisher.Publish(fsPath, typeChecker.Diagnostics);
            _compilationUnits[fsPath] = result;
        }
    }

    public void UpdateFile(DocumentUri path)
    {
        var fsPath = path.GetFileSystemPath();

        var text = File.ReadAllText(fsPath);
        var tokenizer = new Tokenizer();
        tokenizer.Tokenize(fsPath, text);
        diagnosticsPublisher.Publish(path, tokenizer.Diagnostics);

        var parser = new Parser();
        var syntaxTree = parser.Parse(tokenizer.Tokens);
        diagnosticsPublisher.Publish(path, parser.Diagnostics);
        _syntaxTrees[fsPath] = syntaxTree;

        ModuleRepository repository;

        try
        {
            repository = ModuleRepository.Create(_syntaxTrees.Select(x => x.Value).ToList());
        }
        catch (CompileException e)
        {
            diagnosticsPublisher.Publish(path, [e.Diagnostic]);
            return;
        }

        var typeChecker = new TypeChecker();
        var result = typeChecker.Check(syntaxTree, repository);
        diagnosticsPublisher.Publish(fsPath, typeChecker.Diagnostics);
        _compilationUnits[fsPath] = result;
    }

    public List<TopLevelNode>? GetCompilationUnit(DocumentUri path)
    {
        return _compilationUnits.GetValueOrDefault(path.GetFileSystemPath());
    }

    public ModuleRepository? GetModuleRepository()
    {
        try
        {
            return ModuleRepository.Create(_syntaxTrees.Select(x => x.Value).ToList());
        }
        catch (CompileException e)
        {
            return null;
        }
    }
}