test_local.py

"""
Lulluna — Local Test Suite
============================
Run this BEFORE touching the frontend.
Tests every layer: model → prompt → output → parse → JSON shape.

Usage:
    # Full test suite (all traditions, all ages)
    python test_local.py

    # Quick smoke test — one story only
    python test_local.py --quick

    # Test a specific tradition
    python test_local.py --tradition Jataka

    # Test the API endpoint directly (app must be running)
    python test_local.py --api

    # See raw model output
    python test_local.py --quick --verbose
"""

import sys
import time
import json
import argparse
import requests

# ── Colours ──────────────────────────────────────────────────
G = "\033[92m"; R = "\033[91m"; Y = "\033[93m"; B = "\033[94m"
BOLD = "\033[1m"; RESET = "\033[0m"
ok   = lambda m: print(f"{G}  ✓ {m}{RESET}")
fail = lambda m: print(f"{R}  ✗ {m}{RESET}")
warn = lambda m: print(f"{Y}  ⚠ {m}{RESET}")
info = lambda m: print(f"{B}  → {m}{RESET}")
h    = lambda m: print(f"\n{BOLD}{m}{RESET}")

# ── Test cases ────────────────────────────────────────────────
QUICK_CASE = {
    "age": "5-6", "value": "Kindness", "tradition": "Aesop",
    "length": "short", "name": "", "interests": "",
    "min_words": 200, "max_words": 420,
}

FULL_CASES = [
    # age / value / tradition / length / name / interests
    {"age": "3-4",  "value": "Love",         "tradition": "Jataka",         "length": "short",  "name": "Maya",  "interests": ""},
    {"age": "5-6",  "value": "Kindness",      "tradition": "Aesop",          "length": "short",  "name": "",      "interests": "turtles"},
    {"age": "5-6",  "value": "Wisdom",        "tradition": "African",        "length": "medium", "name": "",      "interests": ""},
    {"age": "7-8",  "value": "Courage",       "tradition": "Norse",          "length": "medium", "name": "Leo",   "interests": "dragons"},
    {"age": "7-8",  "value": "Honesty",       "tradition": "Arabian",        "length": "medium", "name": "",      "interests": "the sea"},
    {"age": "9-10", "value": "Perseverance",  "tradition": "Panchatantra",   "length": "long",   "name": "",      "interests": ""},
    {"age": "9-10", "value": "Humility",      "tradition": "Sufi",           "length": "short",  "name": "",      "interests": ""},
    {"age": "5-6",  "value": "Generosity",    "tradition": "Chinese",        "length": "medium", "name": "Mei",   "interests": "the moon"},
    {"age": "7-8",  "value": "Friendship",    "tradition": "Celtic",         "length": "medium", "name": "",      "interests": "the ocean"},
    {"age": "3-4",  "value": "Patience",      "tradition": "Native American","length": "short",  "name": "",      "interests": "rabbits"},
]

EXPECTED_LENGTH = {
    "short":  (180, 420),
    "medium": (420, 800),
    "long":   (750, 1400),
}

BANNED_PHRASES = [
    "i cannot", "i am unable", "as an ai", "language model",
    "i apologize", "i'm sorry but",
    "blood", "death", "kill", "murder", "violence", "terror",
    "sexy", "romantic", "kiss", "boyfriend", "girlfriend",
]

REQUIRED_SECTIONS = ["TITLE:", "EMOJI:", "STORY:"]


# ── Validation ────────────────────────────────────────────────

def validate_result(result: dict, case: dict, verbose: bool) -> list[str]:
    """Check output quality. Returns list of failure strings."""
    failures = []

    # Shape
    for key in ("title", "emoji", "body"):
        if not result.get(key):
            failures.append(f"Missing field: {key}")

    title = result.get("title", "")
    body  = result.get("body", "")

    if len(title) < 3:
        failures.append(f"Title too short: '{title}'")
    if len(title) > 120:
        failures.append(f"Title too long: {len(title)} chars")

    # Word count
    word_count = len(body.split())
    lo, hi = EXPECTED_LENGTH.get(case["length"], (100, 1500))
    if word_count < lo:
        failures.append(f"Too short: {word_count} words (min {lo} for {case['length']})")
    elif word_count > hi * 1.5:   # 50% tolerance over
        failures.append(f"Too long: {word_count} words (max ~{hi} for {case['length']})")

    # Banned content
    body_lower = body.lower()
    for phrase in BANNED_PHRASES:
        if phrase in body_lower:
            failures.append(f"Banned phrase: '{phrase}'")

    # Paragraph structure
    paragraphs = [p.strip() for p in body.split("\n\n") if p.strip()]
    if len(paragraphs) < 2:
        failures.append(f"No paragraph breaks — body is one block")

    # Child name injection
    if case.get("name") and case["name"].lower() not in body.lower():
        failures.append(f"Child name '{case['name']}' not found in story")

    # Calming close — story should end peacefully
    last_para = paragraphs[-1].lower() if paragraphs else ""
    sleep_words = ["sleep", "dream", "rest", "closed", "night", "peaceful",
                   "drifted", "warm", "soft", "quiet", "still", "stars", "moon"]
    if not any(w in last_para for w in sleep_words):
        warn(f"No sleep/dream words in final paragraph — may not close peacefully")
        # Not a hard failure — just a warning

    if verbose:
        print(f"\n  {BOLD}Title:{RESET} {result.get('title')}")
        print(f"  {BOLD}Emoji:{RESET} {result.get('emoji')}")
        print(f"  {BOLD}Words:{RESET} {word_count}")
        print(f"  {BOLD}Body preview:{RESET}")
        for para in paragraphs[:2]:
            print(f"    {para[:120]}...")
        print(f"  {BOLD}Closing:{RESET} ...{last_para[-100:]}")

    return failures


# ── Direct model test (no HTTP) ───────────────────────────────

def run_direct_tests(cases: list[dict], verbose: bool) -> tuple[int, int]:
    """Load the engine directly and test generation."""
    from app import generate_story

    passed = failed = 0
    timings = []

    for i, case in enumerate(cases):
        label = f"[{i+1}/{len(cases)}] {case['tradition']} · {case['value']} · age {case['age']} · {case['length']}"
        h(label)

        t0 = time.time()
        try:
            result = generate_story(
                age=case["age"],
                value=case["value"],
                tradition=case["tradition"],
                length=case["length"],
                child_name=case.get("name", ""),
                interests=case.get("interests", ""),
            )
            elapsed = time.time() - t0
            timings.append(elapsed)

            failures = validate_result(result, case, verbose)

            if failures:
                for f in failures:
                    fail(f)
                failed += 1
            else:
                word_count = len(result["body"].split())
                ok(f"'{result['title']}' · {word_count} words · {elapsed:.1f}s")
                passed += 1

        except Exception as e:
            fail(f"Exception: {e}")
            if verbose:
                import traceback; traceback.print_exc()
            failed += 1

    return passed, failed, timings


# ── API test (app must be running) ────────────────────────────

def run_api_tests(cases: list[dict], base_url: str, verbose: bool) -> tuple[int, int]:
    """
    Test by calling the running Gradio app's API endpoint.
    Mirrors exactly what stories.functions.ts does.
    """
    passed = failed = 0

    for i, case in enumerate(cases):
        label = f"[{i+1}/{len(cases)}] {case['tradition']} · {case['value']} · age {case['age']}"
        h(label)

        # Gradio API format: POST /api/predict with {"data": [...]}
        payload = {
            "data": [
                case["age"],
                case["value"],
                case["tradition"],
                case["length"],
                case.get("name", ""),
                case.get("interests", ""),
            ]
        }

        t0 = time.time()
        try:
            resp = requests.post(
                f"{base_url}/api/generate",
                json=payload,
                timeout=60,
            )
            elapsed = time.time() - t0

            if resp.status_code != 200:
                fail(f"HTTP {resp.status_code}: {resp.text[:200]}")
                failed += 1
                continue

            data = resp.json()
            # Gradio wraps response in {"data": [result]}
            result = data.get("data", [{}])[0]
            if isinstance(result, str):
                result = json.loads(result)

            failures = validate_result(result, case, verbose)

            if failures:
                for f in failures:
                    fail(f)
                failed += 1
            else:
                word_count = len(result["body"].split())
                ok(f"'{result['title']}' · {word_count} words · {elapsed:.1f}s")
                passed += 1

        except requests.exceptions.ConnectionError:
            fail(f"Cannot connect to {base_url} — is the app running? (python app.py)")
            return 0, len(cases), []
        except Exception as e:
            fail(f"Exception: {e}")
            if verbose:
                import traceback; traceback.print_exc()
            failed += 1

    return passed, failed, []


# ── Frontend wiring check ─────────────────────────────────────

def check_frontend_wiring(base_url: str):
    """
    Verify the frontend can reach the backend.
    Simulates exactly what stories.functions.ts sends.
    """
    h("Frontend wiring check")
    info(f"Simulating stories.functions.ts → POST {base_url}/api/generate")

    payload = {
        "data": ["5-6", "Kindness", "Aesop", "short", "Mira", "butterflies"]
    }

    try:
        resp = requests.post(f"{base_url}/api/generate", json=payload, timeout=30)

        if resp.status_code == 200:
            data = resp.json()
            result = data.get("data", [{}])[0]
            if isinstance(result, str):
                result = json.loads(result)
            if result.get("title") and result.get("body"):
                ok(f"Frontend wiring ✓ — story '{result['title']}' generated")
                ok("The frontend will work. Change the API URL and you're done.")
            else:
                fail(f"Response missing title/body: {result}")
        else:
            fail(f"HTTP {resp.status_code}: {resp.text[:200]}")

    except requests.exceptions.ConnectionError:
        fail(f"App not running at {base_url} — start with: python app.py")


# ── Entry point ───────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Lulluna backend tests")
    parser.add_argument("--quick",     action="store_true", help="One story only")
    parser.add_argument("--tradition", type=str, help="Test one tradition only")
    parser.add_argument("--api",       action="store_true", help="Test via HTTP (app must run)")
    parser.add_argument("--wiring",    action="store_true", help="Check frontend wiring only")
    parser.add_argument("--url",       default="http://localhost:7860", help="App URL")
    parser.add_argument("--verbose",   action="store_true", help="Show story content")
    args = parser.parse_args()

    h("Lulluna — Backend Test Suite")

    cases = FULL_CASES
    if args.quick:
        cases = [QUICK_CASE]
    elif args.tradition:
        cases = [c for c in FULL_CASES if c["tradition"].lower() == args.tradition.lower()]
        if not cases:
            print(f"{R}No cases for tradition '{args.tradition}'{RESET}")
            sys.exit(1)

    if args.wiring:
        check_frontend_wiring(args.url)
        return

    if args.api:
        passed, failed, _ = run_api_tests(cases, args.url, args.verbose)
    else:
        passed, failed, timings = run_direct_tests(cases, args.verbose)

    # Summary
    h("Results")
    print(f"  {G}Passed: {passed}{RESET}  {R}Failed: {failed}{RESET}")
    if timings:
        avg = sum(timings) / len(timings)
        fastest = min(timings)
        print(f"  Avg: {avg:.1f}s  Fastest: {fastest:.1f}s")
        if avg > 15:
            warn("Avg > 15s — check Metal backend is active")
        else:
            ok(f"Speed is good ({avg:.1f}s avg on M5)")

    if failed == 0:
        print(f"\n  {G}{BOLD}All tests passed! Backend is ready.{RESET}")
        print(f"\n  {B}Next step:{RESET}")
        print(f"  Change the URL in stories.functions.ts to http://localhost:7860")
        print(f"  Then run: python test_local.py --wiring")
    else:
        print(f"\n  {R}Fix failures before wiring to frontend.{RESET}")
        sys.exit(1)


if __name__ == "__main__":
    main()