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.compiler.*;
import io.gitlab.jfronny.muscript.data.Scope;
import io.gitlab.jfronny.muscript.data.dynamic.*;
import io.gitlab.jfronny.muscript.data.dynamic.additional.DFinal;
import io.gitlab.jfronny.muscript.error.LocationalException;

import java.io.IOException;
import java.util.List;
import java.util.stream.IntStream;
import java.util.stream.Stream;

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(Order.Primary, 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);
        }, this::toString);
    }

    @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 void decompile(ExprWriter writer) throws IOException {
        writer.append("{ ");
        for (int i = 0; i < boundArgs.size(); i++) {
            String arg = boundArgs.get(i);
            if (!Lexer.isValidId(arg)) throw new IllegalArgumentException("Not a valid argument name: " + arg);
            writer.append(arg);
            if (i == boundArgs.size() - 1 && variadic) writer.append("...");
            writer.append(' ');
        }
        writer.append("->").increaseIndent();
        for (Expr<?> expr : stream().toList()) {
            writer.append("\n");
            expr.decompile(writer);
            writer.append(";");
        }
        writer.decreaseIndent().append("\n}");
    }

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

    public Stream<Expr<?>> stream() {
        return Stream.concat(steps.stream(), Stream.of(fin));
    }
}