package p.studio.compiler.services;

import p.studio.compiler.models.IRReservedMetadata;
import p.studio.compiler.pbs.PbsCallableSlotCounter;
import p.studio.compiler.pbs.PbsFrontendCompiler;
import p.studio.compiler.pbs.PbsReservedMetadataExtractor;
import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.semantics.PbsGlobalDependencySupport;
import p.studio.compiler.source.identifiers.FileId;
import p.studio.compiler.source.identifiers.ModuleId;
import p.studio.compiler.source.tables.ModuleTable;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.*;

final class PbsImportedSemanticContextService {
    private final PbsReservedMetadataExtractor reservedMetadataExtractor;
    private final PbsFrontendCompiler frontendCompiler;

    PbsImportedSemanticContextService(
            final PbsReservedMetadataExtractor reservedMetadataExtractor,
            final PbsFrontendCompiler frontendCompiler) {
        this.reservedMetadataExtractor = reservedMetadataExtractor;
        this.frontendCompiler = frontendCompiler;
    }

    Map<FileId, PbsImportedSemanticContext> build(
            final ReadOnlyList<PbsParsedSourceFile> parsedSourceFiles,
            final ModuleTable moduleTable) {
        final Map<ModuleId, ArrayList<PbsParsedSourceFile>> sourcesByModule = new HashMap<>();
        for (final var parsedSource : parsedSourceFiles) {
            sourcesByModule
                    .computeIfAbsent(parsedSource.moduleId(), ignored -> new ArrayList<>())
                    .add(parsedSource);
        }

        final Map<ModuleId, IRReservedMetadata> reservedMetadataByModule = new HashMap<>();
        for (final var entry : sourcesByModule.entrySet()) {
            var merged = IRReservedMetadata.empty();
            for (final var source : entry.getValue()) {
                final var extracted = reservedMetadataExtractor.extract(source.ast(), source.sourceKind());
                merged = mergeReservedMetadata(merged, extracted);
            }
            reservedMetadataByModule.put(entry.getKey(), merged);
        }

        final Map<ModuleId, Map<String, ArrayList<PbsAst.TopDecl>>> topDeclsByNameByModule = new HashMap<>();
        for (final var entry : sourcesByModule.entrySet()) {
            topDeclsByNameByModule.put(entry.getKey(), indexTopDeclsByName(entry.getValue()));
        }
        final var globalInfosByRef = indexGlobalInfos(sourcesByModule, topDeclsByNameByModule, moduleTable);

        final Map<FileId, PbsImportedSemanticContext> contexts = new HashMap<>();
        for (final var parsedSource : parsedSourceFiles) {
            final var supplementalTopDecls = new ArrayList<PbsAst.TopDecl>();
            final var importedCallables = new ArrayList<PbsFrontendCompiler.ImportedCallableSurface>();
            final var importedGlobals = new ArrayList<PbsFrontendCompiler.ImportedGlobalSurface>();
            final var importedCallableKeys = new HashSet<String>();
            final var importedGlobalKeys = new HashSet<String>();
            final var supplementalKeys = new HashSet<String>();
            var importedReservedMetadata = IRReservedMetadata.empty();
            final var directlyVisibleGlobalRefs = new LinkedHashMap<String, PbsGlobalDependencySupport.GlobalRef>();

            for (final var importDecl : parsedSource.ast().imports()) {
                final var moduleRef = importDecl.moduleRef();
                final var importedModuleId = moduleTable.register(
                        new p.studio.compiler.source.tables.ModuleReference(moduleRef.project(), moduleRef.pathSegments()));
                final var importedSources = sourcesByModule.get(importedModuleId);
                if (importedSources == null || importedSources.isEmpty()) {
                    continue;
                }
                final var importedTopDeclsByName = topDeclsByNameByModule.get(importedModuleId);
                if (importedTopDeclsByName == null) {
                    continue;
                }
                importedReservedMetadata = mergeReservedMetadata(
                        importedReservedMetadata,
                        reservedMetadataByModule.getOrDefault(importedModuleId, IRReservedMetadata.empty()));

                for (final var importItem : importDecl.items()) {
                    final var importName = importItem.name();
                    final var localName = importItemLocalName(importItem);
                    final var directDecls = importedTopDeclsByName.getOrDefault(importName, new ArrayList<>());
                    for (final var topDecl : directDecls) {
                        appendSupplementalTopDecl(topDecl, supplementalTopDecls, supplementalKeys);
                        appendTypeDependencyDecls(topDecl, importedTopDeclsByName, supplementalTopDecls, supplementalKeys);
                        if (topDecl instanceof PbsAst.ServiceDecl serviceDecl) {
                            for (final var method : serviceDecl.methods()) {
                                appendImportedCallable(
                                        importedCallables,
                                        importedCallableKeys,
                                        new PbsFrontendCompiler.ImportedCallableSurface(
                                                importedModuleId,
                                                localName + "." + method.name(),
                                                method.parameters().size(),
                                                returnSlotsFor(method),
                                                method.returnKind(),
                                                method.returnType(),
                                                frontendCompiler.callableShapeSurfaceOf(method)));
                            }
                            continue;
                        }
                        if (topDecl instanceof PbsAst.FunctionDecl functionDecl) {
                            appendImportedCallable(
                                    importedCallables,
                                    importedCallableKeys,
                                    new PbsFrontendCompiler.ImportedCallableSurface(
                                            importedModuleId,
                                            localName,
                                            functionDecl.parameters().size(),
                                            returnSlotsFor(functionDecl),
                                            functionDecl.returnKind(),
                                            functionDecl.returnType(),
                                            frontendCompiler.callableShapeSurfaceOf(functionDecl)));
                            continue;
                        }
                        if (topDecl instanceof PbsAst.GlobalDecl globalDecl) {
                            directlyVisibleGlobalRefs.putIfAbsent(
                                    localName,
                                    new PbsGlobalDependencySupport.GlobalRef(importedModuleId, globalDecl.name()));
                        }
                    }
                }
            }

            final var reachableImportedGlobals = collectReachableImportedGlobals(directlyVisibleGlobalRefs, globalInfosByRef);
            for (final var entry : directlyVisibleGlobalRefs.entrySet()) {
                final var info = globalInfosByRef.get(entry.getValue());
                if (info == null) {
                    continue;
                }
                appendImportedGlobal(
                        importedGlobals,
                        importedGlobalKeys,
                        toImportedGlobalSurface(info, entry.getKey(), true));
            }
            for (final var reachableRef : reachableImportedGlobals) {
                final var info = globalInfosByRef.get(reachableRef);
                if (info == null) {
                    continue;
                }
                appendImportedGlobal(
                        importedGlobals,
                        importedGlobalKeys,
                        toImportedGlobalSurface(info, info.globalDecl().name(), false));
            }

            contexts.put(parsedSource.fileId(), new PbsImportedSemanticContext(
                    ReadOnlyList.wrap(supplementalTopDecls),
                    ReadOnlyList.wrap(importedCallables),
                    ReadOnlyList.wrap(importedGlobals),
                    importedReservedMetadata));
        }

        return contexts;
    }

    private Map<String, ArrayList<PbsAst.TopDecl>> indexTopDeclsByName(
            final List<PbsParsedSourceFile> moduleSources) {
        final Map<String, ArrayList<PbsAst.TopDecl>> index = new HashMap<>();
        for (final var source : moduleSources) {
            for (final var topDecl : source.ast().topDecls()) {
                final var declName = topDeclName(topDecl);
                if (declName == null || declName.isBlank()) {
                    continue;
                }
                index.computeIfAbsent(declName, ignored -> new ArrayList<>()).add(topDecl);
            }
        }
        return index;
    }

    private void appendTypeDependencyDecls(
            final PbsAst.TopDecl rootDecl,
            final Map<String, ArrayList<PbsAst.TopDecl>> topDeclsByName,
            final ArrayList<PbsAst.TopDecl> supplementalTopDecls,
            final Set<String> supplementalKeys) {
        final var pendingTypeNames = new ArrayDeque<String>();
        final var visitedTypeNames = new HashSet<String>();
        collectReferencedTypeNames(rootDecl, pendingTypeNames);
        while (!pendingTypeNames.isEmpty()) {
            final var typeName = pendingTypeNames.removeFirst();
            if (!visitedTypeNames.add(typeName)) {
                continue;
            }
            final var candidates = topDeclsByName.get(typeName);
            if (candidates == null) {
                continue;
            }
            for (final var candidate : candidates) {
                if (!isTypeDecl(candidate)) {
                    continue;
                }
                appendSupplementalTopDecl(candidate, supplementalTopDecls, supplementalKeys);
                collectReferencedTypeNames(candidate, pendingTypeNames);
            }
        }
    }

    private void collectReferencedTypeNames(
            final PbsAst.TopDecl topDecl,
            final ArrayDeque<String> sink) {
        if (topDecl instanceof PbsAst.StructDecl structDecl) {
            for (final var field : structDecl.fields()) {
                collectReferencedTypeNames(field.typeRef(), sink);
            }
            for (final var method : structDecl.methods()) {
                collectReferencedTypeNames(method, sink);
            }
            for (final var ctor : structDecl.ctors()) {
                for (final var parameter : ctor.parameters()) {
                    collectReferencedTypeNames(parameter.typeRef(), sink);
                }
            }
            return;
        }
        if (topDecl instanceof PbsAst.BuiltinTypeDecl builtinTypeDecl) {
            for (final var field : builtinTypeDecl.fields()) {
                collectReferencedTypeNames(field.typeRef(), sink);
            }
            for (final var signature : builtinTypeDecl.signatures()) {
                collectReferencedTypeNames(signature, sink);
            }
            return;
        }
        if (topDecl instanceof PbsAst.ServiceDecl serviceDecl) {
            for (final var method : serviceDecl.methods()) {
                collectReferencedTypeNames(method, sink);
            }
            return;
        }
        if (topDecl instanceof PbsAst.HostDecl hostDecl) {
            for (final var signature : hostDecl.signatures()) {
                collectReferencedTypeNames(signature, sink);
            }
            return;
        }
        if (topDecl instanceof PbsAst.ContractDecl contractDecl) {
            for (final var signature : contractDecl.signatures()) {
                collectReferencedTypeNames(signature, sink);
            }
            return;
        }
        if (topDecl instanceof PbsAst.CallbackDecl callbackDecl) {
            collectReferencedTypeNames(callbackDecl.parameters(), callbackDecl.returnType(), callbackDecl.resultErrorType(), sink);
            return;
        }
        if (topDecl instanceof PbsAst.FunctionDecl functionDecl) {
            collectReferencedTypeNames(functionDecl, sink);
            return;
        }
        if (topDecl instanceof PbsAst.ConstDecl constDecl && constDecl.explicitType() != null) {
            collectReferencedTypeNames(constDecl.explicitType(), sink);
            return;
        }
        if (topDecl instanceof PbsAst.GlobalDecl globalDecl && globalDecl.explicitType() != null) {
            collectReferencedTypeNames(globalDecl.explicitType(), sink);
        }
    }

    private void collectReferencedTypeNames(
            final PbsAst.FunctionDecl functionDecl,
            final ArrayDeque<String> sink) {
        collectReferencedTypeNames(
                functionDecl.parameters(),
                functionDecl.returnType(),
                functionDecl.resultErrorType(),
                sink);
    }

    private void collectReferencedTypeNames(
            final PbsAst.FunctionSignature functionSignature,
            final ArrayDeque<String> sink) {
        collectReferencedTypeNames(
                functionSignature.parameters(),
                functionSignature.returnType(),
                functionSignature.resultErrorType(),
                sink);
    }

    private void collectReferencedTypeNames(
            final ReadOnlyList<PbsAst.Parameter> parameters,
            final PbsAst.TypeRef returnType,
            final PbsAst.TypeRef resultErrorType,
            final ArrayDeque<String> sink) {
        for (final var parameter : parameters) {
            collectReferencedTypeNames(parameter.typeRef(), sink);
        }
        collectReferencedTypeNames(returnType, sink);
        collectReferencedTypeNames(resultErrorType, sink);
    }

    private void collectReferencedTypeNames(
            final PbsAst.TypeRef typeRef,
            final ArrayDeque<String> sink) {
        if (typeRef == null) {
            return;
        }
        switch (typeRef.kind()) {
            case SIMPLE -> {
                final var simpleName = typeRef.name();
                if (!isPrimitiveTypeName(simpleName)) {
                    sink.addLast(simpleName);
                }
            }
            case OPTIONAL, GROUP -> collectReferencedTypeNames(typeRef.inner(), sink);
            case NAMED_TUPLE -> {
                for (final var field : typeRef.fields()) {
                    collectReferencedTypeNames(field.typeRef(), sink);
                }
            }
            case UNIT, SELF, ERROR -> {
            }
        }
    }

    private boolean isPrimitiveTypeName(final String typeName) {
        return "int".equals(typeName)
                || "float".equals(typeName)
                || "bool".equals(typeName)
                || "str".equals(typeName)
                || "unit".equals(typeName);
    }

    private boolean isTypeDecl(final PbsAst.TopDecl topDecl) {
        return topDecl instanceof PbsAst.StructDecl
                || topDecl instanceof PbsAst.BuiltinTypeDecl
                || topDecl instanceof PbsAst.ServiceDecl
                || topDecl instanceof PbsAst.HostDecl
                || topDecl instanceof PbsAst.ContractDecl
                || topDecl instanceof PbsAst.CallbackDecl
                || topDecl instanceof PbsAst.EnumDecl
                || topDecl instanceof PbsAst.ErrorDecl;
    }

    private void appendSupplementalTopDecl(
            final PbsAst.TopDecl topDecl,
            final ArrayList<PbsAst.TopDecl> supplementalTopDecls,
            final Set<String> supplementalKeys) {
        final var key = topDeclKey(topDecl);
        if (key == null || key.isBlank()) {
            return;
        }
        if (supplementalKeys.add(key)) {
            supplementalTopDecls.add(topDecl);
        }
    }

    private void appendImportedCallable(
            final ArrayList<PbsFrontendCompiler.ImportedCallableSurface> importedCallables,
            final Set<String> importedCallableKeys,
            final PbsFrontendCompiler.ImportedCallableSurface importedCallableSurface) {
        final var callableKey = importedCallableSurface.moduleId().getIndex()
                + "#"
                + importedCallableSurface.callableName()
                + "#"
                + importedCallableSurface.arity()
                + "#"
                + importedCallableSurface.shapeSurface();
        if (importedCallableKeys.add(callableKey)) {
            importedCallables.add(importedCallableSurface);
        }
    }

    private void appendImportedGlobal(
            final ArrayList<PbsFrontendCompiler.ImportedGlobalSurface> importedGlobals,
            final Set<String> importedGlobalKeys,
            final PbsFrontendCompiler.ImportedGlobalSurface importedGlobalSurface) {
        final var globalKey = importedGlobalSurface.ownerModuleId().getIndex()
                + "#"
                + importedGlobalSurface.globalName()
                + "#"
                + importedGlobalSurface.localName()
                + "#"
                + importedGlobalSurface.directlyVisible();
        if (importedGlobalKeys.add(globalKey)) {
            importedGlobals.add(importedGlobalSurface);
        }
    }

    private String topDeclKey(final PbsAst.TopDecl topDecl) {
        final var declName = topDeclName(topDecl);
        if (declName == null || declName.isBlank()) {
            return null;
        }
        return topDecl.getClass().getSimpleName() + ":" + declName;
    }

    private String topDeclName(final PbsAst.TopDecl topDecl) {
        if (topDecl instanceof PbsAst.StructDecl structDecl) {
            return structDecl.name();
        }
        if (topDecl instanceof PbsAst.BuiltinTypeDecl builtinTypeDecl) {
            return builtinTypeDecl.name();
        }
        if (topDecl instanceof PbsAst.ConstDecl constDecl) {
            return constDecl.name();
        }
        if (topDecl instanceof PbsAst.GlobalDecl globalDecl) {
            return globalDecl.name();
        }
        if (topDecl instanceof PbsAst.ServiceDecl serviceDecl) {
            return serviceDecl.name();
        }
        if (topDecl instanceof PbsAst.HostDecl hostDecl) {
            return hostDecl.name();
        }
        if (topDecl instanceof PbsAst.ContractDecl contractDecl) {
            return contractDecl.name();
        }
        if (topDecl instanceof PbsAst.FunctionDecl functionDecl) {
            return functionDecl.name();
        }
        if (topDecl instanceof PbsAst.CallbackDecl callbackDecl) {
            return callbackDecl.name();
        }
        if (topDecl instanceof PbsAst.EnumDecl enumDecl) {
            return enumDecl.name();
        }
        if (topDecl instanceof PbsAst.ErrorDecl errorDecl) {
            return errorDecl.name();
        }
        return null;
    }

    private IRReservedMetadata mergeReservedMetadata(
            final IRReservedMetadata left,
            final IRReservedMetadata right) {
        final var hostBindings = new ArrayList<IRReservedMetadata.HostMethodBinding>();
        hostBindings.addAll(left.hostMethodBindings().asList());
        hostBindings.addAll(right.hostMethodBindings().asList());

        final var builtinTypeSurfaces = new ArrayList<IRReservedMetadata.BuiltinTypeSurface>();
        builtinTypeSurfaces.addAll(left.builtinTypeSurfaces().asList());
        builtinTypeSurfaces.addAll(right.builtinTypeSurfaces().asList());

        final var builtinConstSurfaces = new ArrayList<IRReservedMetadata.BuiltinConstSurface>();
        builtinConstSurfaces.addAll(left.builtinConstSurfaces().asList());
        builtinConstSurfaces.addAll(right.builtinConstSurfaces().asList());

        final var requiredCapabilities = new HashSet<String>();
        requiredCapabilities.addAll(left.requiredCapabilities().asList());
        requiredCapabilities.addAll(right.requiredCapabilities().asList());

        return new IRReservedMetadata(
                ReadOnlyList.wrap(hostBindings),
                ReadOnlyList.wrap(builtinTypeSurfaces),
                ReadOnlyList.wrap(builtinConstSurfaces),
                ReadOnlyList.wrap(requiredCapabilities.stream().toList()));
    }

    private int returnSlotsFor(final PbsAst.FunctionDecl functionDecl) {
        return PbsCallableSlotCounter.returnSlots(functionDecl.returnKind(), functionDecl.returnType());
    }

    private String importItemLocalName(final PbsAst.ImportItem importItem) {
        if (importItem.alias() == null || importItem.alias().isBlank()) {
            return importItem.name();
        }
        return importItem.alias();
    }

    private Map<PbsGlobalDependencySupport.GlobalRef, GlobalInfo> indexGlobalInfos(
            final Map<ModuleId, ArrayList<PbsParsedSourceFile>> sourcesByModule,
            final Map<ModuleId, Map<String, ArrayList<PbsAst.TopDecl>>> topDeclsByNameByModule,
            final ModuleTable moduleTable) {
        final Map<PbsGlobalDependencySupport.GlobalRef, GlobalInfo> globalsByRef = new LinkedHashMap<>();
        for (final var entry : sourcesByModule.entrySet()) {
            final var moduleId = entry.getKey();
            for (final var source : entry.getValue()) {
                final var importedGlobalsByLocalName = directImportedGlobalsFor(source.ast(), topDeclsByNameByModule, moduleTable);
                for (final var topDecl : source.ast().topDecls()) {
                    if (topDecl instanceof PbsAst.GlobalDecl globalDecl) {
                        globalsByRef.putIfAbsent(
                                new PbsGlobalDependencySupport.GlobalRef(moduleId, globalDecl.name()),
                                new GlobalInfo(moduleId, globalDecl, importedGlobalsByLocalName));
                    }
                }
            }
        }
        return globalsByRef;
    }

    private Map<String, PbsGlobalDependencySupport.GlobalRef> directImportedGlobalsFor(
            final PbsAst.File ast,
            final Map<ModuleId, Map<String, ArrayList<PbsAst.TopDecl>>> topDeclsByNameByModule,
            final ModuleTable moduleTable) {
        final Map<String, PbsGlobalDependencySupport.GlobalRef> importedGlobalsByLocalName = new LinkedHashMap<>();
        for (final var importDecl : ast.imports()) {
            final var importedModuleId = moduleTable.register(
                    new p.studio.compiler.source.tables.ModuleReference(
                            importDecl.moduleRef().project(),
                            importDecl.moduleRef().pathSegments()));
            final var importedTopDeclsByName = topDeclsByNameByModule.get(importedModuleId);
            if (importedTopDeclsByName == null) {
                continue;
            }
            for (final var importItem : importDecl.items()) {
                final var candidates = importedTopDeclsByName.getOrDefault(importItem.name(), new ArrayList<>());
                for (final var topDecl : candidates) {
                    if (topDecl instanceof PbsAst.GlobalDecl globalDecl) {
                        importedGlobalsByLocalName.putIfAbsent(
                                importItemLocalName(importItem),
                                new PbsGlobalDependencySupport.GlobalRef(importedModuleId, globalDecl.name()));
                        break;
                    }
                }
            }
        }
        return importedGlobalsByLocalName;
    }

    private LinkedHashSet<PbsGlobalDependencySupport.GlobalRef> collectReachableImportedGlobals(
            final Map<String, PbsGlobalDependencySupport.GlobalRef> directlyVisibleGlobalRefs,
            final Map<PbsGlobalDependencySupport.GlobalRef, GlobalInfo> globalInfosByRef) {
        final var visited = new LinkedHashSet<PbsGlobalDependencySupport.GlobalRef>();
        final var pending = new ArrayDeque<PbsGlobalDependencySupport.GlobalRef>(directlyVisibleGlobalRefs.values());
        while (!pending.isEmpty()) {
            final var current = pending.removeFirst();
            if (!visited.add(current)) {
                continue;
            }
            final var info = globalInfosByRef.get(current);
            if (info == null) {
                continue;
            }
            final var globalsInSameModule = new LinkedHashSet<String>();
            for (final var entry : globalInfosByRef.entrySet()) {
                if (entry.getKey().moduleId().equals(info.moduleId())) {
                    globalsInSameModule.add(entry.getKey().globalName());
                }
            }
            pending.addAll(PbsGlobalDependencySupport.collectDependencies(
                    info.globalDecl().initializer(),
                    info.moduleId(),
                    globalsInSameModule,
                    info.importedGlobalsByLocalName()));
        }
        return visited;
    }

    private PbsFrontendCompiler.ImportedGlobalSurface toImportedGlobalSurface(
            final GlobalInfo info,
            final String localName,
            final boolean directlyVisible) {
        final var importedDependencies = new ArrayList<PbsFrontendCompiler.ImportedGlobalDependency>();
        for (final var entry : info.importedGlobalsByLocalName().entrySet()) {
            importedDependencies.add(new PbsFrontendCompiler.ImportedGlobalDependency(
                    entry.getKey(),
                    entry.getValue().moduleId(),
                    entry.getValue().globalName()));
        }
        return new PbsFrontendCompiler.ImportedGlobalSurface(
                info.moduleId(),
                info.globalDecl().name(),
                localName,
                info.globalDecl().explicitType(),
                info.globalDecl().initializer(),
                ReadOnlyList.wrap(importedDependencies),
                info.globalDecl().span(),
                directlyVisible);
    }

    private record GlobalInfo(
            ModuleId moduleId,
            PbsAst.GlobalDecl globalDecl,
            Map<String, PbsGlobalDependencySupport.GlobalRef> importedGlobalsByLocalName) {
    }
}
