1#!/usr/bin/env python3
"""
OpenAI API client supporting z.ai and Ollama endpoints.
Sends a sequence of numbers and asks what the first number was.
"""

import argparse
import json
import os
from openai import OpenAI


# Endpoint configurations
ENDPOINTS = {
    "z.ai": {
        "base_url": "https://api.z.ai/api/coding/paas/v4"
    },
    "ollama": {
        "base_url": "https://ollama.com/v1"
    }
}


def list_and_select_model(client):
    """List available models and prompt user to select one."""
    print("Fetching available models...")
    print("-" * 50)
    
    try:
        models = client.models.list()
        model_list = list(models)
        
        if not model_list:
            print("No models available.")
            return None
        
        # Extract model IDs and display them
        model_ids = []
        for i, model in enumerate(model_list, 1):
            model_id = model.id if hasattr(model, 'id') else str(model)
            model_ids.append(model_id)
            print(f"  {i}. {model_id}")
        
        print("-" * 50)
        
        # Prompt user to select a model
        while True:
            try:
                selection = input("Select a model by number (or press Enter for first): ").strip()
                if selection == "":
                    return model_ids[0]
                idx = int(selection) - 1
                if 0 <= idx < len(model_ids):
                    return model_ids[idx]
                else:
                    print(f"Please enter a number between 1 and {len(model_ids)}")
            except ValueError:
                print("Please enter a valid number.")
    except Exception as e:
        print(f"Error fetching models: {e}")
        return None


def main():
    parser = argparse.ArgumentParser(description="OpenAI API client for z.ai or Ollama")
    parser.add_argument(
        "-n", "--number", type=int, default=10,
        help="The upper limit n for the sequence 1 to n (default: 10)"
    )
    parser.add_argument(
        "-m", "--max-tokens", type=int, default=4096,
        help="Maximum tokens in the output (default: 4096)"
    )
    parser.add_argument(
        "-e", "--endpoint", choices=["z.ai", "ollama"], default="z.ai",
        help="API endpoint to use: 'z.ai' or 'ollama' (default: z.ai)"
    )
    parser.add_argument(
        "-t", "--token", type=str, default=None,
        help="API token for authentication (for z.ai endpoint). Can also be set via OAI_TOKEN environment variable."
    )
    args = parser.parse_args()

    # Get endpoint configuration
    endpoint_config = ENDPOINTS[args.endpoint]
    
    api_key = args.token or os.environ.get("OAI_TOKEN")
    if not api_key:
        print("Error: API token required for z.ai endpoint.")
        print("Provide it via --token argument or OAI_TOKEN environment variable.")
        return

    # Initialize the client with the selected endpoint
    client = OpenAI(
        api_key=api_key,
        base_url=endpoint_config["base_url"]
    )
    
    print(f"Using endpoint: {args.endpoint}")
    print(f"Base URL: {endpoint_config['base_url']}")
    print("-" * 50)

    # List models and let user select one
    selected_model = list_and_select_model(client)
    if selected_model is None:
        print("No model selected. Exiting.")
        return
    
    print(f"\nUsing model: {selected_model}")
    print("-" * 50)

    # Build the prompt: numbers 1 to n, then the question
    numbers = ", ".join(str(i) for i in range(1, args.number + 1))
    prompt = f"{numbers}\n\nWhat was the first number?"

    print(f"Sending prompt with numbers 1 to {args.number}...")
    print(f"Prompt:\n{prompt}\n")
    print("-" * 50)

    # Send the chat completion request
    response = client.chat.completions.create(
        model=selected_model,
        messages=[
            {"role": "user", "content": prompt}
        ],
        max_tokens=args.max_tokens,
        stream=False
    )

    # Print the response content
    print("Response:")
    print("-" * 50)
    if response.choices and len(response.choices) > 0:
        print(response.choices[0].message.content)
    else:
        print("No response content received.")

    print("\n" + "-" * 50)
    print("Response Statistics:")
    print("-" * 50)

    # Print usage statistics if available
    if hasattr(response, 'usage') and response.usage:
        print(f"Prompt tokens: {response.usage.prompt_tokens}")
        print(f"Completion tokens: {response.usage.completion_tokens}")
        print(f"Total tokens: {response.usage.total_tokens}")
    else:
        print("No usage statistics available.")

    # Print the full raw response for debugging
    print("\n" + "-" * 50)
    print("Full Response Object (raw):")
    print("-" * 50)
    # Convert to dict if possible, otherwise use __dict__
    try:
        if hasattr(response, 'model_dump'):
            response_dict = response.model_dump()
        elif hasattr(response, 'to_dict'):
            response_dict = response.to_dict()
        else:
            response_dict = response.__dict__
        print(json.dumps(response_dict, indent=2, default=str))
    except Exception as e:
        print(f"Could not serialize response: {e}")
        print(response)


if __name__ == "__main__":
    main()