"""
Test script for the direct Ollama endpoint.
"""

import requests
import json
import argparse
import time

def test_direct_ollama(api_url, model_id, prompt):
    """Test the direct Ollama endpoint."""
    print(f"Testing direct Ollama endpoint at: {api_url}/test-ollama-direct")
    print(f"Model: {model_id}")
    print(f"Prompt: {prompt}")

    try:
        start_time = time.time()

        response = requests.post(
            f"{api_url}/test-ollama-direct",
            headers={"Content-Type": "application/json"},
            json={
                "model": model_id,
                "prompt": prompt
            },
            timeout=300  # 5 minutes timeout
        )

        end_time = time.time()
        elapsed_time = end_time - start_time

        response.raise_for_status()

        result = response.json()
        print(f"\nResponse (took {elapsed_time:.2f} seconds):")

        # Handle different response formats
        content = None
        if 'response' in result and isinstance(result['response'], dict) and 'message' in result['response']:
            content = result['response']['message'].get('content')
        elif 'response' in result and isinstance(result['response'], str):
            content = result['response']
        else:
            content = str(result)

        print(content or 'No content in response')

        return True
    except Exception as e:
        print(f"ERROR: Direct Ollama test failed: {str(e)}")
        return False

def main():
    """Main function."""
    parser = argparse.ArgumentParser(description='Test the direct Ollama endpoint')
    parser.add_argument('--api-url', type=str, default='http://localhost:5252', help='AI service API URL')
    parser.add_argument('--model', type=str, default='llama3.1', help='Model to use for testing')
    parser.add_argument('--prompt', type=str, default='What is the capital of France?', help='Prompt to use for testing')

    args = parser.parse_args()

    print("=== Direct Ollama Test ===")
    print(f"API URL: {args.api_url}")
    print(f"Model: {args.model}")
    print(f"Prompt: {args.prompt}")
    print()

    # Test direct Ollama endpoint
    success = test_direct_ollama(args.api_url, args.model, args.prompt)

    # Print summary
    print("\n=== Test Summary ===")
    print(f"Direct Ollama Test: {'SUCCESS' if success else 'FAILED'}")

if __name__ == "__main__":
    main()