package main.rice.basegen;

import main.rice.node.APyNode;
import main.rice.obj.APyObj;
import main.rice.test.TestCase;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * Generates test cases by combining values from multiple argument nodes.
 * Supports exhaustive (Cartesian product) and random test generation.
 */
public class BaseSetGenerator {

    private final List<APyNode<?>> nodes;
    private final int numRand;
    private Set<TestCase> currentTests = new HashSet<>();

    public BaseSetGenerator(List<APyNode<?>> nodes, int numRand) {
        this.nodes = nodes;
        this.numRand = numRand;
    }

    /**
     * Generates all exhaustive test cases (Cartesian product of each node's exhaustive values).
     */
    public Set<TestCase> genExTests() {
        // Generate exhaustive values for each node
        List<List<? extends APyObj>> allVals = new ArrayList<>();
        for (APyNode<?> node : nodes) {
            allVals.add(new ArrayList<>(node.genExVals()));
        }

        Set<TestCase> tests = new HashSet<>();
        cartesianProduct(allVals, 0, new ArrayList<>(), tests);

        currentTests = new HashSet<>(tests);
        return tests;
    }

    private void cartesianProduct(List<List<? extends APyObj>> allVals, int idx,
                                  List<APyObj> current, Set<TestCase> result) {
        if (idx == allVals.size()) {
            result.add(new TestCase(new ArrayList<>(current)));
            return;
        }
        for (APyObj val : allVals.get(idx)) {
            current.add(val);
            cartesianProduct(allVals, idx + 1, current, result);
            current.remove(current.size() - 1);
        }
    }

    /**
     * Generates up to numRand random test cases that don't duplicate tests in currentTests.
     */
    public Set<TestCase> genRandTests() {
        Set<TestCase> randTests = new HashSet<>();
        if (numRand == 0) {
            return randTests;
        }

        int maxAttempts = numRand * 1000 + 100;
        int attempts = 0;

        while (randTests.size() < numRand && attempts < maxAttempts) {
            List<APyObj> args = new ArrayList<>();
            for (APyNode<?> node : nodes) {
                args.add(node.genRandVal());
            }
            TestCase tc = new TestCase(args);
            if (!currentTests.contains(tc) && !randTests.contains(tc)) {
                randTests.add(tc);
            }
            attempts++;
        }

        return randTests;
    }

    /**
     * Generates the full base set: all exhaustive tests + numRand random tests (no duplicates).
     */
    public List<TestCase> genBaseSet() {
        Set<TestCase> exTests = genExTests();
        Set<TestCase> randTests = genRandTests();

        List<TestCase> result = new ArrayList<>(exTests);
        result.addAll(randTests);
        return result;
    }
}