package p.studio.compiler.backend.bytecode;

import p.studio.utilities.structures.ReadOnlyList;

import java.util.ArrayList;
import java.util.Objects;

public final class BytecodeFunctionLayoutBuilder {
    private BytecodeFunctionLayoutBuilder() {
    }

    public static LayoutResult build(
            final ReadOnlyList<FunctionFragment> functions) {
        final var inputFunctions = functions == null ? ReadOnlyList.<FunctionFragment>empty() : functions;
        final var metas = new ArrayList<BytecodeModule.FunctionMeta>(inputFunctions.size());
        final var names = new ArrayList<BytecodeModule.FunctionName>(inputFunctions.size());
        final var spans = new ArrayList<BytecodeModule.PcToSpan>();
        var totalCodeLen = 0;
        for (var i = 0; i < inputFunctions.size(); i++) {
            final var fn = inputFunctions.get(i);
            final var fnCode = fn.code();
            metas.add(new BytecodeModule.FunctionMeta(
                    totalCodeLen,
                    fnCode.length,
                    fn.paramSlots(),
                    fn.localSlots(),
                    fn.returnSlots(),
                    fn.maxStackSlots()));
            names.add(new BytecodeModule.FunctionName(i, fn.name()));

            for (final var span : fn.instructionSpans()) {
                if (span.functionPc() < 0 || span.functionPc() >= fnCode.length) {
                    throw new BytecodeMarshalingException(
                            BytecodeMarshalingErrorCode.MARSHAL_FORMAT_PC_SPAN_OUT_OF_BOUNDS,
                            "pc_to_span offset is out of bounds for function: " + fn.name());
                }
                spans.add(new BytecodeModule.PcToSpan(
                        totalCodeLen + span.functionPc(),
                        span.span()));
            }
            totalCodeLen += fnCode.length;
        }

        final var flattened = new byte[totalCodeLen];
        var cursor = 0;
        for (final var fn : inputFunctions) {
            final var fnCode = fn.code();
            System.arraycopy(fnCode, 0, flattened, cursor, fnCode.length);
            cursor += fnCode.length;
        }

        return new LayoutResult(
                flattened,
                ReadOnlyList.wrap(metas),
                new BytecodeModule.DebugInfo(
                        ReadOnlyList.wrap(spans),
                        ReadOnlyList.wrap(names)));
    }

    public record FunctionFragment(
            String name,
            byte[] code,
            int paramSlots,
            int localSlots,
            int returnSlots,
            int maxStackSlots,
            ReadOnlyList<InstructionSpan> instructionSpans) {
        public FunctionFragment {
            Objects.requireNonNull(name, "name");
            code = code == null ? new byte[0] : code.clone();
            instructionSpans = instructionSpans == null ? ReadOnlyList.empty() : instructionSpans;
        }
    }

    public record InstructionSpan(
            int functionPc,
            BytecodeModule.SourceSpan span) {
        public InstructionSpan {
            Objects.requireNonNull(span, "span");
        }
    }

    public record LayoutResult(
            byte[] code,
            ReadOnlyList<BytecodeModule.FunctionMeta> functions,
            BytecodeModule.DebugInfo debugInfo) {
        public LayoutResult {
            code = code == null ? new byte[0] : code.clone();
            functions = functions == null ? ReadOnlyList.empty() : functions;
        }
    }
}

