import openai
import argparse
import os
from pathlib import Path

def main():
    parser = argparse.ArgumentParser(description="Test API with different modes.")
    parser.add_argument("--mode", choices=["claude", "local"], default="local", 
                       help="API mode: 'claude' for Anthropic API, 'local' for localhost API")
    parser.add_argument("--port", type=int, default=1333, 
                       help="Port number for localhost API (only used in local mode)")
    parser.add_argument("--model", type=str, default="/proj/inf-scaling/mistralai/Devstral-Small-2507", help="Model name to use (overrides default for mode)")
    args = parser.parse_args()

    if args.mode == "claude":
        # Load environment variables from .env file
        from sweagent.utils.config import load_environment_variables
        from sweagent import REPO_ROOT
        
        print(f"Looking for .env file in:")
        print(f"  Current directory: {Path.cwd() / '.env'}")
        print(f"  Repo root: {REPO_ROOT / '.env'}")
        
        load_environment_variables()
        
        # Get API key from environment variable
        api_key = os.getenv("ANTHROPIC_API_KEY")
        if not api_key:
            print("Error: ANTHROPIC_API_KEY environment variable not set")
            print("Available environment variables:", [k for k in os.environ.keys() if 'ANTHROPIC' in k])
            return

        print(f"API key found: {api_key[:20]}...")

        client = openai.Client(
            base_url="https://api.anthropic.com/v1",
            api_key=api_key
        )
        model = "claude-sonnet-4-20250514"
        
    else:  # local mode
        client = openai.Client(
            base_url=f"http://34.66.144.59:{args.port}/v1",
            api_key="None",
        )
        model = args.model

    print(f"Testing {args.mode} mode")

    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "user", "content": "List 3 countries and their capitals."},
        ],
        temperature=0,
        max_tokens=64,
    )

    print(response)

if __name__ == "__main__":
    main()