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.data.dynamic.type.DTypeCallable;
import io.gitlab.jfronny.muscript.error.LocationalException;

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

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

    public Closure(CodeLocation location, List<String> boundArgs, List<Expr<?>> expressions, boolean variadic) {
        this(location, boundArgs, expressions.subList(0, expressions.size() - 1), expressions.getLast().asDynamicExpr(), variadic);
    }

    private Closure(CodeLocation location, List<String> boundArgs, List<Expr<?>> steps, DynamicExpr fin, boolean variadic) {
        super(Order.Primary, location);
        this.boundArgs = List.copyOf(boundArgs);
        this.steps = List.copyOf(steps);
        this.fin = fin;
        this.variadic = variadic;
    }

    private DTypeCallable getSignature() {
        //TODO take into account inner instructions to properly determine type
        List<DTypeCallable.Arg> args = new LinkedList<>();
        for (int i = 0; i < boundArgs.size(); i++) {
            args.add(new DTypeCallable.Arg(boundArgs.get(i), null, variadic && (i == boundArgs.size() - 1)));
        }
        return new DTypeCallable(args, null);
    }

    @Override
    public DCallable get(Scope dataRoot) {
        return DFinal.of(getSignature(), null, args -> {
            // Compare with ExprGroup.get()
            int ac = args.size();
            int ae = boundArgs.size();
            if (variadic) ae--;
            if (ac < ae) throw new LocationalException(location, "Invoked with too few arguments (expected " + ae + " but got " + ac + ")");
            if (!variadic && ac > ae) throw new LocationalException(location, "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.getLast(), IntStream.range(ae, ac)
                        .mapToObj(args::get)
                        .toList());
            }
            for (Expr<?> step : steps) {
                step.get(fork);
            }
            return fin.get(fork);
        }, () -> this);
    }

    @Override
    public DynamicExpr optimize() {
        return new Closure(
                location,
                boundArgs,
                steps.stream()
                        .map(Expr::optimize)
                        .flatMap(Expr::extractSideEffects)
                        .<Expr<?>>map(Expr::optimize)
                        .toList(),
                fin.optimize(),
                variadic
        );
    }

    @Override
    public Stream<Expr<?>> extractSideEffects() {
        return Stream.empty();
    }

    @Override
    public void decompile(ExprWriter writer) throws IOException {
        writer.append("{");
        boolean first = true;
        for (int i = 0; i < boundArgs.size(); i++) {
            if (!first) writer.append(",");
            first = false;
            writer.append(' ').appendLiteral(boundArgs.get(i));
            if (i == boundArgs.size() - 1 && variadic) 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));
    }
}