package p.studio.compiler.pbs.semantics;

import p.studio.compiler.pbs.PbsFrontendCompiler;
import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.source.Span;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.identifiers.ModuleId;
import p.studio.utilities.structures.DependencyGraphAnaliser;

import java.util.*;

final class PbsGlobalSemanticsValidator {

    Analysis analyse(
            final PbsAst.File ast,
            final ModuleId currentModuleId,
            final p.studio.utilities.structures.ReadOnlyList<PbsFrontendCompiler.ImportedGlobalSurface> importedGlobals) {
        final var normalizedModuleId = currentModuleId == null ? ModuleId.none() : currentModuleId;
        final var localDeclsByName = collectLocalGlobalDecls(ast);
        final var localGlobalNames = new LinkedHashSet<>(localDeclsByName.keySet());
        final var importedDefinitionsByRef = new LinkedHashMap<PbsGlobalDependencySupport.GlobalRef, PbsFrontendCompiler.ImportedGlobalSurface>();
        final var directlyVisibleImports = new LinkedHashMap<String, PbsGlobalDependencySupport.GlobalRef>();

        for (final var importedGlobal : importedGlobals) {
            final var ref = new PbsGlobalDependencySupport.GlobalRef(importedGlobal.ownerModuleId(), importedGlobal.globalName());
            importedDefinitionsByRef.putIfAbsent(ref, importedGlobal);
            if (importedGlobal.directlyVisible()) {
                directlyVisibleImports.putIfAbsent(importedGlobal.localName(), ref);
            }
        }

        final var globalNamesByModule = new LinkedHashMap<ModuleId, LinkedHashSet<String>>();
        globalNamesByModule.put(normalizedModuleId, localGlobalNames);
        for (final var ref : importedDefinitionsByRef.keySet()) {
            globalNamesByModule.computeIfAbsent(ref.moduleId(), ignored -> new LinkedHashSet<>()).add(ref.globalName());
        }

        final var graph = new LinkedHashMap<PbsGlobalDependencySupport.GlobalRef, Set<PbsGlobalDependencySupport.GlobalRef>>();
        final var spansByRef = new LinkedHashMap<PbsGlobalDependencySupport.GlobalRef, Span>();

        for (final var entry : localDeclsByName.entrySet()) {
            final var ref = new PbsGlobalDependencySupport.GlobalRef(normalizedModuleId, entry.getKey());
            spansByRef.put(ref, entry.getValue().span());
            graph.put(ref, PbsGlobalDependencySupport.collectDependencies(
                    entry.getValue().initializer(),
                    normalizedModuleId,
                    localGlobalNames,
                    directlyVisibleImports));
        }

        for (final var entry : importedDefinitionsByRef.entrySet()) {
            final var surface = entry.getValue();
            spansByRef.putIfAbsent(entry.getKey(), surface.span());
            final var importedByLocalName = new LinkedHashMap<String, PbsGlobalDependencySupport.GlobalRef>();
            for (final var importedDependency : surface.importedDependencies()) {
                importedByLocalName.putIfAbsent(
                        importedDependency.localName(),
                        new PbsGlobalDependencySupport.GlobalRef(
                                importedDependency.ownerModuleId(),
                                importedDependency.globalName()));
            }
            graph.putIfAbsent(entry.getKey(), PbsGlobalDependencySupport.collectDependencies(
                    surface.initializer(),
                    surface.ownerModuleId(),
                    globalNamesByModule.getOrDefault(surface.ownerModuleId(), new LinkedHashSet<>()),
                    importedByLocalName));
        }

        final var analysis = new DependencyGraphAnaliser<PbsGlobalDependencySupport.GlobalRef>(
                Comparator.comparingInt((PbsGlobalDependencySupport.GlobalRef ref) -> ref.moduleId().isNone() ? Integer.MAX_VALUE : ref.moduleId().getIndex())
                        .thenComparing(PbsGlobalDependencySupport.GlobalRef::globalName))
                .analyse(graph);
        return new Analysis(analysis, spansByRef);
    }

    void validate(
            final PbsAst.File ast,
            final ModuleId currentModuleId,
            final p.studio.utilities.structures.ReadOnlyList<PbsFrontendCompiler.ImportedGlobalSurface> importedGlobals,
            final DiagnosticSink diagnostics) {
        final var analysis = analyse(ast, currentModuleId, importedGlobals);
        for (final var cycle : analysis.graphAnalysis().cycleComponents()) {
            for (final var ref : cycle) {
                final var span = analysis.spansByRef().get(ref);
                if (span == null) {
                    continue;
                }
                p.studio.compiler.source.diagnostics.Diagnostics.error(
                        diagnostics,
                        PbsSemanticsErrors.E_SEM_GLOBAL_CYCLIC_DEPENDENCY.name(),
                        "Cyclic global dependency detected involving '%s'".formatted(ref.globalName()),
                        span);
            }
        }
    }

    private LinkedHashMap<String, PbsAst.GlobalDecl> collectLocalGlobalDecls(final PbsAst.File ast) {
        final var globals = new LinkedHashMap<String, PbsAst.GlobalDecl>();
        for (final var topDecl : ast.topDecls()) {
            if (topDecl instanceof PbsAst.GlobalDecl globalDecl) {
                globals.putIfAbsent(globalDecl.name(), globalDecl);
            }
        }
        return globals;
    }

    record Analysis(
            DependencyGraphAnaliser.Analysis<PbsGlobalDependencySupport.GlobalRef> graphAnalysis,
            Map<PbsGlobalDependencySupport.GlobalRef, Span> spansByRef) {
    }
}
