package p.studio.compiler.pbs.lowering;

import p.studio.compiler.pbs.PbsCallableSlotCounter;
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.*;
import p.studio.utilities.structures.ReadOnlyList;

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

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

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

    PbsExecutableCallableRegistry create(
            final PbsAst.File ast,
            final ModuleId moduleId,
            final NameTable nameTable,
            final Map<String, String> builtinCanonicalBySourceType,
            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 returnOwnerByCallableId = new HashMap<CallableId, String>();
        final var tupleProjectionByCallableId = new HashMap<CallableId, PbsTupleProjection>();
        final var typeSurfaceTable = new TypeSurfaceTable();
        final var callableShapeTable = new CallableShapeTable();
        final var localCallables = collectLocalLowerableCallables(ast);

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

        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,
                returnOwnerByCallableId,
                tupleProjectionByCallableId,
                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 Map<CallableId, String> returnOwnerByCallableId,
            final Map<CallableId, PbsTupleProjection> tupleProjectionByCallableId,
            final TypeSurfaceTable typeSurfaceTable,
            final CallableShapeTable callableShapeTable,
            final ReadOnlyList<PbsLowerableCallable> localCallables,
            final Map<String, String> builtinCanonicalBySourceType) {
        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));
            final var returnOwner = inferCallableReturnOwner(
                    declaredFn.returnKind(),
                    declaredFn.returnType(),
                    builtinCanonicalBySourceType);
            if (!returnOwner.isBlank()) {
                returnOwnerByCallableId.put(callableId, returnOwner);
            }
            final var tupleProjection = tupleProjectionOf(declaredFn.returnKind(), declaredFn.returnType());
            if (tupleProjection != null) {
                tupleProjectionByCallableId.put(callableId, tupleProjection);
            }
        }
    }

    private void registerImportedCallables(
            final NameTable nameTable,
            final CallableTable callableIdTable,
            final Map<PbsCallableResolutionKey, List<CallableId>> callableIdsByNameAndArity,
            final Map<CallableId, Integer> returnSlotsByCallableId,
            final Map<CallableId, PbsTupleProjection> tupleProjectionByCallableId,
            final Map<CallableId, String> returnOwnerByCallableId,
            final ReadOnlyList<PbsFrontendCompiler.ImportedCallableSurface> importedCallables,
            final Map<String, String> builtinCanonicalBySourceType) {
        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());
            final var returnOwner = inferCallableReturnOwner(
                    importedCallable.returnKind(),
                    importedCallable.returnType(),
                    builtinCanonicalBySourceType);
            if (!returnOwner.isBlank()) {
                returnOwnerByCallableId.put(callableId, returnOwner);
            }
            final var tupleProjection = tupleProjectionOf(importedCallable.returnKind(), importedCallable.returnType());
            if (tupleProjection != null) {
                tupleProjectionByCallableId.put(callableId, tupleProjection);
            }
        }
    }

    private String inferCallableReturnOwner(
            final PbsAst.ReturnKind returnKind,
            final PbsAst.TypeRef returnType,
            final Map<String, String> builtinCanonicalBySourceType) {
        if (returnKind != PbsAst.ReturnKind.PLAIN || returnType == null || builtinCanonicalBySourceType == null) {
            return "";
        }
        final var unwrapped = unwrapGroup(returnType);
        if (unwrapped == null || unwrapped.kind() != PbsAst.TypeRefKind.SIMPLE) {
            return "";
        }
        return builtinCanonicalBySourceType.getOrDefault(unwrapped.name(), "");
    }

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

    private PbsTupleProjection tupleProjectionOf(
            final PbsAst.ReturnKind returnKind,
            final PbsAst.TypeRef returnType) {
        if (returnKind != PbsAst.ReturnKind.PLAIN) {
            return null;
        }
        final var unwrapped = unwrapGroup(returnType);
        if (unwrapped == null || unwrapped.kind() != PbsAst.TypeRefKind.NAMED_TUPLE) {
            return null;
        }
        final var fieldsByLabel = new HashMap<String, PbsTupleFieldProjection>();
        var slotOffset = 0;
        for (final var field : unwrapped.fields()) {
            final var fieldSlotCount = Math.max(1, PbsCallableSlotCounter.returnSlots(PbsAst.ReturnKind.PLAIN, field.typeRef()));
            if (field.label() != null && !field.label().isBlank()) {
                fieldsByLabel.put(field.label(), new PbsTupleFieldProjection(field.label(), slotOffset, fieldSlotCount));
            }
            slotOffset += fieldSlotCount;
        }
        return new PbsTupleProjection(slotOffset, Map.copyOf(fieldsByLabel));
    }

    private PbsAst.TypeRef unwrapGroup(final PbsAst.TypeRef typeRef) {
        if (typeRef == null) {
            return null;
        }
        if (typeRef.kind() != PbsAst.TypeRefKind.GROUP) {
            return typeRef;
        }
        return unwrapGroup(typeRef.inner());
    }
}
