import asyncio
import uuid
import sys
from temporalio.client import Client

async def test_scribe_and_coder(client: Client):
    print("\n=== Testing AI Scribe & Medical Coder ===")
    dictation = (
        "Patient is a 45-year-old male presenting with severe, crushing chest pain radiating to his left arm. "
        "He reports shortness of breath and sweating. He has a past medical history of hypertension and hyperlipidemia. "
        "Current medications include Lisinopril 20mg daily and Atorvastatin 40mg daily. "
        "He is allergic to Penicillin. I suspect acute myocardial infarction. I will order an EKG, Troponin levels, "
        "and administer Aspirin 324mg and sublingual Nitroglycerin immediately. Admitting to Cardiac ICU."
    )
    
    print(f"Sending Dictation: {dictation[:100]}...")
    
    try:
        # Start the Clinical Intake Workflow
        result = await client.execute_workflow(
            "ClinicalIntakeWorkflow",
            dictation,
            id=f"intake-{uuid.uuid4()}",
            task_queue="curaflow-tasks",
        )
        print("Success! Workflow completed.")
        import json
        print(json.dumps(result, indent=2))
    except Exception as e:
        print(f"Workflow failed: {e}")


async def test_lab_watcher(client: Client):
    print("\n=== Testing Lab Result Anomaly Watcher ===")
    patient_id = f"patient-{uuid.uuid4()}"
    
    # 1. Start the long-running workflow with baseline data
    print(f"Starting long-running Lab Result Watcher for patient: {patient_id}")
    baseline_data = {
        "baseline": "Hypertension, CKD Stage 3",
        "medications": ["Lisinopril", "Spironolactone"]
    }
    
    # We use start_workflow instead of execute_workflow because it runs infinitely
    handle = await client.start_workflow(
        "LabResultWatcherWorkflow",
        baseline_data,
        id=patient_id,
        task_queue="curaflow-tasks",
    )
    
    # 2. Wait a moment
    await asyncio.sleep(2)
    
    # 3. Send a completely normal lab result
    normal_lab = "Serum Potassium: 4.2 mEq/L (Normal range 3.6-5.2)"
    print(f"Sending normal lab result: {normal_lab}")
    await handle.signal("add_lab_result", normal_lab)
    
    await asyncio.sleep(10) # Give AI time to process
    
    # 4. Send a dangerous lab result!
    # Lisinopril + Spironolactone + High Potassium = Life threatening hyperkalemia
    dangerous_lab = "Serum Potassium: 6.8 mEq/L (Critical High)"
    print(f"Sending dangerous lab result: {dangerous_lab}")
    await handle.signal("add_lab_result", dangerous_lab)
    
    await asyncio.sleep(10) # Give AI time to process and trigger alert
    
    print("Discharging patient to close the workflow...")
    await handle.signal("discharge_patient")
    
    result = await handle.result()
    print(f"Patient Discharged. Final Workflow Result: {result}")


async def main():
    # Use localhost:7233 if port forwarding, or cluster internal if running inside cluster
    temporal_url = "temporal-frontend.ai-core.svc.cluster.local:7233"
    print(f"Connecting to Temporal at {temporal_url}...")
    try:
        client = await Client.connect(temporal_url, namespace="default")
        print("Connected!")
    except Exception as e:
        print(f"Failed to connect: {e}")
        sys.exit(1)
        
    await test_scribe_and_coder(client)
    await test_lab_watcher(client)

if __name__ == "__main__":
    asyncio.run(main())