package p.studio.compiler.services;

import lombok.extern.slf4j.Slf4j;
import p.studio.compiler.PBSDefinitions;
import p.studio.compiler.messages.BuildingIssueSink;
import p.studio.compiler.messages.FrontendPhaseContext;
import p.studio.compiler.models.IRBackend;
import p.studio.compiler.pbs.PbsFrontendCompiler;
import p.studio.compiler.pbs.PbsReservedMetadataExtractor;
import p.studio.compiler.pbs.stdlib.InterfaceModuleLoader;
import p.studio.compiler.pbs.stdlib.ResourceStdlibEnvironmentResolver;
import p.studio.compiler.pbs.stdlib.StdlibEnvironmentResolver;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.identifiers.ModuleId;
import p.studio.utilities.logs.LogAggregator;

import java.util.*;

@Slf4j
public class PBSFrontendPhaseService implements FrontendPhaseService {
    private final PbsFrontendCompiler frontendCompiler;
    private final PbsModuleAssemblyService moduleAssemblyService;
    private final PbsImportedSemanticContextService importedSemanticContextService;

    public PBSFrontendPhaseService() {
        this(new ResourceStdlibEnvironmentResolver(), new InterfaceModuleLoader());
    }

    PBSFrontendPhaseService(
            final StdlibEnvironmentResolver stdlibEnvironmentResolver,
            final InterfaceModuleLoader interfaceModuleLoader) {
        final var reservedMetadataExtractor = new PbsReservedMetadataExtractor();
        final var moduleVisibilityValidator = new p.studio.compiler.pbs.linking.PbsModuleVisibilityValidator();
        this.frontendCompiler = new PbsFrontendCompiler();
        this.moduleAssemblyService = new PbsModuleAssemblyService(
                stdlibEnvironmentResolver,
                interfaceModuleLoader,
                moduleVisibilityValidator);
        this.importedSemanticContextService = new PbsImportedSemanticContextService(
                reservedMetadataExtractor,
                frontendCompiler);
    }

    @Override
    public IRBackend compile(
            final FrontendPhaseContext ctx,
            final DiagnosticSink diagnostics,
            final LogAggregator logs,
            final BuildingIssueSink issues) {
        final var nameTable = ctx.nameTable();
        final var assembly = moduleAssemblyService.assemble(ctx, nameTable, diagnostics, issues);
        final var parsedSourceFiles = assembly.parsedSourceFiles();
        final var importedSemanticContexts = importedSemanticContextService.build(parsedSourceFiles, assembly.moduleTable());
        final var failedModuleIds = assembly.mutableFailedModuleIds();
        final var moduleDependencyGraph = assembly.moduleDependencyGraph();
        final var canonicalModulePool = assembly.canonicalModulePool();

        final var compiledSourceFiles = new ArrayList<CompiledSourceFile>(parsedSourceFiles.size());
        for (final var parsedSource : parsedSourceFiles) {
            final var blockedModuleIds = blockedModulesByDependency(failedModuleIds, moduleDependencyGraph);
            if (blockedModuleIds.contains(parsedSource.moduleId())) {
                continue;
            }
            final var importedSemanticContext = importedSemanticContexts.getOrDefault(
                    parsedSource.fileId(),
                    PbsImportedSemanticContext.empty());
            final var compileErrorBaseline = diagnostics.errorCount();
            final var irBackendFile = frontendCompiler.compileParsedFile(
                    parsedSource.fileId(),
                    parsedSource.ast(),
                    diagnostics,
                    parsedSource.sourceKind(),
                    parsedSource.moduleId(),
                    canonicalModulePool,
                    ctx.hostAdmissionContext(),
                    nameTable,
                    importedSemanticContext.supplementalTopDecls(),
                    importedSemanticContext.importedCallables(),
                    importedSemanticContext.importedReservedMetadata());
            if (diagnostics.errorCount() > compileErrorBaseline) {
                failedModuleIds.add(parsedSource.moduleId());
            }
            compiledSourceFiles.add(new CompiledSourceFile(parsedSource.moduleId(), irBackendFile));
        }

        final var irBackend = mergeCompiledSources(
                compiledSourceFiles,
                failedModuleIds,
                moduleDependencyGraph);
        logs.using(log).debug("PBS frontend lowered to IR BE:\n%s".formatted(irBackend));
        return irBackend;
    }

    private IRBackend mergeCompiledSources(
            final ArrayList<CompiledSourceFile> compiledSourceFiles,
            final Set<ModuleId> failedModuleIds,
            final Map<ModuleId, Set<ModuleId>> moduleDependencyGraph) {
        final var irBackendAggregator = IRBackend.aggregator();
        final var blockedModuleIds = blockedModulesByDependency(failedModuleIds, moduleDependencyGraph);
        final var entryPointCallableName = PBSDefinitions.PBS.getEntryPointCallableName();
        final var entryPointModuleCandidates = new LinkedHashSet<ModuleId>();

        for (final var compiledSource : compiledSourceFiles) {
            if (blockedModuleIds.contains(compiledSource.moduleId())) {
                continue;
            }
            irBackendAggregator.merge(compiledSource.irBackendFile());
            for (final var executable : compiledSource.irBackendFile().executableFunctions()) {
                if (!entryPointCallableName.equals(executable.callableName())) {
                    continue;
                }
                if (executable.moduleId() != null && !executable.moduleId().isNone()) {
                    entryPointModuleCandidates.add(executable.moduleId());
                }
            }
        }

        irBackendAggregator.entryPointCallableName(entryPointCallableName);
        if (entryPointModuleCandidates.size() == 1) {
            irBackendAggregator.entryPointModuleId(entryPointModuleCandidates.iterator().next());
        }
        return irBackendAggregator.emit();
    }

    private Set<ModuleId> blockedModulesByDependency(
            final Set<ModuleId> failedModuleIds,
            final Map<ModuleId, Set<ModuleId>> dependenciesByModule) {
        final Map<ModuleId, Set<ModuleId>> dependentsByModule = new HashMap<>();
        for (final var entry : dependenciesByModule.entrySet()) {
            final var importer = entry.getKey();
            for (final var dependency : entry.getValue()) {
                dependentsByModule.computeIfAbsent(dependency, ignored -> new HashSet<>()).add(importer);
            }
        }

        final var blocked = new HashSet<ModuleId>(failedModuleIds);
        final var pending = new ArrayDeque<ModuleId>(failedModuleIds);
        while (!pending.isEmpty()) {
            final var failedOrBlocked = pending.removeFirst();
            final var dependents = dependentsByModule.getOrDefault(failedOrBlocked, Set.of());
            for (final var dependent : dependents) {
                if (!blocked.add(dependent)) {
                    continue;
                }
                pending.addLast(dependent);
            }
        }
        return blocked;
    }

    private record CompiledSourceFile(
            ModuleId moduleId,
            p.studio.compiler.models.IRBackendFile irBackendFile) {
    }
}
