"""Comprehensive round-trip test for the forge (refactored) chunker.

Covers:
  1. Repetitive sentence-grouped text (the probe case)
  2. Long single sentence -> triggers _hard_split_long
  3. Mixed: some sentences fit, some need hard split
  4. Edge cases: empty, whitespace-only, single chunk, no overlap
  5. Many random-ish inputs with varying params
"""
import sys
import random
import string
sys.path.insert(0, "/root/.hermes/wizard-outputs/forge")
import chunker

PASS = 0
FAIL = 0

def check(name, text, max_tokens=256, overlap_tokens=32):
    global PASS, FAIL
    if not text.strip():
        expected = ""
    else:
        expected = text.strip()
    cs = chunker.chunk_text(text, max_tokens=max_tokens, overlap_tokens=overlap_tokens)
    r = chunker.reconstruct(cs, overlap_tokens)
    ok = r == expected
    status = "PASS" if ok else "FAIL"
    if ok:
        PASS += 1
    else:
        FAIL += 1
        print(f"[{status}] {name}")
        print(f"   orig_len={len(expected)} recon_len={len(r)}")
        for i in range(min(len(r), len(expected))):
            if r[i] != expected[i]:
                print(f"   FIRST DIFF at char {i}")
                print(f"   orig:  {expected[max(0,i-30):i+30]!r}")
                print(f"   recon: {r[max(0,i-30):i+30]!r}")
                break
        else:
            print(f"   common prefix OK but lengths differ: {len(r)} vs {len(expected)}")
            print(f"   orig tail:  {expected[-30:]!r}")
            print(f"   recon tail: {r[-30:]!r}")
        print(f"   nchunks={len(cs)}")
        # show if any chunk exceeded max_tokens
        for i, c in enumerate(cs):
            # strip the overlap tail for token check
            actual = c
            if chunker._OVERLAP_SENTINEL in actual:
                actual = actual.partition(chunker._OVERLAP_SENTINEL)[2]
            est = chunker.estimate_tokens(actual)
            if est > max_tokens:
                print(f"   chunk[{i}] exceeds max_tokens: est={est} > {max_tokens}, content={actual[:60]!r}")
    return ok


# Case 1: repetitive sentence-grouped text
base = "The quick brown fox jumps over the lazy dog. Pack my box with five dozen liquor jugs. How vexingly quick daft zebras jump. "
text = (base + "Variant sentence here. ") * 20
check("repetitive-sentences", text, max_tokens=64, overlap_tokens=8)

# Case 2: long single sentence -> hard split
long_sent = "Word" + " word" * 200 + " end."  # ~1000 chars, way over 80
check("long-sentence-hard-split", long_sent, max_tokens=20, overlap_tokens=4)

# Case 2b: long single sentence with sentence-ending punctuation
long_sent2 = " ".join(["Sentence"] + ["word"]*300 + ["."])
check("long-sentence-hard-split-v2", long_sent2, max_tokens=30, overlap_tokens=6)

# Case 3: mixed - some short, some long sentences
mixed = "Short one. " + long_sent + " Another short. " + long_sent2 + " Final."
check("mixed-sizes", mixed, max_tokens=20, overlap_tokens=4)

# Case 4: edge cases
check("empty-string", "", max_tokens=64, overlap_tokens=8)
check("whitespace-only", "   \n\t  ", max_tokens=64, overlap_tokens=8)
check("single-sentence", "Just one sentence here.", max_tokens=64, overlap_tokens=8)
check("two-sentences-no-overlap", "First sentence. Second sentence.", max_tokens=256, overlap_tokens=0)
check("small-overlap", "First sentence. Second sentence.", max_tokens=64, overlap_tokens=2)

# Case 5: multiple sentences exactly at boundary
exactly = " ".join(f"Sentence number {i}." for i in range(10))
check("boundary-fit", exactly, max_tokens=10, overlap_tokens=2)

# Case 6: text with special chars, multiple spaces
special = "Hello  world.  How are you?  I am fine.  " * 5
check("multi-space", special, max_tokens=32, overlap_tokens=4)

# Case 7: text starting/ending with whitespace
ws_text = "  The quick brown fox.  The lazy dog.  "
check("leading-trailing-ws", ws_text, max_tokens=64, overlap_tokens=8)

# Case 8: single long sentence that needs hard split, with high overlap
long3 = "A" + " b" * 500 + " end."
check("long-with-high-overlap", long3, max_tokens=16, overlap_tokens=8)

# Case 9: random text fuzz
random.seed(42)
words = ["alpha", "beta", "gamma", "delta", "epsilon", "zeta", "eta", "theta"]
for trial in range(20):
    sents = []
    for _ in range(random.randint(3, 15)):
        nwords = random.randint(3, 50)
        s = " ".join(random.choice(words) for _ in range(nwords))
        sents.append(s + ".")
    text = " ".join(sents)
    check(f"random-trial-{trial}", text, max_tokens=random.randint(16, 64),
          overlap_tokens=random.randint(0, 8))

# Case 10: text with newlines
nl_text = "First paragraph.\n\nSecond paragraph here.\nThird line."
check("newlines", nl_text, max_tokens=64, overlap_tokens=8)

print(f"\n=== Results: {PASS} passed, {FAIL} failed ===")
if FAIL:
    sys.exit(1)