/**
 * Unit test for kafka-0009: GraphGraceSearchUtil diamond recursion O(2^D) blowup.
 *
 * Verifies that:
 *   1. The patched findAndVerifyWindowGrace completes in O(N) time on a D-deep diamond graph.
 *   2. The correct grace period (maximum over all ancestors) is returned.
 *   3. TopologyException is still thrown when no GracePeriodGraphNode is reachable.
 *
 * Compile and run:
 *   javac -d unit KafkaGraphGraceSearchDiamondTest.java
 *   java -ea -cp unit unit.KafkaGraphGraceSearchDiamondTest
 */
package unit;

import java.util.Collection;
import java.util.IdentityHashMap;
import java.util.LinkedHashSet;
import java.util.Map;

public class KafkaGraphGraceSearchDiamondTest {

    // ---- Minimal stub types -------------------------------------------------

    static class TopologyException extends RuntimeException {
        TopologyException(String msg) { super(msg); }
    }

    static class GraphNode {
        private final Collection<GraphNode> parentNodes = new LinkedHashSet<>();
        private final String name;
        GraphNode(String name) { this.name = name; }
        String nodeName() { return name; }
        Collection<GraphNode> parentNodes() { return new LinkedHashSet<>(parentNodes); }
        void addParent(GraphNode p) { parentNodes.add(p); }
        boolean parentNodes_isEmpty() { return parentNodes.isEmpty(); }
    }

    static class GracePeriodGraphNode extends GraphNode {
        private final long grace;
        GracePeriodGraphNode(String name, long grace) { super(name); this.grace = grace; }
        long gracePeriod() { return grace; }
    }

    // ---- Patched implementation (inlined here) --------------------------------

    static long findAndVerifyWindowGrace(GraphNode graphNode) {
        return findAndVerifyWindowGrace(graphNode, "", new IdentityHashMap<>());
    }

    @SuppressWarnings("rawtypes")
    private static long findAndVerifyWindowGrace(GraphNode graphNode,
                                                  String chain,
                                                  Map<GraphNode, Long> memo) {
        if (graphNode == null) {
            throw new TopologyException(
                "Window close time is only defined for windowed computations. Got [" + chain + "].");
        }
        if (graphNode instanceof GracePeriodGraphNode) {
            return ((GracePeriodGraphNode) graphNode).gracePeriod();
        }
        if (memo.containsKey(graphNode)) {
            return memo.get(graphNode);
        }
        final String newChain = chain.isEmpty() ? graphNode.nodeName() : graphNode.nodeName() + "->" + chain;
        if (graphNode.parentNodes_isEmpty()) {
            throw new TopologyException(
                "Window close time is only defined for windowed computations. Got [" + newChain + "].");
        }
        long inheritedGrace = -1;
        for (GraphNode parentNode : graphNode.parentNodes()) {
            long parentGrace = findAndVerifyWindowGrace(parentNode, newChain, memo);
            if (parentGrace > inheritedGrace) inheritedGrace = parentGrace;
        }
        if (inheritedGrace == -1) throw new IllegalStateException();
        memo.put(graphNode, inheritedGrace);
        return inheritedGrace;
    }

    // ---- Also test the DEFECT version (call-counter) -------------------------

    static int callCount;

    static long findAndVerifyWindowGrace_DEFECT(GraphNode graphNode, String chain) {
        callCount++;
        if (graphNode == null) throw new TopologyException("null");
        if (graphNode instanceof GracePeriodGraphNode) return ((GracePeriodGraphNode) graphNode).gracePeriod();
        final String newChain = chain.isEmpty() ? graphNode.nodeName() : graphNode.nodeName() + "->" + chain;
        if (graphNode.parentNodes_isEmpty()) throw new TopologyException("no grace: " + newChain);
        long inheritedGrace = -1;
        for (GraphNode parentNode : graphNode.parentNodes()) {
            long pg = findAndVerifyWindowGrace_DEFECT(parentNode, newChain);
            if (pg > inheritedGrace) inheritedGrace = pg;
        }
        if (inheritedGrace == -1) throw new IllegalStateException();
        return inheritedGrace;
    }

    // ---- Helpers to build diamond graph -------------------------------------

    /**
     * Build a D-deep diamond DAG with a single GracePeriodGraphNode at the top.
     *
     * Shape for D=2:
     *   grace (root)
     *   /          \
     * mid_0_0    mid_0_1
     *   \          /
     *    suppress (leaf)
     *
     * Each level has 2 nodes, each connected to both nodes on the level above.
     * Node count = 2*D + 2 (root + 2D mid nodes + leaf).
     */
    static GraphNode buildDiamond(int depth, long graceMs) {
        GracePeriodGraphNode root = new GracePeriodGraphNode("grace", graceMs);

        // current "top" layer
        GraphNode[] upper = new GraphNode[]{root};

        for (int level = 0; level < depth; level++) {
            GraphNode[] lower = new GraphNode[2];
            lower[0] = new GraphNode("mid_" + level + "_0");
            lower[1] = new GraphNode("mid_" + level + "_1");
            // both lower nodes have both upper nodes as parents
            for (GraphNode u : upper) {
                lower[0].addParent(u);
                lower[1].addParent(u);
            }
            upper = lower;
        }

        // leaf suppress node: has all of upper as parents
        GraphNode suppress = new GraphNode("suppress");
        for (GraphNode u : upper) {
            suppress.addParent(u);
        }
        return suppress;
    }

    // ---- Tests ---------------------------------------------------------------

    static int passed = 0;
    static int failed = 0;

    static void assertEq(String label, long expected, long actual) {
        if (expected == actual) {
            System.out.println("PASS " + label);
            passed++;
        } else {
            System.out.println("FAIL " + label + " expected=" + expected + " got=" + actual);
            failed++;
        }
    }

    static void assertTrue(String label, boolean condition) {
        if (condition) {
            System.out.println("PASS " + label);
            passed++;
        } else {
            System.out.println("FAIL " + label);
            failed++;
        }
    }

    public static void main(String[] args) {
        // Test 1: simple linear chain (D=0: just grace → suppress)
        {
            GracePeriodGraphNode grace = new GracePeriodGraphNode("grace", 5000L);
            GraphNode suppress = new GraphNode("suppress");
            suppress.addParent(grace);
            assertEq("T1 linear-chain grace", 5000L, findAndVerifyWindowGrace(suppress));
        }

        // Test 2: D=1 diamond: one intermediate level, two mid nodes, each parent=grace
        {
            long result = findAndVerifyWindowGrace(buildDiamond(1, 3000L));
            assertEq("T2 D=1 diamond grace", 3000L, result);
        }

        // Test 3: D=3 diamond
        {
            long result = findAndVerifyWindowGrace(buildDiamond(3, 7000L));
            assertEq("T3 D=3 diamond grace", 7000L, result);
        }

        // Test 4: D=20 diamond — patched must complete quickly (linear)
        {
            long start = System.nanoTime();
            long result = findAndVerifyWindowGrace(buildDiamond(20, 9999L));
            long elapsed = System.nanoTime() - start;
            assertEq("T4 D=20 diamond grace value", 9999L, result);
            assertTrue("T4 D=20 patched completes in < 1s", elapsed < 1_000_000_000L);
            System.out.println("    (D=20 patched took " + elapsed / 1_000 + " us)");
        }

        // Test 5: DEFECT version at D=20 should be orders of magnitude slower
        // At D=20 the defect makes 2^20 ~ 1M calls. We measure call count instead.
        {
            callCount = 0;
            findAndVerifyWindowGrace_DEFECT(buildDiamond(10, 1000L), "");
            int defectCalls = callCount;

            // patched version: count calls by wrapping (we count memo hits separately)
            // Approximate: patched should visit O(2*10+2)=22 nodes; defect > 1000
            assertTrue("T5 DEFECT D=10 has exponential calls (> 1000)",
                       defectCalls > 1000);
            System.out.println("    (D=10 defect calls=" + defectCalls + ", expected ~2048)");
        }

        // Test 6: no grace node reachable → TopologyException
        {
            boolean threw = false;
            try {
                GraphNode root = new GraphNode("plain_root");  // no parents
                GraphNode mid = new GraphNode("mid");
                mid.addParent(root);
                findAndVerifyWindowGrace(mid);
            } catch (TopologyException e) {
                threw = true;
            }
            assertTrue("T6 no-grace throws TopologyException", threw);
        }

        // Test 7: maximum grace propagates (two grace parents, max wins)
        {
            GracePeriodGraphNode g1 = new GracePeriodGraphNode("g1", 100L);
            GracePeriodGraphNode g2 = new GracePeriodGraphNode("g2", 500L);
            GraphNode suppress = new GraphNode("suppress");
            suppress.addParent(g1);
            suppress.addParent(g2);
            assertEq("T7 max of two grace parents", 500L, findAndVerifyWindowGrace(suppress));
        }

        // Summary
        System.out.println();
        System.out.println("Results: " + passed + " passed, " + failed + " failed");
        if (failed > 0) System.exit(1);
    }
}