import requests
import argparse
import sys

BASE_URL = "http://127.0.0.1:8000"

def test_openai_style(token):
    headers = {"Authorization": f"Bearer {token}"}

    # Send OpenAI-style Chat Messages
    payload = {
        "messages": [
            {"role": "system", "content": "You are a concise AI assistant."},
            {"role": "user", "content": "What is the main benefit of open-source LLMs?"}
        ],
        "max_tokens": 60
    }

    print(f"🤖 Sending Chat Completion request to {BASE_URL}...")
    try:
        response = requests.post(f"{BASE_URL}/v1/chat/completions", json=payload, headers=headers)

        if response.status_code == 200:
            result = response.json()
            print("\n✨ Model Response:")
            print(result['choices'][0]['message']['content'])
        elif response.status_code == 401:
            print("❌ Error: 401 Unauthorized. Please check if your token is valid and matches the SECRET_KEY.")
        else:
            print(f"❌ Error {response.status_code}: {response.text}")
    except Exception as e:
        print(f"❌ Connection failed: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test the Qwen API with a JWT token.")
    parser.add_argument("--token", type=str, help="The JWT token generated by generate_token.py")
    args = parser.parse_args()

    if not args.token:
        print("❌ Error: Please provide a token using --token <YOUR_TOKEN>")
        sys.exit(1)

    test_openai_style(args.token)