import requests
import json
import argparse

def test_hf_endpoint(endpoint_url, api_token=None):
    """Test the HuggingFace Inference Endpoint with different formats"""

    # Format 1: Wrapped OpenAI format
    payload1 = {
        "inputs": {
            "messages": [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "Explain quantum computing in simple terms."}
            ],
            "max_tokens": 256,
            "temperature": 0.7
        }
    }

    # Format 2: Simple string
    payload2 = {
        "inputs": "Tell me about AI."
    }

    # Prepare headers
    headers = {
        "Content-Type": "application/json",
    }
    
    if api_token:
        headers["Authorization"] = f"Bearer {api_token}"

    # Test Format 1
    print("Testing Format 1: Wrapped OpenAI format...")
    print(f"Payload: {json.dumps(payload1, indent=2)}")
    
    try:
        response = requests.post(endpoint_url, headers=headers, json=payload1)
        print(f"Status Code: {response.status_code}")
        print(f"Response: {json.dumps(response.json(), indent=2)}")
    except Exception as e:
        print(f"Error: {str(e)}")

    print("\n" + "-"*50 + "\n")

    # Test Format 2
    print("Testing Format 2: Simple string...")
    print(f"Payload: {json.dumps(payload2, indent=2)}")
    
    try:
        response = requests.post(endpoint_url, headers=headers, json=payload2)
        print(f"Status Code: {response.status_code}")
        print(f"Response: {json.dumps(response.json(), indent=2)}")
    except Exception as e:
        print(f"Error: {str(e)}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test HuggingFace Inference Endpoint")
    parser.add_argument("--url", type=str, required=True, help="Endpoint URL")
    parser.add_argument("--token", type=str, help="API token (if needed)")
    
    args = parser.parse_args()
    test_hf_endpoint(args.url, args.token)