#!/usr/bin/env python3
"""
Test script for OPRO session-based API
"""

import requests
import json

BASE_URL = "http://127.0.0.1:8010"

def print_section(title):
    print(f"\n{'='*60}")
    print(f"  {title}")
    print(f"{'='*60}\n")

def test_session_workflow():
    """Test the complete session-based workflow."""
    
    print_section("1. Create Session")
    
    # Create a new session
    response = requests.post(f"{BASE_URL}/opro/session/create")
    result = response.json()
    
    if not result.get("success"):
        print(f"❌ Failed to create session: {result}")
        return
    
    session_id = result["data"]["session_id"]
    print(f"✅ Session created: {session_id}")
    print(f"   Session name: {result['data']['session_name']}")
    
    print_section("2. Create First Run in Session")
    
    # Create first run
    create_req = {
        "task_description": "将中文翻译成英文",
        "test_cases": [
            {"input": "你好", "expected_output": "Hello"},
            {"input": "谢谢", "expected_output": "Thank you"}
        ],
        "session_id": session_id
    }
    
    response = requests.post(f"{BASE_URL}/opro/create", json=create_req)
    result = response.json()
    
    if not result.get("success"):
        print(f"❌ Failed to create run: {result}")
        return
    
    run1_id = result["data"]["run_id"]
    print(f"✅ Run 1 created: {run1_id}")
    print(f"   Task: {result['data']['task_description']}")
    
    print_section("3. Create Second Run in Same Session")
    
    # Create second run in same session
    create_req2 = {
        "task_description": "将英文翻译成中文",
        "test_cases": [
            {"input": "Hello", "expected_output": "你好"},
            {"input": "Thank you", "expected_output": "谢谢"}
        ],
        "session_id": session_id
    }
    
    response = requests.post(f"{BASE_URL}/opro/create", json=create_req2)
    result = response.json()
    
    if not result.get("success"):
        print(f"❌ Failed to create run 2: {result}")
        return
    
    run2_id = result["data"]["run_id"]
    print(f"✅ Run 2 created: {run2_id}")
    print(f"   Task: {result['data']['task_description']}")
    
    print_section("4. Get Session Details")
    
    response = requests.get(f"{BASE_URL}/opro/session/{session_id}")
    result = response.json()
    
    if not result.get("success"):
        print(f"❌ Failed to get session: {result}")
        return
    
    print(f"✅ Session details:")
    print(f"   Session ID: {result['data']['session_id']}")
    print(f"   Session name: {result['data']['session_name']}")
    print(f"   Number of runs: {result['data']['num_runs']}")
    print(f"   Runs:")
    for run in result['data']['runs']:
        print(f"     - {run['run_id'][:8]}... : {run['task_description']}")
    
    print_section("5. List All Sessions")
    
    response = requests.get(f"{BASE_URL}/opro/sessions")
    result = response.json()
    
    if not result.get("success"):
        print(f"❌ Failed to list sessions: {result}")
        return
    
    print(f"✅ Total sessions: {len(result['data']['sessions'])}")
    for session in result['data']['sessions']:
        print(f"   - {session['session_name']}: {session['num_runs']} runs")
    
    print_section("✅ All Tests Passed!")

if __name__ == "__main__":
    try:
        # Check if server is running
        response = requests.get(f"{BASE_URL}/health")
        if response.status_code != 200:
            print("❌ Server is not running. Please start it with:")
            print("   uvicorn _qwen_xinference_demo.api:app --host 127.0.0.1 --port 8010")
            exit(1)
        
        test_session_workflow()
        
    except requests.exceptions.ConnectionError:
        print("❌ Cannot connect to server. Please start it with:")
        print("   uvicorn _qwen_xinference_demo.api:app --host 127.0.0.1 --port 8010")
        exit(1)
    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback
        traceback.print_exc()
        exit(1)