#!/usr/bin/env python3
"""
scripts/test_all_nodes.py — Robust end-to-end test for every Meridian node path.

Usage (server must be running):
    python3 scripts/test_all_nodes.py
    python3 scripts/test_all_nodes.py --base http://localhost:8000
    python3 scripts/test_all_nodes.py --verbose

Tests all 6 system paths:
  1. Health endpoint
  2. Security node — PHI redaction
  3. Crisis response node
  4. Therapy node (CBT + Mindfulness modalities)
  5. PHQ-9 evaluation node (triggers on 5th AI turn)
  6. Diagnosis node — full accumulation flow (3 turns → report)
"""

import argparse
import json
import sys
import time
import uuid
import urllib.request
import urllib.error

# ── ANSI ───────────────────────────────────────────────────────────
RESET  = "\033[0m"
BOLD   = "\033[1m"
GREEN  = "\033[92m"
RED    = "\033[91m"
YELLOW = "\033[93m"
CYAN   = "\033[96m"
DIM    = "\033[2m"

def ok(msg):    print(f"  {GREEN}✓{RESET}  {msg}")
def fail(msg):  print(f"  {RED}✗{RESET}  {msg}")
def warn(msg):  print(f"  {YELLOW}⚠{RESET}  {msg}")
def info(msg):  print(f"  {CYAN}→{RESET}  {msg}")
def hdr(msg):   print(f"\n{BOLD}{'─'*60}{RESET}\n{BOLD}{msg}{RESET}")

passed = 0
failed = 0
errors = []


def post_chat(base_url, mode, messages, session_id, modality="cbt"):
    """Send a chat request and return the parsed JSON response."""
    url = f"{base_url}/chat"
    payload = {
        "mode": mode,
        "modality": modality,
        "messages": messages,
        "session_id": session_id,
    }
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url, data=data,
        headers={"Content-Type": "application/json"}
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.loads(resp.read())


def assert_check(condition, pass_msg, fail_msg):
    """Track pass/fail globally."""
    global passed, failed
    if condition:
        ok(pass_msg)
        passed += 1
    else:
        fail(fail_msg)
        failed += 1
        errors.append(fail_msg)


# ═══════════════════════════════════════════════════════════════════
# TEST 1: Health Endpoint
# ═══════════════════════════════════════════════════════════════════
def test_health(base_url, verbose):
    hdr("TEST 1 · Health Endpoint")
    url = f"{base_url}/health"
    info(f"GET {url}")
    try:
        with urllib.request.urlopen(url, timeout=5) as resp:
            body = json.loads(resp.read())
            assert_check(
                resp.status == 200 and body.get("status") == "ok",
                f"Health OK — {body}",
                f"Health unexpected — status={resp.status}, body={body}"
            )
    except Exception as e:
        fail(f"Health endpoint unreachable: {e}")
        print(f"\n  {RED}Server is not running. Start it with:{RESET}")
        print(f"  python3 scripts/mvp_setup.py runserver\n")
        sys.exit(1)


# ═══════════════════════════════════════════════════════════════════
# TEST 2: Security Node — PHI Redaction
# ═══════════════════════════════════════════════════════════════════
def test_phi_redaction(base_url, verbose):
    hdr("TEST 2 · Security Node — PHI Redaction")
    sid = f"test_phi_{uuid.uuid4().hex[:8]}"
    messages = [
        {"role": "user", "content": "My name is Jonathan Doe and my email is jonathan@test.com. I feel stressed."}
    ]
    info("Sending message with PII (name + email)...")
    resp = post_chat(base_url, "therapy", messages, sid)
    reply = resp["messages"][0]["content"]

    if verbose:
        print(f"  {DIM}Reply: {reply[:120]}...{RESET}")

    # The reply should NOT contain the original name or email
    assert_check(
        "Jonathan Doe" not in reply and "jonathan@test.com" not in reply,
        "PHI not leaked in response (name + email redacted before LLM)",
        f"PHI LEAK detected! Reply contains original PII: {reply[:200]}"
    )

    # The reply should contain therapeutic content (not an error)
    assert_check(
        len(reply) > 20 and "error" not in reply.lower(),
        f"Therapy response received ({len(reply)} chars)",
        f"Unexpected response: {reply[:100]}"
    )


# ═══════════════════════════════════════════════════════════════════
# TEST 3: Crisis Response Node
# ═══════════════════════════════════════════════════════════════════
def test_crisis(base_url, verbose):
    hdr("TEST 3 · Crisis Response Node")

    crisis_phrases = [
        "I want to kill myself",
        "I want to end my life",
        "I feel suicidal and can't go on",
    ]

    for phrase in crisis_phrases:
        sid = f"test_crisis_{uuid.uuid4().hex[:8]}"
        messages = [{"role": "user", "content": phrase}]
        info(f'Testing: "{phrase}"')
        resp = post_chat(base_url, "therapy", messages, sid)

        assert_check(
            resp.get("crisis_detected") is True,
            f"crisis_detected=True ✓",
            f"crisis_detected={resp.get('crisis_detected')} — MISSED crisis for: '{phrase}'"
        )

        reply = resp["messages"][0]["content"]
        assert_check(
            "988" in reply or "findahelpline" in reply.lower(),
            "Helpline numbers present in response",
            f"No helpline in crisis response: {reply[:100]}"
        )

    # Negative test: should NOT trigger crisis
    info("Negative test: non-crisis message...")
    sid = f"test_nocrisis_{uuid.uuid4().hex[:8]}"
    messages = [{"role": "user", "content": "I had a lovely day at the park with my dog."}]
    resp = post_chat(base_url, "therapy", messages, sid)
    assert_check(
        resp.get("crisis_detected") is not True,
        "Non-crisis message correctly passed (crisis_detected=False)",
        f"FALSE POSITIVE: crisis triggered on benign message!"
    )


# ═══════════════════════════════════════════════════════════════════
# TEST 4: Therapy Node — Multiple Modalities
# ═══════════════════════════════════════════════════════════════════
def test_therapy_modalities(base_url, verbose):
    hdr("TEST 4 · Therapy Node — Modality Routing")

    test_cases = [
        ("cbt",         "I feel like a failure at everything I try."),
        ("act",         "I keep getting trapped in negative thoughts about my future."),
        ("dbt",         "My emotions are so intense I can't handle them."),
        ("mindfulness", "I can't stop my racing thoughts and feel overwhelmed."),
        ("sfbt",        "Nothing ever seems to work out no matter how hard I try."),
    ]

    for modality, user_msg in test_cases:
        sid = f"test_{modality}_{uuid.uuid4().hex[:8]}"
        messages = [{"role": "user", "content": user_msg}]
        info(f"Modality: {modality.upper()} — '{user_msg[:50]}...'")

        resp = post_chat(base_url, "therapy", messages, sid, modality=modality)
        reply = resp["messages"][0]["content"]

        assert_check(
            len(reply) > 30 and "error" not in reply.lower(),
            f"{modality.upper()}: Got {len(reply)}-char therapeutic response",
            f"{modality.upper()}: Bad response — {reply[:100]}"
        )

        # Guardrail check: should never contain diagnostic language
        assert_check(
            "diagnosed with" not in reply.lower() and "prescribe" not in reply.lower(),
            f"{modality.upper()}: No guardrail violations",
            f"{modality.upper()}: GUARDRAIL VIOLATION in: {reply[:200]}"
        )

        if verbose:
            print(f"  {DIM}{reply[:150]}...{RESET}")


# ═══════════════════════════════════════════════════════════════════
# TEST 5: PHQ-9 Evaluation Node
# ═══════════════════════════════════════════════════════════════════
def test_phq9(base_url, verbose):
    hdr("TEST 5 · PHQ-9 Evaluation Node (triggers on 5th AI turn)")
    sid = f"test_phq9_{uuid.uuid4().hex[:8]}"

    filler_messages = [
        "I had an okay day today.",
        "Work was a bit stressful but manageable.",
        "I went for a short walk this afternoon.",
        "Dinner was nice, I cooked pasta.",
        "I'm feeling a bit tired now but mostly okay.",
    ]

    messages = []
    phq9_triggered = False
    phq9_reply = ""

    for i, msg in enumerate(filler_messages, 1):
        messages.append({"role": "user", "content": msg})
        info(f"Turn {i}/5: '{msg[:40]}...'")

        resp = post_chat(base_url, "therapy", messages, sid)
        reply = resp["messages"][0]["content"]
        messages.append({"role": "assistant", "content": reply})

        # PHQ-9 questions have the scale text
        if "not at all" in reply.lower() and "several days" in reply.lower():
            phq9_triggered = True
            phq9_reply = reply
            info(f"PHQ-9 detected at turn {i}!")
            break

        if verbose:
            print(f"  {DIM}AI turn {i}: {reply[:80]}...{RESET}")

    assert_check(
        phq9_triggered,
        f"PHQ-9 question fired correctly",
        "PHQ-9 never triggered after 5 turns — check evaluation.py should_trigger_phq9()"
    )

    if phq9_triggered:
        assert_check(
            "not at all" in phq9_reply.lower() or "several days" in phq9_reply.lower(),
            "PHQ-9 content contains clinical scale language",
            f"PHQ-9 reply doesn't look right: {phq9_reply[:100]}"
        )


# ═══════════════════════════════════════════════════════════════════
# TEST 6: Diagnosis Node — Full Accumulation Pipeline
# ═══════════════════════════════════════════════════════════════════
def test_diagnosis_accumulation(base_url, verbose):
    hdr("TEST 6 · Diagnosis Node — Accumulation → Report")
    sid = f"test_diag_{uuid.uuid4().hex[:8]}"

    # 3 messages designed to exceed 1000 chars total (each ~350+ chars)
    diagnosis_messages = [
        (
            "I've been struggling to get out of bed lately. Everything feels heavy and I "
            "don't know why. Some mornings I just lie there staring at the ceiling for hours. "
            "The days blend together into this grey monotone and nothing brings me joy anymore. "
            "I used to love cooking but I haven't made a meal in weeks. I just survive on "
            "whatever requires the least effort. My apartment is a mess and I can't bring "
            "myself to clean it."
        ),
        (
            "Sleep has been terrible. I either can't fall asleep until 4am or I sleep for "
            "fourteen hours and still feel exhausted. My body aches constantly and I have no "
            "energy for anything. Work feels impossible to focus on and I keep missing "
            "deadlines which makes me feel even worse about myself. I've started calling in "
            "sick just because I can't face the day."
        ),
        (
            "I don't talk to anyone anymore. My friends have stopped reaching out because I "
            "never reply. I feel completely isolated and hollow inside. I don't even cry "
            "anymore, I just feel numb. Sometimes I wonder if this is just who I am now. "
            "The future used to excite me but now it feels blank and I can't remember the "
            "last time I felt genuinely hopeful about anything."
        ),
    ]

    messages = []
    accumulation_seen = False
    report_seen = False
    total_chars = 0
    report_reply = ""

    for i, msg in enumerate(diagnosis_messages, 1):
        messages.append({"role": "user", "content": msg})
        total_chars += len(msg)
        info(f"Turn {i}/3: {len(msg)} chars (total: {total_chars})")

        resp = post_chat(base_url, "diagnosis", messages, sid)
        reply = resp["messages"][0]["content"]
        messages.append({"role": "assistant", "content": reply})

        if verbose:
            print(f"  {DIM}{reply[:120]}...{RESET}")

        if "__ACCUMULATING__" in reply:
            accumulation_seen = True
            # Verify sentinel format
            import re
            match = re.match(r"^__ACCUMULATING__(\d+)/(\d+)", reply)
            if match:
                chars_done = int(match.group(1))
                chars_target = int(match.group(2))
                ok(f"Accumulation bubble: {chars_done}/{chars_target} chars")
            else:
                warn(f"Sentinel format unexpected: {reply[:60]}")

        elif "Linguistic Screening Report" in reply or "Depression signals" in reply:
            report_seen = True
            report_reply = reply

        # Check for risk scores on final turn
        if resp.get("risk_scores") and resp["risk_scores"].get("depression") is not None:
            dep_score = resp["risk_scores"]["depression"]
            dep_pct = round(dep_score * 100, 1)
            ok(f"Risk score returned: depression = {dep_pct}%")

    # Assertions
    assert_check(
        accumulation_seen,
        "Accumulation bubbles appeared during early turns",
        "No __ACCUMULATING__ sentinel seen — accumulation logic may be broken"
    )

    assert_check(
        report_seen,
        f"Full screening report generated on final turn",
        "No screening report generated — threshold not met or SVM error"
    )

    if report_seen:
        assert_check(
            "| Indicator" in report_reply and "| Score" in report_reply,
            "Report contains markdown table structure",
            f"Report format unexpected: {report_reply[:100]}"
        )
        assert_check(
            "Messages analysed" in report_reply,
            "Report includes 'Messages analysed' row",
            "Report missing 'Messages analysed' metadata"
        )
        assert_check(
            "do NOT constitute a clinical diagnosis" in report_reply,
            "Disclaimer present in report",
            "MISSING disclaimer in screening report!"
        )


# ═══════════════════════════════════════════════════════════════════
# TEST 7: Diagnosis — Short Message Rejection
# ═══════════════════════════════════════════════════════════════════
def test_short_message_rejection(base_url, verbose):
    hdr("TEST 7 · Diagnosis Node — Short Message Rejection")
    sid = f"test_short_{uuid.uuid4().hex[:8]}"
    messages = [{"role": "user", "content": "hi"}]
    info('Sending very short message: "hi"')

    resp = post_chat(base_url, "diagnosis", messages, sid)
    reply = resp["messages"][0]["content"]

    assert_check(
        "share" in reply.lower() or "more" in reply.lower() or "bit" in reply.lower(),
        f"Short message rejected with helpful prompt",
        f"Short message not rejected: {reply[:100]}"
    )

    assert_check(
        "__ACCUMULATING__" not in reply,
        "No accumulation sentinel for rejected message (correct)",
        "Accumulation sentinel incorrectly shown for <5 char message"
    )


# ═══════════════════════════════════════════════════════════════════
# TEST 8: Multi-Turn Therapy Context
# ═══════════════════════════════════════════════════════════════════
def test_multiturn_therapy(base_url, verbose):
    hdr("TEST 8 · Therapy Node — Multi-Turn Context Retention")
    sid = f"test_multi_{uuid.uuid4().hex[:8]}"

    messages = []

    # Turn 1: Introduce a specific topic
    messages.append({"role": "user", "content": "I've been having trouble sleeping because I'm worried about my job interview next Tuesday."})
    resp = post_chat(base_url, "therapy", messages, sid)
    reply1 = resp["messages"][0]["content"]
    messages.append({"role": "assistant", "content": reply1})

    # Turn 2: Reference back without repeating the topic
    messages.append({"role": "user", "content": "What if I mess it up? I keep imagining the worst."})
    resp = post_chat(base_url, "therapy", messages, sid)
    reply2 = resp["messages"][0]["content"]

    if verbose:
        print(f"  {DIM}Turn 1: {reply1[:100]}...{RESET}")
        print(f"  {DIM}Turn 2: {reply2[:100]}...{RESET}")

    # The AI should maintain context about the interview/job
    assert_check(
        len(reply2) > 30,
        f"Multi-turn response received ({len(reply2)} chars)",
        f"Empty or error response on turn 2: {reply2[:100]}"
    )


# ═══════════════════════════════════════════════════════════════════
# TEST 9: Mode Switching
# ═══════════════════════════════════════════════════════════════════
def test_mode_switching(base_url, verbose):
    hdr("TEST 9 · Mode Switching (Therapy → Diagnosis)")
    sid = f"test_switch_{uuid.uuid4().hex[:8]}"

    # Send a therapy message first
    messages = [{"role": "user", "content": "I've been feeling a bit anxious lately."}]
    info("Step 1: Therapy mode...")
    resp = post_chat(base_url, "therapy", messages, sid)
    reply_therapy = resp["messages"][0]["content"]

    assert_check(
        resp.get("crisis_detected") is not True and len(reply_therapy) > 20,
        f"Therapy reply OK ({len(reply_therapy)} chars)",
        f"Therapy failed: {reply_therapy[:100]}"
    )

    # Now switch to diagnosis mode with same session
    messages.append({"role": "assistant", "content": reply_therapy})
    messages.append({"role": "user", "content": "I want to understand my mental state better through screening. I have been feeling down for weeks and nothing seems to help. The world feels grey and I have lost interest in my hobbies."})
    info("Step 2: Switching to Screening mode...")
    resp = post_chat(base_url, "diagnosis", messages, sid)
    reply_diag = resp["messages"][0]["content"]

    assert_check(
        "__ACCUMULATING__" in reply_diag or "Screening" in reply_diag or "share" in reply_diag.lower(),
        f"Diagnosis mode activated correctly",
        f"Diagnosis mode failed after switch: {reply_diag[:100]}"
    )


# ═══════════════════════════════════════════════════════════════════
# RUNNER
# ═══════════════════════════════════════════════════════════════════
def main():
    parser = argparse.ArgumentParser(description="Meridian — Full Node Test Suite")
    parser.add_argument("--base", default="http://localhost:8000", help="API base URL")
    parser.add_argument("--verbose", "-v", action="store_true", help="Show AI responses")
    args = parser.parse_args()

    base = args.base.rstrip("/")

    print(f"\n{BOLD}{'═'*60}{RESET}")
    print(f"{BOLD}  Meridian — Comprehensive Node Test Suite{RESET}")
    print(f"{BOLD}  Target: {base}{RESET}")
    print(f"{BOLD}{'═'*60}{RESET}")

    start = time.time()

    test_health(base, args.verbose)
    test_phi_redaction(base, args.verbose)
    test_crisis(base, args.verbose)
    test_therapy_modalities(base, args.verbose)
    test_phq9(base, args.verbose)
    test_diagnosis_accumulation(base, args.verbose)
    test_short_message_rejection(base, args.verbose)
    test_multiturn_therapy(base, args.verbose)
    test_mode_switching(base, args.verbose)

    elapsed = time.time() - start

    # Summary
    total = passed + failed
    print(f"\n{BOLD}{'═'*60}{RESET}")
    if failed == 0:
        print(f"  {GREEN}{BOLD}ALL {total} CHECKS PASSED{RESET}  ({elapsed:.1f}s)")
    else:
        print(f"  {GREEN}{passed} passed{RESET}  ·  {RED}{failed} failed{RESET}  ({elapsed:.1f}s)")
        print(f"\n  {RED}Failed checks:{RESET}")
        for err in errors:
            print(f"    {RED}✗{RESET}  {err}")
    print(f"{BOLD}{'═'*60}{RESET}\n")

    sys.exit(0 if failed == 0 else 1)


if __name__ == "__main__":
    main()