package p.studio.compiler.pbs.linking;

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.diagnostics.RelatedSpan;
import p.studio.compiler.source.diagnostics.Severity;
import p.studio.compiler.source.identifiers.CallableShapeId;
import p.studio.compiler.source.identifiers.FileId;
import p.studio.compiler.source.identifiers.NameId;
import p.studio.compiler.source.identifiers.TypeSurfaceId;
import p.studio.compiler.source.tables.CallableShapeTable;
import p.studio.compiler.source.tables.NameTable;
import p.studio.compiler.source.tables.TypeSurfaceTable;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.*;

import static p.studio.compiler.pbs.PbsReservedMetadataExtractor.typeSurfaceKey;

public final class PbsModuleVisibilityValidator {

    public void validate(
            final ReadOnlyList<ModuleUnit> modules,
            final DiagnosticSink diagnostics) {
        validate(modules, new NameTable(), diagnostics);
    }

    public void validate(
            final ReadOnlyList<ModuleUnit> modules,
            final NameTable nameTable,
            final DiagnosticSink diagnostics) {
        final var signatureInterning = new SignatureInterning(new TypeSurfaceTable(), new CallableShapeTable());
        final var exportsByModule = new HashMap<ModuleRefKey, ModuleExports>();

        for (final var module : modules) {
            final var exports = validateModule(module, nameTable, signatureInterning, diagnostics);
            exportsByModule.put(toModuleRef(module.coordinates()), exports);
        }

        for (final var module : modules) {
            validateImports(module, exportsByModule, nameTable, diagnostics);
        }
    }

    private ModuleExports validateModule(
            final ModuleUnit module,
            final NameTable nameTable,
            final SignatureInterning signatureInterning,
            final DiagnosticSink diagnostics) {
        final var exports = new ModuleExports();

        if (module.sourceFiles().isEmpty()) {
            if (!module.barrelFiles().isEmpty()) {
                p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, 
                        PbsLinkErrors.E_LINK_BARREL_WITHOUT_SOURCE.name(),
                        "Module %s has mod.barrel but no .pbs source files".formatted(displayModule(module.coordinates())),
                        module.barrelFiles().getFirst().ast().span());
            }
            return exports;
        }

        if (module.barrelFiles().isEmpty()) {
            p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, 
                    PbsLinkErrors.E_LINK_MISSING_BARREL.name(),
                    "Module %s is missing mod.barrel".formatted(displayModule(module.coordinates())),
                    module.sourceFiles().getFirst().ast().span());
            return exports;
        }

        if (module.barrelFiles().size() > 1) {
            final var first = module.barrelFiles().getFirst();
            for (int i = 1; i < module.barrelFiles().size(); i++) {
                final var duplicate = module.barrelFiles().get(i);
                p.studio.compiler.source.diagnostics.Diagnostics.report(diagnostics, 
                        Severity.Error,
                        PbsLinkErrors.E_LINK_DUPLICATE_BARREL_FILE.name(),
                        "Module %s has multiple mod.barrel files".formatted(displayModule(module.coordinates())),
                        duplicate.ast().span(),
                        List.of(new RelatedSpan("First mod.barrel is here", first.ast().span())));
            }
        }

        final var barrel = module.barrelFiles().getFirst();
        final var declarations = collectDeclarations(module, nameTable, signatureInterning);
        exports.declaredNameIds.addAll(declarations.declaredNameIds);
        final Map<NonFunctionSymbolKey, Span> seenNonFunctionEntries = new HashMap<>();
        final Map<FunctionSymbolKey, Span> seenFunctionEntries = new HashMap<>();

        for (final var item : barrel.ast().items()) {
            if (item instanceof PbsAst.BarrelFunctionItem(
                    PbsAst.Visibility visibility, String name, ReadOnlyList<PbsAst.TypeRef> parameterTypes,
                    PbsAst.ReturnKind returnKind, PbsAst.TypeRef returnType, PbsAst.TypeRef resultErrorType, Span span
            )) {
                final var functionKey = functionKey(
                        name,
                        parameterTypes.asList(),
                        returnKind,
                        returnType,
                        resultErrorType,
                        nameTable,
                        signatureInterning);
                final var firstFunctionEntry = seenFunctionEntries.putIfAbsent(functionKey, span);
                if (firstFunctionEntry != null) {
                    p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, 
                            PbsLinkErrors.E_LINK_DUPLICATE_BARREL_ENTRY.name(),
                            "Duplicate barrel function entry '%s' in module %s".formatted(
                                    name,
                                    displayModule(module.coordinates())),
                            span,
                            List.of(new RelatedSpan("First barrel function entry is here", firstFunctionEntry)));
                    continue;
                }

                final var matches = declarations.functionsBySignature.getOrDefault(functionKey, List.of());
                if (matches.isEmpty()) {
                    p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, 
                            PbsLinkErrors.E_LINK_UNRESOLVED_BARREL_ENTRY.name(),
                            "Barrel function entry '%s' in module %s does not resolve to a declaration".formatted(
                                    name,
                                    displayModule(module.coordinates())),
                            span);
                    continue;
                }
                if (matches.size() > 1) {
                    p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, 
                            PbsLinkErrors.E_LINK_AMBIGUOUS_BARREL_ENTRY.name(),
                            "Barrel function entry '%s' in module %s resolves ambiguously".formatted(
                                    name,
                                    displayModule(module.coordinates())),
                            span,
                            List.of(new RelatedSpan("One matching function declaration is here", matches.getFirst())));
                    continue;
                }

                if (visibility == PbsAst.Visibility.PUB) {
                    exports.publicNameIds.add(functionKey.nameId());
                }
                continue;
            }

            final var symbolKind = nonFunctionKind(item);
            if (symbolKind == null) {
                continue;
            }

            final var symbolKey = new NonFunctionSymbolKey(symbolKind, nameTable.register(nonFunctionName(item)));
            final var firstNonFunctionEntry = seenNonFunctionEntries.putIfAbsent(symbolKey, item.span());
            if (firstNonFunctionEntry != null) {
                p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, 
                        PbsLinkErrors.E_LINK_DUPLICATE_BARREL_ENTRY.name(),
                        "Duplicate barrel entry '%s' in module %s".formatted(
                                nonFunctionName(item),
                                displayModule(module.coordinates())),
                        item.span(),
                        List.of(new RelatedSpan("First barrel entry is here", firstNonFunctionEntry)));
                continue;
            }

            final var matches = declarations.nonFunctionsByKindAndName.getOrDefault(symbolKey, List.of());
            if (matches.isEmpty()) {
                p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, 
                        PbsLinkErrors.E_LINK_UNRESOLVED_BARREL_ENTRY.name(),
                        "Barrel entry '%s' in module %s does not resolve to a declaration".formatted(
                                nonFunctionName(item),
                                displayModule(module.coordinates())),
                        item.span());
                continue;
            }
            if (matches.size() > 1) {
                p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, 
                        PbsLinkErrors.E_LINK_AMBIGUOUS_BARREL_ENTRY.name(),
                        "Barrel entry '%s' in module %s resolves ambiguously".formatted(
                                nonFunctionName(item),
                                displayModule(module.coordinates())),
                        item.span(),
                        List.of(new RelatedSpan("One matching declaration is here", matches.getFirst())));
                continue;
            }

            if (nonFunctionVisibility(item) == PbsAst.Visibility.PUB) {
                exports.publicNameIds.add(symbolKey.nameId());
            }
        }

        return exports;
    }

    private void validateImports(
            final ModuleUnit module,
            final Map<ModuleRefKey, ModuleExports> exportsByModule,
            final NameTable nameTable,
            final DiagnosticSink diagnostics) {
        for (final var source : module.sourceFiles()) {
            for (final var importDecl : source.ast().imports()) {
                final var targetModule = new ModuleRefKey(
                        importDecl.moduleRef().project(),
                        importDecl.moduleRef().pathSegments());
                final var targetExports = exportsByModule.get(targetModule);
                if (targetExports == null) {
                    p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics,
                            PbsLinkErrors.E_LINK_IMPORT_MODULE_NOT_FOUND.name(),
                            "Import target module %s was not found".formatted(displayModule(targetModule)),
                            importDecl.moduleRef().span());
                    continue;
                }

                if (importDecl.items().isEmpty()) {
                    continue;
                }

                for (final var importItem : importDecl.items()) {
                    final var importedNameId = nameTable.register(importItem.name());
                    if (targetExports.publicNameIds.contains(importedNameId)) {
                        continue;
                    }
                    if (!targetExports.declaredNameIds.contains(importedNameId)) {
                        p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics,
                                PbsLinkErrors.E_LINK_IMPORT_SYMBOL_UNRESOLVED.name(),
                                "Symbol '%s' does not exist in module %s".formatted(
                                        importItem.name(),
                                        displayModule(targetModule)),
                                importItem.span());
                        continue;
                    }
                    if (!targetExports.publicNameIds.contains(importedNameId)) {
                        p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, 
                                PbsLinkErrors.E_LINK_IMPORT_SYMBOL_NOT_PUBLIC.name(),
                                "Symbol '%s' is not public in module %s".formatted(
                                        importItem.name(),
                                        displayModule(targetModule)),
                                importItem.span());
                    }
                }
            }
        }
    }

    private ModuleDeclarations collectDeclarations(
            final ModuleUnit module,
            final NameTable nameTable,
            final SignatureInterning signatureInterning) {
        final var declarations = new ModuleDeclarations();

        for (final var sourceFile : module.sourceFiles()) {
            for (final var topDecl : sourceFile.ast().topDecls()) {
                if (topDecl instanceof PbsAst.FunctionDecl functionDecl) {
                    final var key = functionKey(
                            functionDecl.name(),
                            functionDecl.parameters().stream().map(PbsAst.Parameter::typeRef).toList(),
                            functionDecl.returnKind(),
                            functionDecl.returnType(),
                            functionDecl.resultErrorType(),
                            nameTable,
                            signatureInterning);
                    declarations.functionsBySignature
                            .computeIfAbsent(key, ignored -> new ArrayList<>())
                            .add(functionDecl.span());
                    declarations.declaredNameIds.add(key.nameId());
                    continue;
                }

                if (topDecl instanceof PbsAst.StructDecl structDecl) {
                    registerNonFunctionDeclaration(declarations, NonFunctionKind.STRUCT, structDecl.name(), structDecl.span(), nameTable);
                } else if (topDecl instanceof PbsAst.BuiltinTypeDecl builtinTypeDecl) {
                    registerNonFunctionDeclaration(declarations, NonFunctionKind.STRUCT, builtinTypeDecl.name(), builtinTypeDecl.span(), nameTable);
                } else if (topDecl instanceof PbsAst.ContractDecl contractDecl) {
                    registerNonFunctionDeclaration(declarations, NonFunctionKind.CONTRACT, contractDecl.name(), contractDecl.span(), nameTable);
                } else if (topDecl instanceof PbsAst.HostDecl hostDecl) {
                    registerNonFunctionDeclaration(declarations, NonFunctionKind.HOST, hostDecl.name(), hostDecl.span(), nameTable);
                } else if (topDecl instanceof PbsAst.ErrorDecl errorDecl) {
                    registerNonFunctionDeclaration(declarations, NonFunctionKind.ERROR, errorDecl.name(), errorDecl.span(), nameTable);
                } else if (topDecl instanceof PbsAst.EnumDecl enumDecl) {
                    registerNonFunctionDeclaration(declarations, NonFunctionKind.ENUM, enumDecl.name(), enumDecl.span(), nameTable);
                } else if (topDecl instanceof PbsAst.ServiceDecl serviceDecl) {
                    registerNonFunctionDeclaration(declarations, NonFunctionKind.SERVICE, serviceDecl.name(), serviceDecl.span(), nameTable);
                } else if (topDecl instanceof PbsAst.GlobalDecl globalDecl) {
                    registerNonFunctionDeclaration(declarations, NonFunctionKind.GLOBAL, globalDecl.name(), globalDecl.span(), nameTable);
                } else if (topDecl instanceof PbsAst.ConstDecl constDecl) {
                    registerNonFunctionDeclaration(declarations, NonFunctionKind.CONST, constDecl.name(), constDecl.span(), nameTable);
                } else if (topDecl instanceof PbsAst.CallbackDecl callbackDecl) {
                    registerNonFunctionDeclaration(declarations, NonFunctionKind.CALLBACK, callbackDecl.name(), callbackDecl.span(), nameTable);
                }
            }
        }

        return declarations;
    }

    private void registerNonFunctionDeclaration(
            final ModuleDeclarations declarations,
            final NonFunctionKind kind,
            final String name,
            final Span span,
            final NameTable nameTable) {
        final var key = new NonFunctionSymbolKey(kind, nameTable.register(name));
        declarations.nonFunctionsByKindAndName
                .computeIfAbsent(key, ignored -> new ArrayList<>())
                .add(span);
        declarations.declaredNameIds.add(key.nameId());
    }

    private FunctionSymbolKey functionKey(
            final String name,
            final List<PbsAst.TypeRef> parameterTypes,
            final PbsAst.ReturnKind returnKind,
            final PbsAst.TypeRef returnType,
            final PbsAst.TypeRef resultErrorType,
            final NameTable nameTable,
            final SignatureInterning signatureInterning) {
        return new FunctionSymbolKey(
                nameTable.register(name),
                callableShapeId(parameterTypes, returnKind, returnType, resultErrorType, signatureInterning));
    }

    private CallableShapeId callableShapeId(
            final List<PbsAst.TypeRef> parameterTypes,
            final PbsAst.ReturnKind returnKind,
            final PbsAst.TypeRef returnType,
            final PbsAst.TypeRef resultErrorType,
            final SignatureInterning signatureInterning) {
        final var parameterSurfaceIds = new ArrayList<TypeSurfaceId>(parameterTypes.size());
        for (final var parameterType : parameterTypes) {
            parameterSurfaceIds.add(signatureInterning.typeSurfaces().register(typeSurfaceKey(parameterType)));
        }
        final var outputSurfaceId = signatureInterning.typeSurfaces().register(outputSurfaceKey(returnKind, returnType, resultErrorType));
        return signatureInterning.callableShapes().register(ReadOnlyList.wrap(parameterSurfaceIds), outputSurfaceId);
    }

    private String outputSurfaceKey(
            final PbsAst.ReturnKind returnKind,
            final PbsAst.TypeRef returnType,
            final PbsAst.TypeRef resultErrorType) {
        return switch (returnKind) {
            case INFERRED_UNIT -> "infer_unit";
            case EXPLICIT_UNIT -> "unit";
            case PLAIN -> typeSurfaceKey(returnType);
            case RESULT -> "result<" + typeSurfaceKey(resultErrorType) + ">" + typeSurfaceKey(returnType);
        };
    }

    private NonFunctionKind nonFunctionKind(final PbsAst.BarrelItem item) {
        if (item instanceof PbsAst.BarrelStructItem) {
            return NonFunctionKind.STRUCT;
        }
        if (item instanceof PbsAst.BarrelContractItem) {
            return NonFunctionKind.CONTRACT;
        }
        if (item instanceof PbsAst.BarrelHostItem) {
            return NonFunctionKind.HOST;
        }
        if (item instanceof PbsAst.BarrelErrorItem) {
            return NonFunctionKind.ERROR;
        }
        if (item instanceof PbsAst.BarrelEnumItem) {
            return NonFunctionKind.ENUM;
        }
        if (item instanceof PbsAst.BarrelServiceItem) {
            return NonFunctionKind.SERVICE;
        }
        if (item instanceof PbsAst.BarrelGlobalItem) {
            return NonFunctionKind.GLOBAL;
        }
        if (item instanceof PbsAst.BarrelConstItem) {
            return NonFunctionKind.CONST;
        }
        if (item instanceof PbsAst.BarrelCallbackItem) {
            return NonFunctionKind.CALLBACK;
        }
        return null;
    }

    private String nonFunctionName(final PbsAst.BarrelItem item) {
        if (item instanceof PbsAst.BarrelStructItem structItem) {
            return structItem.name();
        }
        if (item instanceof PbsAst.BarrelContractItem contractItem) {
            return contractItem.name();
        }
        if (item instanceof PbsAst.BarrelHostItem hostItem) {
            return hostItem.name();
        }
        if (item instanceof PbsAst.BarrelErrorItem errorItem) {
            return errorItem.name();
        }
        if (item instanceof PbsAst.BarrelEnumItem enumItem) {
            return enumItem.name();
        }
        if (item instanceof PbsAst.BarrelServiceItem serviceItem) {
            return serviceItem.name();
        }
        if (item instanceof PbsAst.BarrelGlobalItem globalItem) {
            return globalItem.name();
        }
        if (item instanceof PbsAst.BarrelConstItem constItem) {
            return constItem.name();
        }
        if (item instanceof PbsAst.BarrelCallbackItem callbackItem) {
            return callbackItem.name();
        }
        return "<unknown>";
    }

    private PbsAst.Visibility nonFunctionVisibility(final PbsAst.BarrelItem item) {
        if (item instanceof PbsAst.BarrelStructItem structItem) {
            return structItem.visibility();
        }
        if (item instanceof PbsAst.BarrelContractItem contractItem) {
            return contractItem.visibility();
        }
        if (item instanceof PbsAst.BarrelHostItem hostItem) {
            return hostItem.visibility();
        }
        if (item instanceof PbsAst.BarrelErrorItem errorItem) {
            return errorItem.visibility();
        }
        if (item instanceof PbsAst.BarrelEnumItem enumItem) {
            return enumItem.visibility();
        }
        if (item instanceof PbsAst.BarrelServiceItem serviceItem) {
            return serviceItem.visibility();
        }
        if (item instanceof PbsAst.BarrelGlobalItem globalItem) {
            return globalItem.visibility();
        }
        if (item instanceof PbsAst.BarrelConstItem constItem) {
            return constItem.visibility();
        }
        if (item instanceof PbsAst.BarrelCallbackItem callbackItem) {
            return callbackItem.visibility();
        }
        return PbsAst.Visibility.MOD;
    }

    private ModuleRefKey toModuleRef(final ModuleCoordinates coordinates) {
        return new ModuleRefKey(coordinates.project(), coordinates.pathSegments());
    }

    private String displayModule(final ModuleCoordinates coordinates) {
        return displayModule(new ModuleRefKey(coordinates.project(), coordinates.pathSegments()));
    }

    private String displayModule(final ModuleRefKey moduleRef) {
        if (moduleRef.pathSegments().isEmpty()) {
            return "@%s:<root>".formatted(moduleRef.project());
        }
        return "@%s:%s".formatted(moduleRef.project(), String.join("/", moduleRef.pathSegments().asList()));
    }

    public record ModuleCoordinates(
            String project,
            ReadOnlyList<String> pathSegments) {
    }

    public record SourceFile(
            FileId fileId,
            PbsAst.File ast) {
    }

    public record BarrelFile(
            FileId fileId,
            PbsAst.BarrelFile ast) {
    }

    public record ModuleUnit(
            ModuleCoordinates coordinates,
            ReadOnlyList<SourceFile> sourceFiles,
            ReadOnlyList<BarrelFile> barrelFiles) {
    }

    private record ModuleRefKey(
            String project,
            ReadOnlyList<String> pathSegments) {
    }

    private enum NonFunctionKind {
        STRUCT,
        CONTRACT,
        HOST,
        ERROR,
        ENUM,
        SERVICE,
        GLOBAL,
        CONST,
        CALLBACK
    }

    private record NonFunctionSymbolKey(
            NonFunctionKind kind,
            NameId nameId) {
    }

    private record FunctionSymbolKey(
            NameId nameId,
            CallableShapeId callableShapeId) {
    }

    private record SignatureInterning(
            TypeSurfaceTable typeSurfaces,
            CallableShapeTable callableShapes) {
    }

    private static final class ModuleDeclarations {
        private final Map<NonFunctionSymbolKey, List<Span>> nonFunctionsByKindAndName = new HashMap<>();
        private final Map<FunctionSymbolKey, List<Span>> functionsBySignature = new HashMap<>();
        private final Set<NameId> declaredNameIds = new HashSet<>();
    }

    private static final class ModuleExports {
        private final Set<NameId> declaredNameIds = new HashSet<>();
        private final Set<NameId> publicNameIds = new HashSet<>();
    }
}
