package io.gitlab.jfronny.muscript.ast.dynamic;

import io.gitlab.jfronny.muscript.ast.DynamicExpr;
import io.gitlab.jfronny.muscript.ast.Expr;
import io.gitlab.jfronny.muscript.data.Scope;
import io.gitlab.jfronny.muscript.data.dynamic.*;
import io.gitlab.jfronny.muscript.error.LocationalException;

import java.util.List;
import java.util.stream.IntStream;

public class Closure extends DynamicExpr {
    private final List<String> boundArgs;
    private final List<Expr<?>> steps;
    private final DynamicExpr fin;
    private final boolean variadic;

    public Closure(int chStart, int chEnd, List<String> boundArgs, List<Expr<?>> expressions, boolean variadic) {
        this(chStart, chEnd, boundArgs, expressions.subList(0, expressions.size() - 1), expressions.get(expressions.size() - 1).asDynamicExpr(), variadic);
    }

    private Closure(int chStart, int chEnd, List<String> boundArgs, List<Expr<?>> steps, DynamicExpr fin, boolean variadic) {
        super(chStart, chEnd);
        this.boundArgs = List.copyOf(boundArgs);
        this.steps = List.copyOf(steps);
        this.fin = fin;
        this.variadic = variadic;
    }

    @Override
    public DCallable get(Scope dataRoot) {
        return DFinal.of(args -> {
            int ac = args.size();
            int ae = boundArgs.size();
            if (variadic) ae--;
            if (ac < ae) throw new LocationalException(chStart, chEnd, "Invoked with too few arguments (expected " + ae + " but got " + ac + ")");
            if (!variadic && ac > ae) throw new LocationalException(chStart, chEnd, "Invoked with too many arguments (expected " + ae + " but got " + ac + ")");
            Scope fork = dataRoot.fork();
            for (int i = 0; i < ae; i++) fork.set(boundArgs.get(i), args.get(i));
            if (variadic) {
                fork.set(boundArgs.get(boundArgs.size() - 1), IntStream.range(ae, ac)
                        .<Dynamic<?>>mapToObj(args::get)
                        .toList());
            }
            for (Expr<?> step : steps) {
                step.get(fork);
            }
            return fin.get(fork);
        });
    }

    @Override
    public DynamicExpr optimize() {
        // Eliminating side effect free steps might be possible, but would be too much work imo
        return new Closure(chStart, chEnd, boundArgs, steps.stream().<Expr<?>>map(Expr::optimize).toList(), fin.optimize(), variadic);
    }

    @Override
    public boolean equals(Object obj) {
        return obj instanceof Closure closure
                && boundArgs.equals(closure.boundArgs)
                && steps.equals(closure.steps)
                && fin.equals(closure.fin);
    }
}