package p.studio.compiler.pbs;

import p.studio.compiler.messages.HostAdmissionContext;
import p.studio.compiler.models.IRBackendFile;
import p.studio.compiler.models.IRFunction;
import p.studio.compiler.models.IRReservedMetadata;
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.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.identifiers.FileId;
import p.studio.compiler.source.identifiers.ModuleId;
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;

public final class PbsFrontendCompiler {
    private final PbsFlowSemanticsValidator flowSemanticsValidator = new PbsFlowSemanticsValidator();
    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) {
        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,
                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 NameTable nameTable) {
        return compileParsedFile(
                fileId,
                ast,
                diagnostics,
                sourceKind,
                moduleId,
                modulePool,
                hostAdmissionContext,
                nameTable,
                ReadOnlyList.empty(),
                ReadOnlyList.empty(),
                IRReservedMetadata.empty());
    }

    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 NameTable nameTable,
            final ReadOnlyList<PbsAst.TopDecl> supplementalTopDecls,
            final ReadOnlyList<ImportedCallableSurface> importedCallables,
            final IRReservedMetadata importedReservedMetadata) {
        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 effectiveImportedReservedMetadata = importedReservedMetadata == null
                ? IRReservedMetadata.empty()
                : importedReservedMetadata;
        final var semanticsErrorBaseline = diagnostics.errorCount();
        new PbsDeclarationSemanticsValidator(effectiveNameTable).validate(ast, sourceKind, diagnostics);
        flowSemanticsValidator.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 var loweringErrorBaseline = diagnostics.errorCount();
        final var executableLowering = executableLoweringService.lower(
                fileId,
                ast,
                effectiveSupplementalTopDecls,
                effectiveModuleId,
                mergeReservedMetadata(reservedMetadata, effectiveImportedReservedMetadata),
                effectiveNameTable,
                diagnostics,
                effectiveImportedCallables);
        if (diagnostics.errorCount() > loweringErrorBaseline) {
            return IRBackendFile.empty(fileId);
        }
        return new IRBackendFile(
                fileId,
                functions,
                executableLowering.executableFunctions(),
                reservedMetadata,
                effectiveModulePool,
                executableLowering.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 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,
            String shapeSurface) {
        public ImportedCallableSurface {
            moduleId = moduleId == null ? ModuleId.none() : moduleId;
        }
    }
}
