package p.studio.compiler.models;

import lombok.Builder;
import lombok.Getter;
import p.studio.compiler.source.identifiers.CallableId;
import p.studio.compiler.source.identifiers.IntrinsicId;
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.LinkedHashSet;

@Builder
@Getter
public class IRBackend implements IRBackendReader {
    @Builder.Default
    private final String entryPointCallableName = "main";
    @Builder.Default
    private final ModuleId entryPointModuleId = ModuleId.none();
    @Builder.Default
    private final ReadOnlyList<IRFunction> functions = ReadOnlyList.empty();
    @Builder.Default
    private final ReadOnlyList<IRSyntheticFunction> syntheticFunctions = ReadOnlyList.empty();
    @Builder.Default
    private final ReadOnlyList<IRGlobal> globals = ReadOnlyList.empty();
    @Builder.Default
    private final ReadOnlyList<IRBackendExecutableFunction> executableFunctions = ReadOnlyList.empty();
    @Builder.Default
    private final ReadOnlyList<ModuleReference> modulePool = ReadOnlyList.empty();
    @Builder.Default
    private final ReadOnlyList<CallableSignatureRef> callableSignatures = ReadOnlyList.empty();
    @Builder.Default
    private final ReadOnlyList<IntrinsicReference> intrinsicPool = ReadOnlyList.empty();
    @Builder.Default
    private final IRReservedMetadata reservedMetadata = IRReservedMetadata.empty();

    public static IRBackendAggregator aggregator() {
        return new IRBackendAggregator();
    }

    public static final class IRBackendAggregator {
        private String entryPointCallableName;
        private ModuleId entryPointModuleId = ModuleId.none();
        private final ArrayList<IRFunction> functions = new ArrayList<>();
        private final ArrayList<IRSyntheticFunction> syntheticFunctions = new ArrayList<>();
        private final ArrayList<IRGlobal> globals = new ArrayList<>();
        private final ArrayList<IRBackendExecutableFunction> executableFunctions = new ArrayList<>();
        private final ModuleTable moduleTable = new ModuleTable();
        private final CallableTable callableTable = new CallableTable();
        private final IntrinsicTable intrinsicTable = new IntrinsicTable();
        private final ArrayList<IRReservedMetadata.HostMethodBinding> hostMethodBindings = new ArrayList<>();
        private final ArrayList<IRReservedMetadata.BuiltinTypeSurface> builtinTypeSurfaces = new ArrayList<>();
        private final ArrayList<IRReservedMetadata.BuiltinConstSurface> builtinConstSurfaces = new ArrayList<>();
        private final LinkedHashSet<String> requiredCapabilities = new LinkedHashSet<>();

        public void merge(final IRBackendFile backendFile) {
            if (backendFile == null) {
                return;
            }
            functions.addAll(backendFile.functions().asList());
            final var moduleRemap = reindexModules(backendFile.modulePool());
            for (final var syntheticFunction : backendFile.syntheticFunctions()) {
                syntheticFunctions.add(new IRSyntheticFunction(
                        remapModuleId(syntheticFunction.moduleId(), moduleRemap, "synthetic function"),
                        syntheticFunction.callableName(),
                        syntheticFunction.kind(),
                        syntheticFunction.origin()));
            }
            for (final var global : backendFile.globals()) {
                globals.add(new IRGlobal(
                        global.fileId(),
                        remapModuleId(global.moduleId(), moduleRemap, "global"),
                        global.name(),
                        global.declaredTypeSurface(),
                        global.slot(),
                        global.visibility(),
                        global.isHidden(),
                        global.hiddenKind(),
                        new IRGlobalOrigin(
                                global.origin().derivedFromFile(),
                                remapModuleId(global.origin().derivedFromModule(), moduleRemap, "global origin"),
                                global.origin().derivedFromUserSymbol(),
                                global.origin().primarySpan(),
                                global.origin().fallbackSyntheticLabel()),
                        global.span()));
            }
            final var callableRemap = reindexCallables(backendFile.callableSignatures(), moduleRemap);
            final var intrinsicRemap = reindexIntrinsics(backendFile.intrinsicPool());
            for (final var function : backendFile.executableFunctions()) {
                executableFunctions.add(remapExecutableFunction(function, moduleRemap, callableRemap, intrinsicRemap));
            }
            final var metadata = backendFile.reservedMetadata();
            hostMethodBindings.addAll(metadata.hostMethodBindings().asList());
            builtinTypeSurfaces.addAll(metadata.builtinTypeSurfaces().asList());
            builtinConstSurfaces.addAll(metadata.builtinConstSurfaces().asList());
            requiredCapabilities.addAll(metadata.requiredCapabilities().asList());
        }

        public IRBackendAggregator entryPointCallableName(final String callableName) {
            if (callableName != null && !callableName.isBlank()) {
                entryPointCallableName = callableName;
            }
            return this;
        }

        public IRBackendAggregator entryPointModuleId(final ModuleId moduleId) {
            if (moduleId != null) {
                entryPointModuleId = moduleId;
            }
            return this;
        }

        private ModuleId[] reindexModules(final ReadOnlyList<ModuleReference> localModulePool) {
            if (localModulePool == null || localModulePool.isEmpty()) {
                return new ModuleId[0];
            }
            final var remap = new ModuleId[localModulePool.size()];
            for (var i = 0; i < localModulePool.size(); i++) {
                remap[i] = moduleTable.register(localModulePool.get(i));
            }
            return remap;
        }

        private CallableId[] reindexCallables(
                final ReadOnlyList<CallableSignatureRef> localCallableSignatures,
                final ModuleId[] moduleRemap) {
            if (localCallableSignatures == null || localCallableSignatures.isEmpty()) {
                return new CallableId[0];
            }
            final var remap = new CallableId[localCallableSignatures.size()];
            for (var i = 0; i < localCallableSignatures.size(); i++) {
                final var signature = localCallableSignatures.get(i);
                final var remappedModuleId = remapModuleId(
                        signature.moduleId(),
                        moduleRemap,
                        "callable signature");
                remap[i] = callableTable.register(
                        remappedModuleId,
                        signature.callableName(),
                        signature.arity(),
                        signature.typeShape());
            }
            return remap;
        }

        private IRBackendExecutableFunction remapExecutableFunction(
                final IRBackendExecutableFunction function,
                final ModuleId[] moduleRemap,
                final CallableId[] callableRemap,
                final IntrinsicId[] intrinsicRemap) {
            final var remappedInstructions = new ArrayList<IRBackendExecutableFunction.Instruction>(function.instructions().size());
            final var remappedFunctionModuleId = remapModuleId(
                    function.moduleId(),
                    moduleRemap,
                    "function");
            for (final var instruction : function.instructions()) {
                final CallableId remappedCallee;
                if (instruction.calleeCallableId() == null) {
                    remappedCallee = null;
                } else {
                    remappedCallee = remapCallableId(instruction.calleeCallableId(), callableRemap, "callee");
                }
                final IRBackendExecutableFunction.IntrinsicCallMetadata remappedIntrinsic;
                if (instruction.kind() == IRBackendExecutableFunction.InstructionKind.CALL_INTRINSIC && instruction.intrinsicCall() != null) {
                    final var intrinsic = instruction.intrinsicCall();
                    remappedIntrinsic = new IRBackendExecutableFunction.IntrinsicCallMetadata(
                            intrinsic.canonicalName(),
                            intrinsic.canonicalVersion(),
                            remapIntrinsicId(intrinsic.intrinsicId(), intrinsicRemap));
                } else {
                    remappedIntrinsic = instruction.intrinsicCall();
                }
                final var remappedCalleeModuleId = remapModuleId(
                        instruction.calleeModuleId(),
                        moduleRemap,
                        "callee");
                remappedInstructions.add(new IRBackendExecutableFunction.Instruction(
                        instruction.kind(),
                        remappedCalleeModuleId,
                        instruction.calleeCallableName(),
                        remappedCallee,
                        instruction.hostCall(),
                        remappedIntrinsic,
                        instruction.label(),
                        instruction.targetLabel(),
                        instruction.expectedArgSlots(),
                        instruction.expectedRetSlots(),
                        instruction.span()));
            }
            return new IRBackendExecutableFunction(
                    function.fileId(),
                    remappedFunctionModuleId,
                    function.callableName(),
                    remapCallableId(function.callableId(), callableRemap, "function"),
                    function.sourceStart(),
                    function.sourceEnd(),
                    function.paramSlots(),
                    function.localSlots(),
                    function.returnSlots(),
                    function.maxStackSlots(),
                    ReadOnlyList.wrap(remappedInstructions),
                    function.span());
        }

        private ModuleId remapModuleId(
                final ModuleId localModuleId,
                final ModuleId[] moduleRemap,
                final String label) {
            if (localModuleId != null && !localModuleId.isNone()) {
                final var localIndex = localModuleId.getIndex();
                if (localIndex < 0 || localIndex >= moduleRemap.length) {
                    throw new IllegalArgumentException("invalid local " + label + " module id: " + localModuleId);
                }
                return moduleRemap[localIndex];
            }
            return ModuleId.none();
        }

        private CallableId remapCallableId(
                final CallableId localCallableId,
                final CallableId[] callableRemap,
                final String label) {
            final var localIndex = localCallableId.getIndex();
            if (localIndex < 0 || localIndex >= callableRemap.length) {
                throw new IllegalArgumentException("invalid local " + label + " callable id: " + localCallableId);
            }
            return callableRemap[localIndex];
        }

        private IntrinsicId[] reindexIntrinsics(final ReadOnlyList<IntrinsicReference> localIntrinsicPool) {
            if (localIntrinsicPool == null || localIntrinsicPool.isEmpty()) {
                return new IntrinsicId[0];
            }
            final var remap = new IntrinsicId[localIntrinsicPool.size()];
            for (var i = 0; i < localIntrinsicPool.size(); i++) {
                remap[i] = intrinsicTable.register(localIntrinsicPool.get(i));
            }
            return remap;
        }

        private IntrinsicId remapIntrinsicId(
                final IntrinsicId localIntrinsicId,
                final IntrinsicId[] intrinsicRemap) {
            final var localIndex = localIntrinsicId.getIndex();
            if (localIndex < 0 || localIndex >= intrinsicRemap.length) {
                throw new IllegalArgumentException("invalid local intrinsic id: " + localIntrinsicId);
            }
            return intrinsicRemap[localIndex];
        }

        private ReadOnlyList<CallableSignatureRef> emitCallableSignatures() {
            final var signatures = new ArrayList<CallableSignatureRef>(callableTable.size());
            for (final var callableId : callableTable.identifiers()) {
                signatures.add(callableTable.get(callableId));
            }
            return ReadOnlyList.wrap(signatures);
        }

        private ReadOnlyList<ModuleReference> emitModulePool() {
            final var modules = new ArrayList<ModuleReference>(moduleTable.size());
            for (final var moduleId : moduleTable.identifiers()) {
                modules.add(moduleTable.get(moduleId));
            }
            return ReadOnlyList.wrap(modules);
        }

        private ReadOnlyList<IntrinsicReference> emitIntrinsicPool() {
            final var pool = new ArrayList<IntrinsicReference>(intrinsicTable.size());
            for (final var intrinsicId : intrinsicTable.identifiers()) {
                pool.add(intrinsicTable.get(intrinsicId));
            }
            return ReadOnlyList.wrap(pool);
        }

        public IRBackend emit() {
            return IRBackend
                    .builder()
                    .entryPointCallableName(resolveEntryPointCallableName())
                    .entryPointModuleId(resolveEntryPointModuleId())
                    .functions(ReadOnlyList.wrap(functions))
                    .syntheticFunctions(ReadOnlyList.wrap(syntheticFunctions))
                    .globals(ReadOnlyList.wrap(globals))
                    .executableFunctions(ReadOnlyList.wrap(executableFunctions))
                    .modulePool(emitModulePool())
                    .callableSignatures(emitCallableSignatures())
                    .intrinsicPool(emitIntrinsicPool())
                    .reservedMetadata(new IRReservedMetadata(
                            ReadOnlyList.wrap(hostMethodBindings),
                            ReadOnlyList.wrap(builtinTypeSurfaces),
                            ReadOnlyList.wrap(builtinConstSurfaces),
                            ReadOnlyList.wrap(requiredCapabilities.stream().toList())))
                    .build();
        }

        private String resolveEntryPointCallableName() {
            return entryPointCallableName == null || entryPointCallableName.isBlank()
                    ? "main"
                    : entryPointCallableName;
        }

        private ModuleId resolveEntryPointModuleId() {
            return entryPointModuleId == null ? ModuleId.none() : entryPointModuleId;
        }
    }

    @Override
    public String toString() {
        final var sb = new StringBuilder();
        sb.append("IRBackend{entrypoint=").append(entryPointCallableName)
                .append('@')
                .append(entryPointModuleId == null || entryPointModuleId.isNone() ? "-" : entryPointModuleId.getIndex())
                .append(", functions=").append(functions.size())
                .append(", syntheticFunctions=").append(syntheticFunctions.size())
                .append(", globals=").append(globals.size())
                .append(", executableFunctions=").append(executableFunctions.size())
                .append(", modulePool=").append(modulePool.size())
                .append(", callableSignatures=").append(callableSignatures.size())
                .append(", intrinsicPool=").append(intrinsicPool.size())
                .append(", hostBindings=").append(reservedMetadata.hostMethodBindings().size())
                .append(", builtinTypes=").append(reservedMetadata.builtinTypeSurfaces().size())
                .append(", builtinConsts=").append(reservedMetadata.builtinConstSurfaces().size())
                .append(", requiredCapabilities=").append(reservedMetadata.requiredCapabilities().size())
                .append('}');

        if (functions.isEmpty()) {
            return sb.toString();
        }

        sb.append(System.lineSeparator());
        for (final var function : functions) {
            sb.append("  - fn ")
                    .append(function.name())
                    .append('(')
                    .append(function.parameterCount())
                    .append(" params)")
                    .append(" returnType=")
                    .append(function.hasExplicitReturnType() ? "explicit" : "inferred")
                    .append(" fileId=")
                    .append(function.fileId())
                    .append(" span=[")
                    .append(function.span().getStart())
                    .append(',')
                    .append(function.span().getEnd())
                    .append(']')
                    .append(System.lineSeparator());
        }

        return sb.toString().trim();
    }

}
