package p.studio.compiler.services;

import p.studio.compiler.messages.BuildingIssueSink;
import p.studio.compiler.messages.FrontendPhaseContext;
import p.studio.compiler.models.IRGlobalVisibility;
import p.studio.compiler.models.ProjectDescriptor;
import p.studio.compiler.models.SourceHandle;
import p.studio.compiler.models.SourceKind;
import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.lexer.PbsLexer;
import p.studio.compiler.pbs.linking.PbsLinkErrors;
import p.studio.compiler.pbs.linking.PbsModuleVisibilityValidator;
import p.studio.compiler.pbs.parser.PbsBarrelParser;
import p.studio.compiler.pbs.parser.PbsParser;
import p.studio.compiler.pbs.stdlib.InterfaceModuleLoader;
import p.studio.compiler.pbs.stdlib.StdlibEnvironmentResolver;
import p.studio.compiler.source.Span;
import p.studio.compiler.source.diagnostics.DiagnosticPhase;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.identifiers.FileId;
import p.studio.compiler.source.identifiers.ModuleId;
import p.studio.compiler.source.identifiers.ProjectId;
import p.studio.compiler.source.tables.FileTable;
import p.studio.compiler.source.tables.ModuleReference;
import p.studio.compiler.source.tables.ModuleTable;
import p.studio.compiler.source.tables.NameTable;
import p.studio.utilities.structures.ReadOnlyList;

import java.nio.charset.StandardCharsets;
import java.nio.file.Path;
import java.util.*;

final class PbsModuleAssemblyService {
    private final StdlibEnvironmentResolver stdlibEnvironmentResolver;
    private final InterfaceModuleLoader interfaceModuleLoader;
    private final PbsModuleVisibilityValidator moduleVisibilityValidator;

    PbsModuleAssemblyService(
            final StdlibEnvironmentResolver stdlibEnvironmentResolver,
            final InterfaceModuleLoader interfaceModuleLoader,
            final PbsModuleVisibilityValidator moduleVisibilityValidator) {
        this.stdlibEnvironmentResolver = stdlibEnvironmentResolver;
        this.interfaceModuleLoader = interfaceModuleLoader;
        this.moduleVisibilityValidator = moduleVisibilityValidator;
    }

    PbsModuleAssembly assemble(
            final FrontendPhaseContext ctx,
            final NameTable nameTable,
            final DiagnosticSink diagnostics,
            final BuildingIssueSink issues) {
        if (!(ctx.fileTable instanceof FileTable writableFileTable)) {
            throw new IllegalStateException("PBS frontend requires writable FileTable for synthetic stdlib sources");
        }

        final var parsedSourceFiles = new ArrayList<PbsParsedSourceFile>();
        final Map<PbsModuleVisibilityValidator.ModuleCoordinates, MutableModuleUnit> modulesByCoordinates = new LinkedHashMap<>();
        final var moduleTable = new ModuleTable();
        final var moduleIdByFile = new HashMap<FileId, ModuleId>();
        final var failedModuleIds = new HashSet<ModuleId>();
        final var projectIdByName = new HashMap<String, ProjectId>();
        final var defaultSyntheticOwnerProjectId = ctx.stack.reverseTopologicalOrder.isEmpty()
                ? null
                : ctx.stack.reverseTopologicalOrder.getFirst();

        for (final var projectId : ctx.stack.reverseTopologicalOrder) {
            final var projectDescriptor = ctx.projectTable.get(projectId);
            final var projectSourceKind = ctx.sourceKind(projectId);
            final var fileIds = ctx.fileTable.getFiles(projectId);
            projectIdByName.putIfAbsent(projectDescriptor.getName(), projectId);

            for (final var fileId : fileIds) {
                final var sourceHandle = ctx.fileTable.get(fileId);
                sourceHandle.readUtf8().ifPresentOrElse(
                        utf8Content -> consumeSourceFile(
                                projectDescriptor,
                                projectSourceKind,
                                sourceHandle,
                                fileId,
                                utf8Content,
                                diagnostics,
                                modulesByCoordinates,
                                moduleTable,
                                moduleIdByFile,
                                failedModuleIds,
                                parsedSourceFiles),
                        () -> issues.report(builder -> builder
                                .error(true)
                                .message("Failed to read file content: %s".formatted(sourceHandle.toString()))));
            }
        }

        loadReservedStdlibModules(
                modulesByCoordinates,
                parsedSourceFiles,
                moduleIdByFile,
                moduleTable,
                writableFileTable,
                projectIdByName,
                defaultSyntheticOwnerProjectId,
                diagnostics,
                ctx.stdlibVersion());

        final var modules = new ArrayList<PbsModuleVisibilityValidator.ModuleUnit>(modulesByCoordinates.size());
        for (final var entry : modulesByCoordinates.entrySet()) {
            modules.add(new PbsModuleVisibilityValidator.ModuleUnit(
                    entry.getKey(),
                    ReadOnlyList.wrap(entry.getValue().sources),
                    ReadOnlyList.wrap(entry.getValue().barrels)));
        }

        moduleVisibilityValidator.validate(ReadOnlyList.wrap(modules), nameTable, diagnostics);
        markModulesWithLinkingErrors(diagnostics, moduleIdByFile, failedModuleIds);

        return new PbsModuleAssembly(
                ReadOnlyList.wrap(parsedSourceFiles),
                buildModuleDependencyGraph(parsedSourceFiles, moduleTable),
                emitModulePool(moduleTable),
                failedModuleIds,
                moduleTable,
                collectGlobalVisibilityByModule(modulesByCoordinates, moduleTable));
    }

    private Map<ModuleId, Map<String, IRGlobalVisibility>> collectGlobalVisibilityByModule(
            final Map<PbsModuleVisibilityValidator.ModuleCoordinates, MutableModuleUnit> modulesByCoordinates,
            final ModuleTable moduleTable) {
        final var visibilityByModule = new LinkedHashMap<ModuleId, Map<String, IRGlobalVisibility>>();
        for (final var entry : modulesByCoordinates.entrySet()) {
            final var moduleId = moduleId(moduleTable, entry.getKey());
            final var visibilityByGlobal = new LinkedHashMap<String, IRGlobalVisibility>();
            for (final var barrelFile : entry.getValue().barrels) {
                for (final var item : barrelFile.ast().items()) {
                    if (!(item instanceof PbsAst.BarrelGlobalItem globalItem)) {
                        continue;
                    }
                    visibilityByGlobal.put(
                            globalItem.name(),
                            globalItem.visibility() == PbsAst.Visibility.PUB
                                    ? IRGlobalVisibility.PUBLIC
                                    : IRGlobalVisibility.MODULE);
                }
            }
            if (!visibilityByGlobal.isEmpty()) {
                visibilityByModule.put(moduleId, Map.copyOf(visibilityByGlobal));
            }
        }
        return Map.copyOf(visibilityByModule);
    }

    private void consumeSourceFile(
            final ProjectDescriptor projectDescriptor,
            final SourceKind projectSourceKind,
            final SourceHandle sourceHandle,
            final FileId fileId,
            final String utf8Content,
            final DiagnosticSink diagnostics,
            final Map<PbsModuleVisibilityValidator.ModuleCoordinates, MutableModuleUnit> modulesByCoordinates,
            final ModuleTable moduleTable,
            final Map<FileId, ModuleId> moduleIdByFile,
            final Set<ModuleId> failedModuleIds,
            final ArrayList<PbsParsedSourceFile> parsedSourceFiles) {
        final var coordinates = resolveModuleCoordinates(projectDescriptor, sourceHandle);
        final var moduleId = moduleId(moduleTable, coordinates);
        final var moduleUnit = modulesByCoordinates.computeIfAbsent(coordinates, ignored -> new MutableModuleUnit());

        switch (sourceHandle.getExtension()) {
            case "pbs" -> {
                moduleIdByFile.put(fileId, moduleId);
                final var parseErrorBaseline = diagnostics.errorCount();
                final var ast = parseSourceFile(fileId, utf8Content, diagnostics, projectSourceKind);
                moduleUnit.sources.add(new PbsModuleVisibilityValidator.SourceFile(fileId, ast));
                parsedSourceFiles.add(new PbsParsedSourceFile(fileId, ast, moduleId, projectSourceKind));
                if (diagnostics.errorCount() > parseErrorBaseline) {
                    failedModuleIds.add(moduleId);
                }
            }
            case "barrel" -> {
                moduleIdByFile.put(fileId, moduleId);
                if ("mod.barrel".equals(sourceHandle.getFilename())) {
                    final var parseErrorBaseline = diagnostics.errorCount();
                    final var barrelAst = parseBarrelFile(fileId, utf8Content, diagnostics);
                    moduleUnit.barrels.add(new PbsModuleVisibilityValidator.BarrelFile(fileId, barrelAst));
                    if (diagnostics.errorCount() > parseErrorBaseline) {
                        failedModuleIds.add(moduleId);
                    }
                } else {
                    p.studio.compiler.source.diagnostics.Diagnostics.error(
                            diagnostics,
                            PbsLinkErrors.E_LINK_INVALID_BARREL_FILENAME.name(),
                            "Only 'mod.barrel' is allowed as barrel filename",
                            new Span(fileId, 0, utf8Content.getBytes(StandardCharsets.UTF_8).length));
                    failedModuleIds.add(moduleId);
                }
            }
            default -> {
            }
        }
    }

    private void loadReservedStdlibModules(
            final Map<PbsModuleVisibilityValidator.ModuleCoordinates, MutableModuleUnit> modulesByCoordinates,
            final ArrayList<PbsParsedSourceFile> parsedSourceFiles,
            final Map<FileId, ModuleId> moduleIdByFile,
            final ModuleTable moduleTable,
            final FileTable fileTable,
            final Map<String, ProjectId> projectIdByName,
            final ProjectId defaultSyntheticOwnerProjectId,
            final DiagnosticSink diagnostics,
            final int stdlibVersion) {
        final var stdlibEnvironment = stdlibEnvironmentResolver.resolve(stdlibVersion);
        final var pending = new ArrayDeque<PbsModuleVisibilityValidator.ModuleCoordinates>();
        final var resolved = new HashSet<ModuleId>();

        enqueueReservedImportsFromKnownModules(modulesByCoordinates, pending);
        while (!pending.isEmpty()) {
            final var target = pending.removeFirst();
            final var targetId = moduleId(moduleTable, target);
            if (!resolved.add(targetId)) {
                continue;
            }
            if (modulesByCoordinates.containsKey(target)) {
                continue;
            }

            final var moduleSource = stdlibEnvironment.resolveModule(target.project(), target.pathSegments());
            if (moduleSource.isEmpty()) {
                continue;
            }

            final var syntheticOwnerProjectId = projectIdByName.getOrDefault(target.project(), defaultSyntheticOwnerProjectId);
            if (syntheticOwnerProjectId == null) {
                continue;
            }
            final var loadedModule = interfaceModuleLoader.load(
                    moduleSource.get(),
                    syntheticOwnerProjectId,
                    fileTable,
                    diagnostics);
            final var moduleData = new MutableModuleUnit();
            moduleData.sources.addAll(loadedModule.sourceFiles().asList());
            moduleData.barrels.addAll(loadedModule.barrelFiles().asList());
            modulesByCoordinates.put(loadedModule.coordinates(), moduleData);
            for (final var sourceFile : loadedModule.sourceFiles()) {
                moduleIdByFile.put(sourceFile.fileId(), targetId);
                parsedSourceFiles.add(new PbsParsedSourceFile(
                        sourceFile.fileId(),
                        sourceFile.ast(),
                        targetId,
                        SourceKind.SDK_INTERFACE));
            }
            for (final var barrelFile : loadedModule.barrelFiles()) {
                moduleIdByFile.put(barrelFile.fileId(), targetId);
            }
            enqueueReservedImportsFromSourceFiles(loadedModule.sourceFiles().asList(), pending);
        }
    }

    private void enqueueReservedImportsFromKnownModules(
            final Map<PbsModuleVisibilityValidator.ModuleCoordinates, MutableModuleUnit> modulesByCoordinates,
            final ArrayDeque<PbsModuleVisibilityValidator.ModuleCoordinates> pending) {
        for (final var moduleUnit : modulesByCoordinates.values()) {
            enqueueReservedImportsFromSourceFiles(moduleUnit.sources, pending);
        }
    }

    private void enqueueReservedImportsFromSourceFiles(
            final Iterable<PbsModuleVisibilityValidator.SourceFile> sourceFiles,
            final ArrayDeque<PbsModuleVisibilityValidator.ModuleCoordinates> pending) {
        for (final var source : sourceFiles) {
            for (final var importDecl : source.ast().imports()) {
                final var moduleRef = importDecl.moduleRef();
                if (!isReservedProjectSpace(moduleRef.project())) {
                    continue;
                }
                pending.addLast(new PbsModuleVisibilityValidator.ModuleCoordinates(
                        moduleRef.project(),
                        moduleRef.pathSegments()));
            }
        }
    }

    private boolean isReservedProjectSpace(final String projectName) {
        return "core".equals(projectName) || "sdk".equals(projectName);
    }

    private PbsAst.File parseSourceFile(
            final FileId fileId,
            final String source,
            final DiagnosticSink diagnostics,
            final SourceKind sourceKind) {
        final var tokens = PbsLexer.lex(source, fileId, diagnostics);
        final var parseMode = sourceKind == SourceKind.SDK_INTERFACE
                ? PbsParser.ParseMode.INTERFACE_MODULE
                : PbsParser.ParseMode.ORDINARY;
        return PbsParser.parse(tokens, fileId, diagnostics, parseMode);
    }

    private PbsAst.BarrelFile parseBarrelFile(
            final FileId fileId,
            final String source,
            final DiagnosticSink diagnostics) {
        final var tokens = PbsLexer.lex(source, fileId, diagnostics);
        return PbsBarrelParser.parse(tokens, fileId, diagnostics);
    }

    private PbsModuleVisibilityValidator.ModuleCoordinates resolveModuleCoordinates(
            final ProjectDescriptor projectDescriptor,
            final SourceHandle sourceHandle) {
        final var sourceRelativePath = sourceRelativePath(projectDescriptor, sourceHandle);
        final var modulePath = sourceRelativePath.getParent();
        final var pathSegments = new ArrayList<String>();
        if (modulePath != null) {
            for (final var segment : modulePath) {
                pathSegments.add(segment.toString());
            }
        }

        return new PbsModuleVisibilityValidator.ModuleCoordinates(
                projectDescriptor.getName(),
                ReadOnlyList.wrap(pathSegments));
    }

    private Path sourceRelativePath(
            final ProjectDescriptor projectDescriptor,
            final SourceHandle sourceHandle) {
        final var canonPath = sourceHandle.getCanonPath();
        for (final var sourceRoot : projectDescriptor.getSourceRoots()) {
            if (canonPath.startsWith(sourceRoot)) {
                return sourceRoot.relativize(canonPath);
            }
        }
        return sourceHandle.getRelativePath();
    }

    private void markModulesWithLinkingErrors(
            final DiagnosticSink diagnostics,
            final Map<FileId, ModuleId> moduleIdByFile,
            final Set<ModuleId> failedModuleIds) {
        for (final var diagnostic : diagnostics) {
            if (!diagnostic.getSeverity().isError()) {
                continue;
            }
            if (diagnostic.getPhase() != DiagnosticPhase.LINKING) {
                continue;
            }
            final var moduleId = moduleIdByFile.get(diagnostic.getSpan().getFileId());
            if (moduleId != null) {
                failedModuleIds.add(moduleId);
            }
        }
    }

    private Map<ModuleId, Set<ModuleId>> buildModuleDependencyGraph(
            final ArrayList<PbsParsedSourceFile> parsedSourceFiles,
            final ModuleTable moduleTable) {
        final Map<ModuleId, Set<ModuleId>> dependenciesByModule = new HashMap<>();
        for (final var parsedSource : parsedSourceFiles) {
            final var moduleDependencies = dependenciesByModule.computeIfAbsent(
                    parsedSource.moduleId(),
                    ignored -> new HashSet<>());
            for (final var importDecl : parsedSource.ast().imports()) {
                final var moduleRef = importDecl.moduleRef();
                moduleDependencies.add(moduleId(moduleTable, moduleRef.project(), moduleRef.pathSegments()));
            }
        }
        return dependenciesByModule;
    }

    private ModuleId moduleId(
            final ModuleTable moduleTable,
            final PbsModuleVisibilityValidator.ModuleCoordinates coordinates) {
        return moduleId(moduleTable, coordinates.project(), coordinates.pathSegments());
    }

    private ModuleId moduleId(
            final ModuleTable moduleTable,
            final String project,
            final ReadOnlyList<String> pathSegments) {
        return moduleTable.register(new ModuleReference(project, pathSegments));
    }

    private ReadOnlyList<ModuleReference> emitModulePool(final ModuleTable moduleTable) {
        final var pool = new ArrayList<ModuleReference>(moduleTable.size());
        for (final var moduleId : moduleTable.identifiers()) {
            pool.add(moduleTable.get(moduleId));
        }
        return ReadOnlyList.wrap(pool);
    }

    private static final class MutableModuleUnit {
        private final ArrayList<PbsModuleVisibilityValidator.SourceFile> sources = new ArrayList<>();
        private final ArrayList<PbsModuleVisibilityValidator.BarrelFile> barrels = new ArrayList<>();
    }
}
