package p.studio.compiler.pbs.lowering;

import p.studio.compiler.pbs.PbsFrontendCompiler;
import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.source.identifiers.CallableId;
import p.studio.compiler.source.identifiers.ModuleId;
import p.studio.compiler.source.tables.CallableShapeTable;
import p.studio.compiler.source.tables.CallableSignatureRef;
import p.studio.compiler.source.tables.CallableTable;
import p.studio.compiler.source.tables.NameTable;
import p.studio.compiler.source.tables.TypeSurfaceTable;
import p.studio.utilities.structures.ReadOnlyList;

import static p.studio.compiler.pbs.lowering.PbsExecutableLoweringModels.*;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

final class PbsExecutableCallableRegistryFactory {
    private final PbsCallableShapeSurfaceService callableShapeSurfaceService = new PbsCallableShapeSurfaceService();

    PbsExecutableCallableRegistry create(
            final PbsAst.File ast,
            final ModuleId moduleId,
            final NameTable nameTable,
            final ReadOnlyList<PbsFrontendCompiler.ImportedCallableSurface> importedCallables) {
        final var callableIdTable = new CallableTable();
        final var callableIdsByNameAndArity = new HashMap<PbsCallableResolutionKey, List<CallableId>>();
        final var callableIdByDeclaration = new HashMap<PbsLowerableCallable, CallableId>();
        final var returnSlotsByCallableId = new HashMap<CallableId, Integer>();
        final var typeSurfaceTable = new TypeSurfaceTable();
        final var callableShapeTable = new CallableShapeTable();
        final var localCallables = collectLocalLowerableCallables(ast);

        registerLocalCallables(
                moduleId,
                nameTable,
                callableIdTable,
                callableIdsByNameAndArity,
                callableIdByDeclaration,
                returnSlotsByCallableId,
                typeSurfaceTable,
                callableShapeTable,
                localCallables);
        registerImportedCallables(
                nameTable,
                callableIdTable,
                callableIdsByNameAndArity,
                returnSlotsByCallableId,
                importedCallables);

        final var callableSignatureByCallableId = new HashMap<CallableId, CallableSignatureRef>();
        final var callableSignatures = new ArrayList<CallableSignatureRef>(callableIdTable.size());
        for (final var callableId : callableIdTable.identifiers()) {
            final var signature = callableIdTable.get(callableId);
            callableSignatureByCallableId.put(callableId, signature);
            callableSignatures.add(signature);
        }
        return new PbsExecutableCallableRegistry(
                localCallables,
                callableIdByDeclaration,
                callableIdsByNameAndArity,
                callableSignatureByCallableId,
                returnSlotsByCallableId,
                ReadOnlyList.wrap(callableSignatures));
    }

    private ReadOnlyList<PbsLowerableCallable> collectLocalLowerableCallables(final PbsAst.File ast) {
        final var callables = new ArrayList<PbsLowerableCallable>();
        for (final var functionDecl : ast.functions()) {
            callables.add(new PbsLowerableCallable(functionDecl.name(), functionDecl));
        }
        for (final var topDecl : ast.topDecls()) {
            if (!(topDecl instanceof PbsAst.ServiceDecl serviceDecl)) {
                continue;
            }
            for (final var method : serviceDecl.methods()) {
                callables.add(new PbsLowerableCallable(serviceDecl.name() + "." + method.name(), method));
            }
        }
        return ReadOnlyList.wrap(callables);
    }

    private void registerLocalCallables(
            final ModuleId moduleId,
            final NameTable nameTable,
            final CallableTable callableIdTable,
            final Map<PbsCallableResolutionKey, List<CallableId>> callableIdsByNameAndArity,
            final Map<PbsLowerableCallable, CallableId> callableIdByDeclaration,
            final Map<CallableId, Integer> returnSlotsByCallableId,
            final TypeSurfaceTable typeSurfaceTable,
            final CallableShapeTable callableShapeTable,
            final ReadOnlyList<PbsLowerableCallable> localCallables) {
        for (final var declaredCallable : localCallables) {
            final var declaredFn = declaredCallable.functionDecl();
            final var callableShapeId = callableShapeSurfaceService.callableShapeId(
                    declaredFn,
                    typeSurfaceTable,
                    callableShapeTable);
            final var callableId = callableIdTable.register(
                    moduleId,
                    declaredCallable.callableName(),
                    declaredFn.parameters().size(),
                    callableShapeSurfaceService.callableShapeSurface(callableShapeId, typeSurfaceTable, callableShapeTable));
            callableIdByDeclaration.put(declaredCallable, callableId);
            callableIdsByNameAndArity
                    .computeIfAbsent(
                            new PbsCallableResolutionKey(nameTable.register(declaredCallable.callableName()), declaredFn.parameters().size()),
                            ignored -> new ArrayList<>())
                    .add(callableId);
            returnSlotsByCallableId.put(callableId, returnSlotsFor(declaredFn));
        }
    }

    private void registerImportedCallables(
            final NameTable nameTable,
            final CallableTable callableIdTable,
            final Map<PbsCallableResolutionKey, List<CallableId>> callableIdsByNameAndArity,
            final Map<CallableId, Integer> returnSlotsByCallableId,
            final ReadOnlyList<PbsFrontendCompiler.ImportedCallableSurface> importedCallables) {
        for (final var importedCallable : importedCallables) {
            final var callableId = callableIdTable.register(
                    importedCallable.moduleId(),
                    importedCallable.callableName(),
                    importedCallable.arity(),
                    importedCallable.shapeSurface());
            callableIdsByNameAndArity
                    .computeIfAbsent(
                            new PbsCallableResolutionKey(
                                    nameTable.register(importedCallable.callableName()),
                                    importedCallable.arity()),
                            ignored -> new ArrayList<>())
                    .add(callableId);
            returnSlotsByCallableId.put(callableId, importedCallable.returnSlots());
        }
    }

    private int returnSlotsFor(final PbsAst.FunctionDecl functionDecl) {
        return switch (functionDecl.returnKind()) {
            case INFERRED_UNIT, EXPLICIT_UNIT -> 0;
            case PLAIN, RESULT -> 1;
        };
    }
}
