"""
Unit tests for distlib CWE-407 defects.

distlib-0001: get_dependent_dists — dep list O(N²) membership check
  File: distlib/database.py ~line 1284
  Fix:  parallel dep_set for O(1) succ-not-in-dep guard

distlib-0002: get_steps — result.remove O(N) + todo.pop(0) O(N)
  File: distlib/util.py ~line 1127
  Fix:  OrderedDict.move_to_end O(1) + deque.popleft O(1)

Tests use operation counters to verify algorithmic complexity.
No distlib import needed — implementations are inlined as stubs.
"""

import sys
import time
from collections import deque, OrderedDict


# ---------------------------------------------------------------------------
# distlib-0001: get_dependent_dists
# ---------------------------------------------------------------------------

def get_dependent_dists_before(reverse_list, dist):
    """
    Defective implementation: O(N²) due to `succ not in dep` list scan.
    Returns (result_list, op_count).
    """
    ops = 0
    dep = [dist]
    todo = list(reverse_list.get(dist, []))
    while todo:
        d = todo.pop()
        dep.append(d)
        for succ in reverse_list.get(d, []):
            ops += len(dep)          # cost of `succ not in dep`
            if succ not in dep:
                todo.append(succ)
    dep.pop(0)
    return dep, ops


def get_dependent_dists_after(reverse_list, dist):
    """
    Fixed implementation: O(N) with parallel dep_set for O(1) membership.
    Returns (result_list, op_count).
    """
    ops = 0
    dep = [dist]
    dep_set = {dist}
    todo = list(reverse_list.get(dist, []))
    while todo:
        d = todo.pop()
        dep.append(d)
        dep_set.add(d)
        for succ in reverse_list.get(d, []):
            ops += 1                 # O(1) set lookup
            if succ not in dep_set:
                todo.append(succ)
    dep.pop(0)
    return dep, ops


def make_fan_in_graph(n):
    """
    Build a reverse_list that forces O(N²) membership checks.

    Structure: a "collector" node c is depended on by p0..p(N-1) (in
    reverse_list terms: reverse_list[c] = [p0..p(N-1)]).
    Each pi is in turn depended on by a unique "leaf" node li
    (reverse_list[pi] = [li]).
    Additionally, every leaf li lists *all* pj (j!=i) as its successors,
    so when we visit li, we check all N pj nodes against dep.
    By the time we visit leaf li, dep already contains c, p0..p(i), and
    the previously-visited leaves, so each `pj not in dep` check runs
    against a growing list.

    Total membership checks: N leaves × N pj nodes = O(N²).
    All checks correctly return False the first time (pj nodes added from
    c's reverse_list are in dep) or True when pj is already there.
    Result size = 1 (c) + N (pi) + N (li) = 2N+1 dependents.
    """
    reverse_list = {}
    # c's direct dependents: p0..p(N-1)
    reverse_list["c"] = [f"p{i}" for i in range(n)]
    for i in range(n):
        # pi is depended on by leaf li
        reverse_list[f"p{i}"] = [f"l{i}"]
        # leaf li lists all pj as successors (all already in dep when li visited)
        reverse_list[f"l{i}"] = [f"p{j}" for j in range(n)]
    return reverse_list


def make_simple_chain(n):
    """
    Simple linear chain: base <- p0 <- p1 <- ... <- p(n-1).
    reverse_list[base] = [p0], reverse_list[pi] = [p(i+1)], last has [].
    Result of get_dependent_dists(base) = [p0, p1, ..., p(n-1)].
    """
    rl = {"base": ["p0"]}
    for i in range(n - 1):
        rl[f"p{i}"] = [f"p{i+1}"]
    rl[f"p{n-1}"] = []
    return rl


def test_distlib_0001_correctness():
    """Both implementations produce the same result on a simple chain."""
    N = 20
    rl = make_simple_chain(N)
    result_before, _ = get_dependent_dists_before(rl, "base")
    result_after, _ = get_dependent_dists_after(rl, "base")
    assert sorted(result_before) == sorted(result_after), (
        f"Results differ:\n  before={sorted(result_before)}\n  after={sorted(result_after)}"
    )
    assert len(result_before) == N, (
        f"Expected {N} dependents, got {len(result_before)}"
    )
    print(f"  correctness: both return {len(result_before)} dependents")


def test_distlib_0001_complexity():
    """
    BEFORE: op count grows as O(N²); AFTER: op count grows as O(N).

    At N=500, each of N leaves checks N peers against a dep list of size ~N
    → ~N² = 250,000 membership-check ops in BEFORE; each check is O(1) in
    AFTER → ~N² total ops but each counted as 1 → ratio reflects list vs set.
    Expected ratio: >= 50x.
    """
    N = 500
    rl = make_fan_in_graph(N)

    _, ops_before = get_dependent_dists_before(rl, "c")
    _, ops_after = get_dependent_dists_after(rl, "c")

    ratio = ops_before / ops_after if ops_after > 0 else float("inf")
    print(f"  N={N}: ops_before={ops_before:,}, ops_after={ops_after:,}, ratio={ratio:.1f}x")

    assert ratio >= 50, (
        f"Expected >= 50x op ratio at N={N}, got {ratio:.1f}x "
        f"(before={ops_before:,}, after={ops_after:,})"
    )
    print(f"  PASS: {ratio:.1f}x operation-count reduction")


def test_distlib_0001_no_infinite_loop():
    """
    dep / dep_set prevent revisiting nodes in a graph with shared sub-deps.
    Diamond: base <- {a, b}, a <- c, b <- c (c depends on both a and b).
    """
    reverse_list = {
        "base": ["a", "b"],
        "a":    ["c"],
        "b":    ["c"],
        "c":    [],
    }
    result_before, _ = get_dependent_dists_before(reverse_list, "base")
    result_after, _ = get_dependent_dists_after(reverse_list, "base")

    # c should appear exactly once in both (dedup prevents double-add)
    assert result_before.count("c") == 1, f"BEFORE: c appears {result_before.count('c')} times"
    assert result_after.count("c") == 1, f"AFTER: c appears {result_after.count('c')} times"
    print(f"  diamond dedup: BEFORE={result_before}, AFTER={result_after}")


# ---------------------------------------------------------------------------
# distlib-0002: get_steps
# ---------------------------------------------------------------------------

def get_steps_before(preds, final):
    """
    Defective implementation:
      - result.remove(step) is O(N) list scan
      - todo.pop(0) is O(N) list shift
    Returns (result_list, op_count).
    """
    ops = 0
    result = []
    todo = []
    seen = set()
    todo.append(final)
    while todo:
        ops += len(todo)             # cost of pop(0): shifts all remaining elements
        step = todo.pop(0)
        if step in seen:
            if step != final:
                ops += len(result)   # cost of result.remove(step)
                result.remove(step)
                result.append(step)
        else:
            seen.add(step)
            result.append(step)
            preds_list = preds.get(step, ())
            todo.extend(preds_list)
    return list(reversed(result)), ops


def get_steps_after(preds, final):
    """
    Fixed implementation:
      - result is OrderedDict; move_to_end is O(1)
      - todo is deque; popleft is O(1)
    Returns (result_list, op_count).
    """
    ops = 0
    result = OrderedDict()
    todo = deque()
    seen = set()
    todo.append(final)
    while todo:
        ops += 1                     # O(1) deque.popleft
        step = todo.popleft()
        if step in seen:
            if step != final:
                ops += 1             # O(1) OrderedDict.move_to_end
                result.move_to_end(step)
        else:
            seen.add(step)
            result[step] = None
            preds_list = preds.get(step, ())
            todo.extend(preds_list)
    return list(reversed(list(result))), ops


def make_diamond_steps():
    """
    Diamond dependency graph for steps:
      final -> a -> c
      final -> b -> c
    Step c appears as a predecessor of both a and b.
    When c is first seen via a, then revisited via b, result.remove(c)
    is called in the BEFORE version.
    """
    return {
        "final": ("a", "b"),
        "a":     ("c",),
        "b":     ("c",),
        "c":     (),
    }


def make_chain_steps(n):
    """
    Linear chain: final -> s0 -> s1 -> ... -> s(n-1).
    Every step is visited exactly once; no remove() calls.
    Used to stress-test pop(0) vs popleft().
    """
    preds = {"final": (f"s0",)}
    for i in range(n - 1):
        preds[f"s{i}"] = (f"s{i+1}",)
    preds[f"s{n-1}"] = ()
    return preds


def make_wide_diamond_steps(width, depth):
    """
    Wide diamond: final has `width` direct predecessors, each of which
    has the same `depth` common predecessors.  Each common predecessor
    is revisited `width` times — triggering result.remove() in BEFORE.
    """
    preds = {}
    shared = [f"shared_{d}" for d in range(depth)]
    branches = [f"branch_{w}" for w in range(width)]

    preds["final"] = tuple(branches)
    for b in branches:
        preds[b] = tuple(shared)
    for s in shared:
        preds[s] = ()
    return preds, shared


def test_distlib_0002_correctness_diamond():
    """Both implementations produce the same topological order on a diamond."""
    preds = make_diamond_steps()
    result_before, _ = get_steps_before(preds, "final")
    result_after, _ = get_steps_after(preds, "final")

    assert result_before == result_after, (
        f"Results differ:\n  before={result_before}\n  after={result_after}"
    )
    # c should be first (deepest common dep), final last
    assert result_before[0] == "c", f"Expected 'c' first, got {result_before[0]}"
    assert result_before[-1] == "final", f"Expected 'final' last, got {result_before[-1]}"
    print(f"  diamond order: {result_before}")


def test_distlib_0002_correctness_chain():
    """Linear chain: order must be [s(n-1), ..., s0, final]."""
    N = 20
    preds = make_chain_steps(N)
    result_before, _ = get_steps_before(preds, "final")
    result_after, _ = get_steps_after(preds, "final")

    assert result_before == result_after, (
        f"Chain results differ:\n  before={result_before}\n  after={result_after}"
    )
    expected_first = f"s{N-1}"
    assert result_before[0] == expected_first, (
        f"Expected '{expected_first}' first, got {result_before[0]}"
    )
    print(f"  chain N={N}: first={result_before[0]}, last={result_before[-1]}")


def test_distlib_0002_complexity():
    """
    BEFORE: op count grows as O(N²) for wide-diamond graphs.
    AFTER: op count grows as O(N).

    Build a graph with width=200 branches all sharing depth=100 common steps.
    BEFORE: each of 100 shared nodes is remove()'d 199 times → ~100×200 = 20,000
            remove ops, each scanning a result list of up to 300 items → ~6,000,000 ops.
    AFTER: each move_to_end is O(1) → ~200×100 = 20,000 ops total.
    Expected ratio: >= 50x.
    """
    WIDTH = 200
    DEPTH = 100
    preds, _ = make_wide_diamond_steps(WIDTH, DEPTH)

    _, ops_before = get_steps_before(preds, "final")
    _, ops_after = get_steps_after(preds, "final")

    ratio = ops_before / ops_after if ops_after > 0 else float("inf")
    print(f"  width={WIDTH}, depth={DEPTH}: ops_before={ops_before:,}, "
          f"ops_after={ops_after:,}, ratio={ratio:.1f}x")

    assert ratio >= 50, (
        f"Expected >= 50x op ratio, got {ratio:.1f}x "
        f"(before={ops_before:,}, after={ops_after:,})"
    )
    print(f"  PASS: {ratio:.1f}x operation-count reduction")


def test_distlib_0002_popleft_complexity():
    """
    todo.pop(0) on a list is O(N); deque.popleft() is O(1).

    To expose the difference, we need a graph where todo stays large during
    traversal.  A wide fan-out graph achieves this: final has N direct
    predecessors, all leaf nodes.  When each branch is popped from todo, the
    remaining todo list still holds all the sibling branches — pop(0) must
    shift them all.  At N predecessors the total shift-work is O(N²).

    We count shifts explicitly by accumulating len(todo) before each pop(0).
    """
    N = 2000
    # final -> p0, p1, ..., p(N-1); each pi is a leaf
    preds = {"final": tuple(f"p{i}" for i in range(N))}
    for i in range(N):
        preds[f"p{i}"] = ()

    _, ops_before = get_steps_before(preds, "final")
    _, ops_after = get_steps_after(preds, "final")

    ratio = ops_before / ops_after if ops_after > 0 else float("inf")
    print(f"  fan-out N={N}: ops_before={ops_before:,}, ops_after={ops_after:,}, ratio={ratio:.1f}x")

    assert ratio >= 50, (
        f"Expected >= 50x op ratio at N={N}, got {ratio:.1f}x "
        f"(before={ops_before:,}, after={ops_after:,})"
    )
    print(f"  PASS: {ratio:.1f}x operation-count reduction")


# ---------------------------------------------------------------------------
# Runner
# ---------------------------------------------------------------------------

TESTS = [
    # distlib-0001
    test_distlib_0001_correctness,
    test_distlib_0001_complexity,
    test_distlib_0001_no_infinite_loop,
    # distlib-0002
    test_distlib_0002_correctness_diamond,
    test_distlib_0002_correctness_chain,
    test_distlib_0002_complexity,
    test_distlib_0002_popleft_complexity,
]

if __name__ == "__main__":
    failed = 0
    for t in TESTS:
        print(f"=== {t.__name__} ===")
        try:
            t()
            print("  PASS")
        except AssertionError as e:
            print(f"  FAIL: {e}", file=sys.stderr)
            failed += 1
        except Exception as e:
            print(f"  ERROR: {e}", file=sys.stderr)
            failed += 1

    print()
    if failed:
        print(f"FAILED: {failed}/{len(TESTS)} tests failed", file=sys.stderr)
        sys.exit(1)
    else:
        print(f"ALL {len(TESTS)} TESTS PASSED")