#!/usr/bin/env python3
import os
import argparse
from pathlib import Path
import dotenv
from openai import OpenAI    # pip install --upgrade openai>=1.0.0

# Load environment variables from .env file
dotenv.load_dotenv()

def test_rits_inference(endpoint: str, model_name: str) -> None:
    """
    Call the RITS endpoint and print a short completion.
    """
    base_url = (f"https://inference-3scale-apicast-production.apps.rits.fmaas.res.ibm.com/{endpoint}/v1")

    api_key = os.getenv("RITS_API_KEY")
    if not api_key:
        raise RuntimeError(
            "RITS_API_KEY not found! Please create a .env file in the SWE-agent directory "
            "with the line: RITS_API_KEY=your_api_key_here"
        )

    client = OpenAI(
        base_url=base_url,
        api_key=api_key,                     # required by the SDK, though ignored by RITS
        default_headers={"RITS_API_KEY": api_key},
        timeout=60,
    )

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

    print("=== Response ===")
    print(resp.choices[0].message.content)
    print(f"Usage: {resp.usage}")

def main() -> None:
    parser = argparse.ArgumentParser(description="Quick RITS inference test")
    parser.add_argument(
        "--endpoint",
        default="deepseek-v3-h200",
        help="Endpoint slug as shown in the RITS UI (e.g. deepseek-v3-h200)",
    )
    parser.add_argument(
        "--model",
        default="deepseek-ai/DeepSeek-V3",
        help="Fully-qualified model name (e.g. deepseek-ai/DeepSeek-V3)",
    )
    args = parser.parse_args()
    test_rits_inference(args.endpoint, args.model)

if __name__ == "__main__":
    main()
