package p.studio.compiler.pbs.lowering;

import p.studio.compiler.models.IRBackendExecutableFunction;
import p.studio.compiler.messages.Addressable;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.identifiers.NameId;
import p.studio.compiler.source.tables.IntrinsicTable;
import p.studio.compiler.source.tables.NameTable;

import java.util.*;

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

final class PbsExecutableLoweringContext {
    private final DiagnosticSink diagnostics;
    private final NameTable nameTable;
    private final PbsExecutableMetadataIndex metadataIndex;
    private final PbsExecutableCallableRegistry callableRegistry;
    private final IntrinsicTable intrinsicIdTable;
    private final Map<NameId, Integer> localSlotByNameId;
    private final Map<NameId, PbsTupleProjection> localTupleProjectionByNameId = new HashMap<>();
    private final Map<NameId, String> localOwnerByNameId;
    private final ArrayList<IRBackendExecutableFunction.Instruction> instructions = new ArrayList<>();
    private final ArrayDeque<PbsLoopTargets> loopTargets = new ArrayDeque<>();
    private int nextLabelId = 0;
    private int nextLocalSlot;

    PbsExecutableLoweringContext(
            final DiagnosticSink diagnostics,
            final NameTable nameTable,
            final PbsExecutableMetadataIndex metadataIndex,
            final PbsExecutableCallableRegistry callableRegistry,
            final IntrinsicTable intrinsicIdTable,
            final Map<NameId, Integer> localSlotByNameId,
            final Map<NameId, String> localOwnerByNameId,
            final int initialLocalSlot) {
        this.diagnostics = diagnostics;
        this.nameTable = nameTable;
        this.metadataIndex = metadataIndex;
        this.callableRegistry = callableRegistry;
        this.intrinsicIdTable = intrinsicIdTable;
        this.localSlotByNameId = localSlotByNameId == null ? new HashMap<>() : localSlotByNameId;
        this.localOwnerByNameId = localOwnerByNameId == null ? new HashMap<>() : localOwnerByNameId;
        this.nextLocalSlot = Math.max(0, initialLocalSlot);
    }

    DiagnosticSink diagnostics() {
        return diagnostics;
    }

    NameTable nameTable() {
        return nameTable;
    }

    Map<NameId, List<p.studio.compiler.models.IRReservedMetadata.HostMethodBinding>> hostByMethodName() {
        return metadataIndex.hostByMethodName();
    }

    Map<NameId, String> builtinConstOwnerByNameId() {
        return metadataIndex.builtinConstOwnerByNameId();
    }

    Map<NameId, Integer> globalSlotByNameId() {
        return metadataIndex.globalSlotByNameId();
    }

    Map<NameId, p.studio.compiler.pbs.ast.PbsAst.ConstDecl> constDeclByNameId() {
        return metadataIndex.constDeclByNameId();
    }

    Map<String, Addressable> addressableByAddress() {
        return metadataIndex.addressableByAddress();
    }

    Map<String, String> builtinCanonicalBySourceType() {
        return metadataIndex.builtinCanonicalBySourceType();
    }

    Map<PbsIntrinsicOwnerMethodKey, List<p.studio.compiler.models.IRReservedMetadata.IntrinsicSurface>> intrinsicByOwnerAndMethod() {
        return metadataIndex.intrinsicByOwnerAndMethod();
    }

    Map<PbsIntrinsicCanonicalKey, String> intrinsicReturnOwnerByCanonical() {
        return metadataIndex.intrinsicReturnOwnerByCanonical();
    }

    Map<PbsCallableResolutionKey, List<p.studio.compiler.source.identifiers.CallableId>> callableIdsByNameAndArity() {
        return callableRegistry.callableIdsByNameAndArity();
    }

    Map<p.studio.compiler.source.identifiers.CallableId, p.studio.compiler.source.tables.CallableSignatureRef> callableSignatureByCallableId() {
        return callableRegistry.callableSignatureByCallableId();
    }

    Map<p.studio.compiler.source.identifiers.CallableId, Integer> returnSlotsByCallableId() {
        return callableRegistry.returnSlotsByCallableId();
    }

    Map<p.studio.compiler.source.identifiers.CallableId, String> returnOwnerByCallableId() {
        return callableRegistry.returnOwnerByCallableId();
    }

    Map<p.studio.compiler.source.identifiers.CallableId, PbsTupleProjection> tupleProjectionByCallableId() {
        return callableRegistry.tupleProjectionByCallableId();
    }

    IntrinsicTable intrinsicIdTable() {
        return intrinsicIdTable;
    }

    Map<NameId, Integer> localSlotByNameId() {
        return localSlotByNameId;
    }

    String resolveLocalOwner(final String localName) {
        if (localName == null || localName.isBlank()) {
            return "";
        }
        return localOwnerByNameId.getOrDefault(nameTable.register(localName), "");
    }

    void bindLocalOwner(
            final String localName,
            final String ownerCanonicalName) {
        if (localName == null || localName.isBlank()) {
            return;
        }
        final var nameId = nameTable.register(localName);
        if (ownerCanonicalName == null || ownerCanonicalName.isBlank()) {
            localOwnerByNameId.remove(nameId);
            return;
        }
        localOwnerByNameId.put(nameId, ownerCanonicalName);
    }

    Integer resolveLocalSlot(final String localName) {
        if (localName == null || localName.isBlank()) {
            return null;
        }
        return localSlotByNameId.get(nameTable.register(localName));
    }

    PbsTupleProjection resolveLocalTupleProjection(final String localName) {
        if (localName == null || localName.isBlank()) {
            return null;
        }
        return localTupleProjectionByNameId.get(nameTable.register(localName));
    }

    Integer resolveGlobalSlot(final String globalName) {
        if (globalName == null || globalName.isBlank()) {
            return null;
        }
        return globalSlotByNameId().get(nameTable.register(globalName));
    }

    p.studio.compiler.pbs.ast.PbsAst.ConstDecl resolveConstDecl(final String constName) {
        if (constName == null || constName.isBlank()) {
            return null;
        }
        return constDeclByNameId().get(nameTable.register(constName));
    }

    Integer resolveAssetId(final String address) {
        if (address == null || address.isBlank()) {
            return null;
        }
        final var addressable = addressableByAddress().get(address);
        return addressable == null ? null : addressable.assetId();
    }

    int declareLocalSlot(final String localName) {
        return declareLocalSlots(localName, 1);
    }

    int declareLocalSlots(
            final String localName,
            final int slotCount) {
        final var nameId = nameTable.register(localName);
        final var existing = localSlotByNameId.get(nameId);
        if (existing != null) {
            return existing;
        }
        final var slot = nextLocalSlot;
        nextLocalSlot += Math.max(1, slotCount);
        localSlotByNameId.put(nameId, slot);
        return slot;
    }

    void bindLocalTupleProjection(
            final String localName,
            final PbsTupleProjection tupleProjection) {
        if (localName == null || localName.isBlank()) {
            return;
        }
        final var nameId = nameTable.register(localName);
        if (tupleProjection == null) {
            localTupleProjectionByNameId.remove(nameId);
            return;
        }
        localTupleProjectionByNameId.put(nameId, tupleProjection);
    }

    int nextLocalSlot() {
        return nextLocalSlot;
    }

    ArrayList<IRBackendExecutableFunction.Instruction> instructions() {
        return instructions;
    }

    ArrayDeque<PbsLoopTargets> loopTargets() {
        return loopTargets;
    }

    void pushLoop(
            final String continueLabel,
            final String breakLabel) {
        loopTargets.push(new PbsLoopTargets(continueLabel, breakLabel));
    }

    void popLoop() {
        loopTargets.pop();
    }

    String nextLabel(final String prefix) {
        return "__" + prefix + "_" + nextLabelId++;
    }
}
