import httpx
import time

BASE_URL = "http://127.0.0.1:8000/api/v1"
API_KEY = "storyline-secret-key-123"

client = httpx.Client(timeout=30.0)

def test_health():
    print("Testing Health Check...")
    response = client.get(f"{BASE_URL}/storyline/health")
    print(f"Status: {response.status_code}, Body: {response.json()}")

def test_chat_unauthorized():
    print("\nTesting Chat (No Auth)...")
    response = client.post(f"{BASE_URL}/storyline/chat", json={"prompt": "Hello"})
    print(f"Status: {response.status_code}, Body: {response.json()}")

def test_chat_authorized():
    print("\nTesting Chat (Authorized)...")
    headers = {"X-API-Key": API_KEY}
    response = client.post(
        f"{BASE_URL}/storyline/chat", 
        json={"prompt": "How do I use Storyline?", "context": "Module 1"},
        headers=headers
    )
    print(f"Status: {response.status_code}, Body: {response.json()}")

def test_rate_limit():
    print("\nTesting Rate Limit (making multiple calls)...")
    headers = {"X-API-Key": API_KEY}
    for i in range(22):
        response = httpx.post(
            f"{BASE_URL}/storyline/chat", 
            json={"prompt": f"Test {i}"},
            headers=headers
        )
        if response.status_code == 429:
            print(f"Rate limit hit at call {i+1}!")
            return
    print("Rate limit not hit (this is unexpected if limit is 20).")

def test_gemini():
    print("\nTesting Gemini Endpoint...")
    headers = {"X-API-Key": API_KEY}
    response = client.post(f"{BASE_URL}/gemini/chat", json={"prompt": "Hello Gemini"}, headers=headers)
    print(f"Status: {response.status_code}, Body: {response.json()}")

def test_openai():
    print("\nTesting OpenAI Endpoint...")
    headers = {"X-API-Key": API_KEY}
    response = client.post(f"{BASE_URL}/openai/chat", json={"prompt": "Hello OpenAI"}, headers=headers)
    print(f"Status: {response.status_code}, Body: {response.json()}")

def test_module_management():
    print("\nTesting Module Management...")
    headers = {"X-API-Key": API_KEY}
    
    # 1. Create a module
    print("Creating module 'Articulate 101'...")
    response = client.post(
        f"{BASE_URL}/admin/modules", 
        json={"name": "Articulate 101"}, 
        headers=headers
    )
    if response.status_code == 400:
        print("Module already exists, getting list...")
        response = client.get(f"{BASE_URL}/admin/modules", headers=headers)
        module_data = response.json()[0]
    else:
        module_data = response.json()
    
    new_module_key = module_data["secret_key"]
    print(f"Module created/found. Secret Key: {new_module_key}")

    # 2. Use the new module key to call Gemini
    print("\nTesting Gemini with NEW module key...")
    module_headers = {"X-API-Key": new_module_key}
    response = client.post(
        f"{BASE_URL}/gemini/chat", 
        json={"prompt": "Can you see this with my module key?"}, 
        headers=module_headers
    )
    print(f"Status: {response.status_code}, Body: {response.json()}")

if __name__ == "__main__":
    try:
        test_health()
        test_chat_unauthorized()
        test_chat_authorized()
        test_gemini()
        test_openai()
        test_module_management()
        # test_rate_limit()
    except Exception as e:
        print(f"Error: {e}")