package p.studio.compiler.pbs.lowering;

import p.studio.compiler.models.IRBackendExecutableFunction;
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, 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<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();
    }

    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));
    }

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

    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++;
    }
}
