"""Quick validation tests for MedTriage environment endpoints."""
import httpx
import json

BASE = "http://127.0.0.1:7860"
ALL_PASS = True


def check(name, condition, detail=""):
    global ALL_PASS
    detail_str = f" ({detail})" if detail else ""
    if condition:
        print(f"  PASS: {name}{detail_str}")
    else:
        print(f"  FAIL: {name}{detail_str}")
        ALL_PASS = False


print("\n=== MedTriage Validation Tests ===\n")

# ── Health + Tasks ──────────────────────────────────────────────────────────
r = httpx.get(f"{BASE}/health", timeout=10)
check("Health endpoint returns 200", r.status_code == 200, r.json().get("status"))

r = httpx.get(f"{BASE}/tasks", timeout=10)
tasks = [t["id"] for t in r.json()["tasks"]]
check("3 tasks available", len(tasks) == 3, str(tasks))
check("vital-triage task exists", "vital-triage" in tasks)
check("differential-diagnosis task exists", "differential-diagnosis" in tasks)
check("treatment-safety task exists", "treatment-safety" in tasks)

# ── Task 1 Tests ─────────────────────────────────────────────────────────────
print("\n── Task 1: vital-triage ──")
r = httpx.post(f"{BASE}/reset", json={"task": "vital-triage"}, timeout=10)
check("Task 1 reset returns 200", r.status_code == 200)
obs = r.json()["observation"]
check("Observation has patient", "patient" in r.json()["observation"])
check("Observation has task_instruction", bool(obs.get("task_instruction")))
check("Observation has vitals", obs["patient"]["vitals"]["heart_rate"] is not None)

# Correct ESI-1
r = httpx.post(f"{BASE}/step", json={"esi_level": 1, "triage_reason": "hypotension severe"}, timeout=10)
check("Step returns 200", r.status_code == 200)
step = r.json()
check("Step returns reward in [0,1]", 0.0 <= step["reward"] <= 1.0, f"reward={step['reward']}")
check("Step done=True after action", step["done"] == True)

# ESI penalty test
httpx.post(f"{BASE}/reset", json={"task": "vital-triage"}, timeout=10)
r = httpx.post(f"{BASE}/step", json={"esi_level": 5, "triage_reason": "seems ok"}, timeout=10)
step = r.json()
check("Wrong ESI gets lower reward", step["reward"] <= 0.5, f"reward={step['reward']}")

# ── Task 2 Tests ─────────────────────────────────────────────────────────────
print("\n── Task 2: differential-diagnosis ──")
r = httpx.post(f"{BASE}/reset", json={"task": "differential-diagnosis"}, timeout=10)
check("Task 2 reset returns 200", r.status_code == 200)

r = httpx.post(f"{BASE}/step", json={
    "diagnoses": ["STEMI", "NSTEMI", "Aortic dissection"],
    "red_flags": ["ST elevation", "elevated troponin", "hypotension", "diaphoresis"],
    "recommended_tests": ["12-lead ECG", "Troponin", "CXR"]
}, timeout=10)
step = r.json()
check("Task 2 step returns 200", r.status_code == 200)
check("Task 2 reward in [0,1]", 0.0 <= step["reward"] <= 1.0, f"reward={step['reward']}")

# ── Task 3 Tests ─────────────────────────────────────────────────────────────
print("\n-- Task 3: treatment-safety --")
r = httpx.post(f"{BASE}/reset", json={"task": "treatment-safety"}, timeout=10)
check("Task 3 reset returns 200", r.status_code == 200)
obs = r.json()["observation"]
check("Diagnosis context provided", "diagnosis_context" in obs.get("context", {}))

# Correct safe treatment — loop until we hit T001 (STEMI patient where aspirin is correct)
t3_correct_pass = False
t3_contra_pass = False
for _ in range(20):
    r = httpx.post(f"{BASE}/reset", json={"task": "treatment-safety"}, timeout=10)
    pid = r.json()["observation"]["patient"]["patient_id"]
    if pid == "T001" and not t3_correct_pass:
        r2 = httpx.post(f"{BASE}/step", json={
            "diagnosis": "STEMI",
            "drug_name": "aspirin",
            "dose_mg": 300,
            "route": "PO",
            "rationale": "dual antiplatelet therapy"
        }, timeout=10)
        step = r2.json()
        check("Task 3 correct drug (aspirin for STEMI) reward >= 0.6", step["reward"] >= 0.6, f"reward={step['reward']}")
        t3_correct_pass = True

    if pid == "T004" and not t3_contra_pass:
        r2 = httpx.post(f"{BASE}/step", json={
            "diagnosis": "Status asthmaticus",
            "drug_name": "aspirin",
            "dose_mg": 300,
            "route": "PO",
            "rationale": "anti-inflammatory"
        }, timeout=10)
        step = r2.json()
        check("Contraindicated drug (aspirin in NSAID-allergy) returns 0.0", step["reward"] == 0.0, f"reward={step['reward']}")
        t3_contra_pass = True

    if t3_correct_pass and t3_contra_pass:
        break

if not t3_correct_pass:
    check("Task 3 correct drug test ran", False, "T001 patient not drawn in 20 tries")
if not t3_contra_pass:
    check("Task 3 contraindication test ran", False, "T004 patient not drawn in 20 tries")


# ── State endpoint ────────────────────────────────────────────────────────────
print("\n── State endpoint ──")
r = httpx.get(f"{BASE}/state", timeout=10)
check("State returns 200", r.status_code == 200)
state = r.json()
check("State has episode_id", bool(state.get("episode_id")))
check("State has step_count", "step_count" in state)
check("State has cumulative_reward", "cumulative_reward" in state)

# ── Summary ───────────────────────────────────────────────────────────────────
print()
if ALL_PASS:
    print("✅ ALL TESTS PASSED — Environment is working correctly!")
else:
    print("❌ SOME TESTS FAILED — Check output above.")