package p.studio.compiler.pbs;

import p.studio.compiler.messages.HostAdmissionContext;
import p.studio.compiler.messages.FESurfaceContext;
import p.studio.compiler.models.IRBackendExecutableFunction;
import p.studio.compiler.models.IRBackendFile;
import p.studio.compiler.models.IRFunction;
import p.studio.compiler.models.IRGlobal;
import p.studio.compiler.models.IRGlobalOrigin;
import p.studio.compiler.models.IRGlobalVisibility;
import p.studio.compiler.models.IRReservedMetadata;
import p.studio.compiler.models.IRSyntheticCallableKind;
import p.studio.compiler.models.IRSyntheticFunction;
import p.studio.compiler.models.IRSyntheticOrigin;
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.lowering.PbsExecutableLoweringService;
import p.studio.compiler.pbs.parser.PbsParser;
import p.studio.compiler.pbs.semantics.PbsDeclarationSemanticsValidator;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticsValidator;
import p.studio.compiler.pbs.semantics.PbsLifecycleSemanticsValidator;
import p.studio.compiler.source.Span;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.identifiers.CallableId;
import p.studio.compiler.source.identifiers.FileId;
import p.studio.compiler.source.identifiers.ModuleId;
import p.studio.compiler.source.tables.CallableSignatureRef;
import p.studio.compiler.source.tables.ModuleReference;
import p.studio.compiler.source.tables.NameTable;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.Map;

public final class PbsFrontendCompiler {
    private final PbsFlowSemanticsValidator flowSemanticsValidator = new PbsFlowSemanticsValidator();
    private final PbsLifecycleSemanticsValidator lifecycleSemanticsValidator = new PbsLifecycleSemanticsValidator();
    private final PbsReservedMetadataExtractor reservedMetadataExtractor = new PbsReservedMetadataExtractor();
    private final PbsHostAdmissionValidator hostAdmissionValidator = new PbsHostAdmissionValidator();
    private final PbsExecutableLoweringService executableLoweringService = new PbsExecutableLoweringService();

    public IRBackendFile compileFile(
            final FileId fileId,
            final String source,
            final DiagnosticSink diagnostics) {
        return compileFile(fileId, source, diagnostics, SourceKind.PROJECT);
    }

    public IRBackendFile compileFile(
            final FileId fileId,
            final String source,
            final DiagnosticSink diagnostics,
            final SourceKind sourceKind) {
        return compileFile(fileId, source, diagnostics, sourceKind, HostAdmissionContext.permissiveDefault());
    }

    public IRBackendFile compileFile(
            final FileId fileId,
            final String source,
            final DiagnosticSink diagnostics,
            final SourceKind sourceKind,
            final HostAdmissionContext hostAdmissionContext) {
        return compileFile(fileId, source, diagnostics, sourceKind, hostAdmissionContext, FESurfaceContext.empty());
    }

    public IRBackendFile compileFile(
            final FileId fileId,
            final String source,
            final DiagnosticSink diagnostics,
            final SourceKind sourceKind,
            final HostAdmissionContext hostAdmissionContext,
            final FESurfaceContext feSurfaceContext) {
        final var nameTable = new NameTable();
        final var admissionBaseline = diagnostics.errorCount();
        final var tokens = PbsLexer.lex(source, fileId, diagnostics);
        final var parseMode = sourceKind == SourceKind.SDK_INTERFACE
                ? PbsParser.ParseMode.INTERFACE_MODULE
                : PbsParser.ParseMode.ORDINARY;
        final var ast = PbsParser.parse(tokens, fileId, diagnostics, parseMode);
        final var irBackendFile = compileParsedFile(
                fileId,
                ast,
                diagnostics,
                sourceKind,
                ModuleId.none(),
                ReadOnlyList.empty(),
                hostAdmissionContext,
                feSurfaceContext,
                nameTable);
        if (diagnostics.errorCount() > admissionBaseline) {
            return IRBackendFile.empty(fileId);
        }
        return irBackendFile;
    }

    public IRBackendFile compileParsedFile(
            final FileId fileId,
            final PbsAst.File ast,
            final DiagnosticSink diagnostics,
            final SourceKind sourceKind,
            final ModuleId moduleId,
            final ReadOnlyList<ModuleReference> modulePool,
            final HostAdmissionContext hostAdmissionContext,
            final FESurfaceContext feSurfaceContext,
            final NameTable nameTable) {
        return compileParsedFile(
                fileId,
                ast,
                diagnostics,
                sourceKind,
                moduleId,
                modulePool,
                hostAdmissionContext,
                feSurfaceContext,
                nameTable,
                ReadOnlyList.empty(),
                ReadOnlyList.empty(),
                ReadOnlyList.empty(),
                IRReservedMetadata.empty(),
                Map.of());
    }

    public IRBackendFile compileParsedFile(
            final FileId fileId,
            final PbsAst.File ast,
            final DiagnosticSink diagnostics,
            final SourceKind sourceKind,
            final ModuleId moduleId,
            final ReadOnlyList<ModuleReference> modulePool,
            final HostAdmissionContext hostAdmissionContext,
            final FESurfaceContext feSurfaceContext,
            final NameTable nameTable,
            final ReadOnlyList<PbsAst.TopDecl> supplementalTopDecls,
            final ReadOnlyList<ImportedCallableSurface> importedCallables,
            final ReadOnlyList<ImportedGlobalSurface> importedGlobals,
            final IRReservedMetadata importedReservedMetadata,
            final Map<String, IRGlobalVisibility> declaredGlobalVisibilityByName) {
        final var effectiveModuleId = moduleId == null ? ModuleId.none() : moduleId;
        final var effectiveModulePool = modulePool == null ? ReadOnlyList.<ModuleReference>empty() : modulePool;
        final var effectiveNameTable = nameTable == null ? new NameTable() : nameTable;
        final var effectiveSupplementalTopDecls = supplementalTopDecls == null
                ? ReadOnlyList.<PbsAst.TopDecl>empty()
                : supplementalTopDecls;
        final var effectiveImportedCallables = importedCallables == null
                ? ReadOnlyList.<ImportedCallableSurface>empty()
                : importedCallables;
        final var effectiveImportedGlobals = importedGlobals == null
                ? ReadOnlyList.<ImportedGlobalSurface>empty()
                : importedGlobals;
        final var effectiveFESurfaceContext = feSurfaceContext == null
                ? FESurfaceContext.empty()
                : feSurfaceContext;
        final var effectiveImportedReservedMetadata = importedReservedMetadata == null
                ? IRReservedMetadata.empty()
                : importedReservedMetadata;
        final var effectiveDeclaredGlobalVisibilityByName = declaredGlobalVisibilityByName == null
                ? Map.<String, IRGlobalVisibility>of()
                : Map.copyOf(declaredGlobalVisibilityByName);
        final var semanticsErrorBaseline = diagnostics.errorCount();
        new PbsDeclarationSemanticsValidator(effectiveNameTable).validate(
                ast,
                sourceKind,
                effectiveModuleId,
                effectiveImportedGlobals,
                diagnostics);
        flowSemanticsValidator.validate(ast, effectiveSupplementalTopDecls, effectiveFESurfaceContext, diagnostics);
        lifecycleSemanticsValidator.validate(ast, effectiveSupplementalTopDecls, diagnostics);
        if (diagnostics.errorCount() > semanticsErrorBaseline) {
            return IRBackendFile.empty(fileId);
        }

        final var reservedMetadata = validateAndBuildReservedMetadata(
                ast,
                sourceKind,
                hostAdmissionContext,
                diagnostics);
        if (reservedMetadata == null) {
            return IRBackendFile.empty(fileId);
        }

        final ReadOnlyList<IRFunction> functions = sourceKind == SourceKind.SDK_INTERFACE
                ? ReadOnlyList.empty()
                : lowerFunctions(fileId, ast);
        final ReadOnlyList<IRSyntheticFunction> syntheticFunctions = sourceKind == SourceKind.SDK_INTERFACE
                ? ReadOnlyList.empty()
                : lowerSyntheticFunctions(fileId, effectiveModuleId, ast);
        final ReadOnlyList<IRGlobal> globals = sourceKind == SourceKind.SDK_INTERFACE
                ? ReadOnlyList.empty()
                : lowerGlobals(fileId, effectiveModuleId, ast, effectiveDeclaredGlobalVisibilityByName);
        final var loweringErrorBaseline = diagnostics.errorCount();
        final var executableLowering = executableLoweringService.lower(
                fileId,
                ast,
                effectiveSupplementalTopDecls,
                effectiveModuleId,
                mergeReservedMetadata(reservedMetadata, effectiveImportedReservedMetadata),
                effectiveNameTable,
                effectiveFESurfaceContext,
                diagnostics,
                effectiveImportedCallables);
        if (diagnostics.errorCount() > loweringErrorBaseline) {
            return IRBackendFile.empty(fileId);
        }
        final var lifecycleArtifacts = scanLifecycleArtifacts(ast);
        final var executableFunctions = new ArrayList<>(executableLowering.executableFunctions().asList());
        final var callableSignatures = new ArrayList<>(executableLowering.callableSignatures().asList());
        final var fileInitExecutable = lowerFileInitExecutable(
                fileId,
                effectiveModuleId,
                lifecycleArtifacts,
                executableFunctions,
                callableSignatures);
        if (fileInitExecutable != null) {
            executableFunctions.add(fileInitExecutable);
        }
        return new IRBackendFile(
                fileId,
                functions,
                syntheticFunctions,
                globals,
                ReadOnlyList.wrap(executableFunctions),
                reservedMetadata,
                effectiveModulePool,
                ReadOnlyList.wrap(callableSignatures),
                executableLowering.intrinsicPool());
    }

    public String callableShapeSurfaceOf(final PbsAst.FunctionDecl functionDecl) {
        return executableLoweringService.callableShapeSurfaceOf(functionDecl);
    }

    private IRReservedMetadata validateAndBuildReservedMetadata(
            final PbsAst.File ast,
            final SourceKind sourceKind,
            final HostAdmissionContext hostAdmissionContext,
            final DiagnosticSink diagnostics) {
        final var extractedReservedMetadata = reservedMetadataExtractor.extract(ast, sourceKind);
        final var hostAdmissionErrorBaseline = diagnostics.errorCount();
        final var requiredCapabilities = hostAdmissionValidator.validate(
                extractedReservedMetadata,
                hostAdmissionContext,
                diagnostics);
        if (diagnostics.errorCount() > hostAdmissionErrorBaseline) {
            return null;
        }
        return new IRReservedMetadata(
                extractedReservedMetadata.hostMethodBindings(),
                extractedReservedMetadata.builtinTypeSurfaces(),
                extractedReservedMetadata.builtinConstSurfaces(),
                requiredCapabilities);
    }

    private ReadOnlyList<IRFunction> lowerFunctions(final FileId fileId, final PbsAst.File ast) {
        final var functions = new ArrayList<IRFunction>(ast.functions().size());
        for (final var fn : ast.functions()) {
            functions.add(new IRFunction(
                    fileId,
                    fn.name(),
                    fn.parameters().size(),
                    fn.returnKind() != PbsAst.ReturnKind.INFERRED_UNIT,
                    fn.span()));
        }
        return ReadOnlyList.wrap(functions);
    }

    private ReadOnlyList<IRGlobal> lowerGlobals(
            final FileId fileId,
            final ModuleId moduleId,
            final PbsAst.File ast,
            final Map<String, IRGlobalVisibility> declaredGlobalVisibilityByName) {
        final var globals = new ArrayList<IRGlobal>();
        var nextSlot = 0;
        for (final var topDecl : ast.topDecls()) {
            if (!(topDecl instanceof PbsAst.GlobalDecl globalDecl)) {
                continue;
            }
            globals.add(new IRGlobal(
                    fileId,
                    moduleId,
                    globalDecl.name(),
                    PbsReservedMetadataExtractor.typeSurfaceKey(globalDecl.explicitType()),
                    nextSlot++,
                    declaredGlobalVisibilityByName.getOrDefault(globalDecl.name(), IRGlobalVisibility.MODULE),
                    false,
                    p.studio.compiler.models.IRHiddenGlobalKind.NONE,
                    new IRGlobalOrigin(
                            fileId,
                            moduleId,
                            globalDecl.name(),
                            globalDecl.span(),
                            globalDecl.name()),
                    globalDecl.span()));
        }
        return ReadOnlyList.wrap(globals);
    }

    private ReadOnlyList<IRSyntheticFunction> lowerSyntheticFunctions(
            final FileId fileId,
            final ModuleId moduleId,
            final PbsAst.File ast) {
        final var syntheticFunctions = new ArrayList<IRSyntheticFunction>();
        final var lifecycleArtifacts = scanLifecycleArtifacts(ast);
        final var initDecl = lifecycleArtifacts.initDecl();
        final var frameDecl = lifecycleArtifacts.frameDecl();
        final var firstGlobalDecl = lifecycleArtifacts.firstGlobalDecl();

        final var fileInitAnchorSpan = initDecl != null ? initDecl.span() : firstGlobalDecl == null ? null : firstGlobalDecl.span();
        final var fileInitAnchorName = initDecl != null ? initDecl.name() : firstGlobalDecl == null ? "" : firstGlobalDecl.name();
        if (fileInitAnchorSpan != null) {
            syntheticFunctions.add(new IRSyntheticFunction(
                    moduleId,
                    fileInitCallableName(fileId),
                    IRSyntheticCallableKind.FILE_INIT_FRAGMENT,
                    new IRSyntheticOrigin(fileId, fileInitAnchorName, fileInitAnchorSpan)));
            syntheticFunctions.add(new IRSyntheticFunction(
                    moduleId,
                    moduleInitCallableName(moduleId),
                    IRSyntheticCallableKind.MODULE_INIT,
                    new IRSyntheticOrigin(fileId, fileInitAnchorName, fileInitAnchorSpan)));
        }
        if (initDecl != null && frameDecl != null) {
            syntheticFunctions.add(new IRSyntheticFunction(
                    moduleId,
                    projectInitCallableName(moduleId),
                    IRSyntheticCallableKind.PROJECT_INIT,
                    new IRSyntheticOrigin(fileId, initDecl.name(), initDecl.span())));
        }
        if (frameDecl != null) {
            syntheticFunctions.add(new IRSyntheticFunction(
                    moduleId,
                    frameWrapperCallableName(moduleId),
                    IRSyntheticCallableKind.PUBLISHED_FRAME_WRAPPER,
                    new IRSyntheticOrigin(fileId, frameDecl.name(), frameDecl.span())));
        }
        return ReadOnlyList.wrap(syntheticFunctions);
    }

    private IRBackendExecutableFunction lowerFileInitExecutable(
            final FileId fileId,
            final ModuleId moduleId,
            final PbsLifecycleArtifacts lifecycleArtifacts,
            final ArrayList<IRBackendExecutableFunction> executableFunctions,
            final ArrayList<CallableSignatureRef> callableSignatures) {
        if (lifecycleArtifacts.firstGlobalDecl() == null
                && lifecycleArtifacts.initDecl() == null) {
            return null;
        }
        final var syntheticCallableName = fileInitCallableName(fileId);
        final var syntheticCallableId = new CallableId(callableSignatures.size());
        callableSignatures.add(new CallableSignatureRef(moduleId, syntheticCallableName, 0, "() -> unit"));

        final var instructions = new ArrayList<IRBackendExecutableFunction.Instruction>();
        final var moduleInitSource = lifecycleArtifacts.initDecl();
        if (moduleInitSource != null && moduleInitSource != lifecycleArtifacts.frameDecl()) {
            final var initExecutable = resolveLifecycleExecutable(executableFunctions, moduleId, moduleInitSource.name());
            if (initExecutable != null) {
                instructions.add(new IRBackendExecutableFunction.Instruction(
                        IRBackendExecutableFunction.InstructionKind.CALL_FUNC,
                        moduleId,
                        initExecutable.callableName(),
                        initExecutable.callableId(),
                        null,
                        null,
                        0,
                        0,
                        moduleInitSource.span()));
            }
        }
        instructions.add(new IRBackendExecutableFunction.Instruction(
                IRBackendExecutableFunction.InstructionKind.RET,
                "",
                null,
                null,
                lifecycleArtifacts.anchorSpan()));
        final var span = lifecycleArtifacts.anchorSpan();
        return new IRBackendExecutableFunction(
                fileId,
                moduleId,
                syntheticCallableName,
                syntheticCallableId,
                safeToInt(span.getStart()),
                safeToInt(span.getEnd()),
                0,
                0,
                0,
                1,
                ReadOnlyList.wrap(instructions),
                span);
    }

    private IRBackendExecutableFunction resolveLifecycleExecutable(
            final ArrayList<IRBackendExecutableFunction> executableFunctions,
            final ModuleId moduleId,
            final String callableName) {
        for (final var executable : executableFunctions) {
            if (!callableName.equals(executable.callableName())) {
                continue;
            }
            if (!moduleIdsMatch(moduleId, executable.moduleId())) {
                continue;
            }
            return executable;
        }
        return null;
    }

    private boolean moduleIdsMatch(
            final ModuleId left,
            final ModuleId right) {
        final var normalizedLeft = left == null ? ModuleId.none() : left;
        final var normalizedRight = right == null ? ModuleId.none() : right;
        if (normalizedLeft.isNone() && normalizedRight.isNone()) {
            return true;
        }
        return !normalizedLeft.isNone()
                && !normalizedRight.isNone()
                && normalizedLeft.getIndex() == normalizedRight.getIndex();
    }

    private PbsLifecycleArtifacts scanLifecycleArtifacts(final PbsAst.File ast) {
        PbsAst.FunctionDecl initDecl = null;
        PbsAst.FunctionDecl frameDecl = null;
        PbsAst.GlobalDecl firstGlobalDecl = null;
        for (final var topDecl : ast.topDecls()) {
            if (topDecl instanceof PbsAst.FunctionDecl functionDecl) {
                if (functionDecl.lifecycleMarker() == PbsAst.LifecycleMarker.INIT && initDecl == null) {
                    initDecl = functionDecl;
                }
                if (functionDecl.lifecycleMarker() == PbsAst.LifecycleMarker.FRAME && frameDecl == null) {
                    frameDecl = functionDecl;
                }
            }
            if (topDecl instanceof PbsAst.GlobalDecl globalDecl && firstGlobalDecl == null) {
                firstGlobalDecl = globalDecl;
            }
        }
        return new PbsLifecycleArtifacts(initDecl, frameDecl, firstGlobalDecl);
    }

    private static String fileInitCallableName(final FileId fileId) {
        return "__pbs.file_init$f" + fileId.getId();
    }

    public static String moduleInitCallableName(final ModuleId moduleId) {
        return "__pbs.module_init$m" + normalizedModuleIndex(moduleId);
    }

    public static String projectInitCallableName(final ModuleId moduleId) {
        return "__pbs.project_init$m" + normalizedModuleIndex(moduleId);
    }

    public static String frameWrapperCallableName(final ModuleId moduleId) {
        return "__pbs.frame_wrapper$m" + normalizedModuleIndex(moduleId);
    }

    public static String bootGuardGlobalName(final ModuleId moduleId) {
        return "__pbs.boot_guard$m" + normalizedModuleIndex(moduleId);
    }

    private static int normalizedModuleIndex(final ModuleId moduleId) {
        final var normalized = moduleId == null ? ModuleId.none() : moduleId;
        return normalized.isNone() ? -1 : normalized.getIndex();
    }

    private int safeToInt(final long value) {
        if (value > Integer.MAX_VALUE) {
            return Integer.MAX_VALUE;
        }
        if (value < Integer.MIN_VALUE) {
            return Integer.MIN_VALUE;
        }
        return (int) value;
    }

    private record PbsLifecycleArtifacts(
            PbsAst.FunctionDecl initDecl,
            PbsAst.FunctionDecl frameDecl,
            PbsAst.GlobalDecl firstGlobalDecl) {
        private Span anchorSpan() {
            if (initDecl != null) {
                return initDecl.span();
            }
            if (frameDecl != null) {
                return frameDecl.span();
            }
            if (firstGlobalDecl != null) {
                return firstGlobalDecl.span();
            }
            return Span.none();
        }
    }

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

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

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

        final var requiredCapabilities = new HashSet<String>();
        requiredCapabilities.addAll(primary.requiredCapabilities().asList());
        requiredCapabilities.addAll(imported.requiredCapabilities().asList());

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

    public record ImportedCallableSurface(
            ModuleId moduleId,
            String callableName,
            int arity,
            int returnSlots,
            PbsAst.ReturnKind returnKind,
            PbsAst.TypeRef returnType,
            String shapeSurface) {
        public ImportedCallableSurface {
            moduleId = moduleId == null ? ModuleId.none() : moduleId;
        }
    }

    public record ImportedGlobalDependency(
            String localName,
            ModuleId ownerModuleId,
            String globalName) {
        public ImportedGlobalDependency {
            ownerModuleId = ownerModuleId == null ? ModuleId.none() : ownerModuleId;
        }
    }

    public record ImportedGlobalSurface(
            ModuleId ownerModuleId,
            String globalName,
            String localName,
            PbsAst.TypeRef explicitType,
            PbsAst.Expression initializer,
            ReadOnlyList<ImportedGlobalDependency> importedDependencies,
            p.studio.compiler.source.Span span,
            boolean directlyVisible) {
        public ImportedGlobalSurface {
            ownerModuleId = ownerModuleId == null ? ModuleId.none() : ownerModuleId;
            importedDependencies = importedDependencies == null ? ReadOnlyList.empty() : importedDependencies;
        }
    }
}
