using NubLang.Ast;
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, CompilationUnit> _compilationUnits = new();

    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(path, text);

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

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

            _syntaxTrees[path] = parseResult;
        }

        foreach (var (fsPath, syntaxTree) in _syntaxTrees)
        {
            var modules = Module.Collect(_syntaxTrees.Select(x => x.Value).ToList());

            var typeChecker = new TypeChecker(syntaxTree, modules);
            var result = typeChecker.Check();
            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(fsPath, text);
        tokenizer.Tokenize();
        diagnosticsPublisher.Publish(path, tokenizer.Diagnostics);

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

        var modules = Module.Collect(_syntaxTrees.Select(x => x.Value).ToList());

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

    public void RemoveFile(DocumentUri path)
    {
        var fsPath = path.GetFileSystemPath();
        _syntaxTrees.Remove(fsPath);
        _compilationUnits.Remove(fsPath);
    }

    public CompilationUnit? GetCompilationUnit(DocumentUri path)
    {
        return _compilationUnits.GetValueOrDefault(path.GetFileSystemPath());
    }
}