"""
Test script for connecting to a remote Ollama server.
"""

import os
import sys
import requests
import json
import argparse

# Import local test configuration
from local_test_config import config

def test_ollama_connection():
    """Test the connection to the Ollama API."""
    print(f"Testing connection to Ollama API at: {config.OLLAMA_API_URL}")

    try:
        # Try to connect to Ollama API via OpenWebUI
        # OpenWebUI exposes Ollama models through its own API
        response = requests.get(f"{config.OPENWEBUI_URL}/api/models", timeout=config.API_TIMEOUT)
        response.raise_for_status()

        # Print the models from OpenWebUI
        print("Connection successful!")
        print("Available models:")
        models_data = response.json()
        if isinstance(models_data, list):
            for model in models_data:
                print(f"- {model.get('id')}")
        else:
            print(f"Unexpected response format: {models_data}")

        return True
    except requests.exceptions.Timeout as e:
        print(f"ERROR: Timeout connecting to Ollama API: {str(e)}. The request exceeded the {config.API_TIMEOUT} second timeout.")
        return False
    except requests.exceptions.ConnectionError as e:
        print(f"ERROR: Connection error to Ollama API: {str(e)}. Please check if Ollama is running at {config.OLLAMA_API_URL}.")
        return False
    except Exception as e:
        print(f"ERROR: Error connecting to Ollama API: {str(e)}")
        return False

def test_generate_response(model_id=None, prompt=None):
    """Test generating a response from the model."""
    if model_id is None:
        model_id = config.DEFAULT_MODEL

    if prompt is None:
        prompt = "What is the capital of France?"

    print(f"Testing model: {model_id}")
    print(f"Prompt: {prompt}")

    try:
        # Prepare the request for OpenWebUI's chat completions API
        request_json = {
            "model": model_id,
            "messages": [
                {"role": "user", "content": prompt}
            ],
            "stream": False,
            "temperature": 0.7,
            "max_tokens": 500
        }

        # Make the API call to OpenWebUI's chat completions API
        print(f"Sending request to OpenWebUI API at: {config.OPENWEBUI_URL}/api/chat/completions")

        response = requests.post(
            f"{config.OPENWEBUI_URL}/api/chat/completions",
            headers={"Content-Type": "application/json"},
            json=request_json,
            timeout=config.API_TIMEOUT
        )

        response.raise_for_status()

        # Parse the response
        response_data = response.json()

        # Print the response (OpenWebUI uses OpenAI-compatible format)
        print("\nResponse:")
        content = response_data.get('choices', [{}])[0].get('message', {}).get('content', 'No content in response')
        print(content)

        return True
    except requests.exceptions.Timeout as e:
        print(f"ERROR: Timeout generating response: {str(e)}. The request exceeded the {config.API_TIMEOUT} second timeout.")
        return False
    except requests.exceptions.ConnectionError as e:
        print(f"ERROR: Connection error to Ollama API: {str(e)}. Please check if Ollama is running at {config.OLLAMA_API_URL}.")
        return False
    except Exception as e:
        print(f"ERROR: Error generating response: {str(e)}")
        return False

def test_openwebui_connection():
    """Test the connection to OpenWebUI."""
    print(f"Testing connection to OpenWebUI at: {config.OPENWEBUI_URL}")

    try:
        # Try to connect to OpenWebUI
        response = requests.get(f"{config.OPENWEBUI_URL}/api/health", timeout=config.API_TIMEOUT)
        response.raise_for_status()

        # Print the response
        print("Connection successful!")
        print(f"OpenWebUI status: {response.text}")

        return True
    except requests.exceptions.Timeout as e:
        print(f"ERROR: Timeout connecting to OpenWebUI: {str(e)}. The request exceeded the {config.API_TIMEOUT} second timeout.")
        return False
    except requests.exceptions.ConnectionError as e:
        print(f"ERROR: Connection error to OpenWebUI: {str(e)}. Please check if OpenWebUI is running at {config.OPENWEBUI_URL}.")
        return False
    except Exception as e:
        print(f"ERROR: Error connecting to OpenWebUI: {str(e)}")
        return False

def main():
    """Main function."""
    parser = argparse.ArgumentParser(description='Test remote Ollama and OpenWebUI connections')
    parser.add_argument('--model', type=str, default=config.DEFAULT_MODEL, 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')
    parser.add_argument('--ollama-url', type=str, default=config.OLLAMA_API_URL, help='Ollama API URL')
    parser.add_argument('--openwebui-url', type=str, default=config.OPENWEBUI_URL, help='OpenWebUI URL')
    parser.add_argument('--timeout', type=int, default=config.API_TIMEOUT, help='API timeout in seconds')

    args = parser.parse_args()

    # Update configuration with command-line arguments
    config.OLLAMA_API_URL = args.ollama_url
    config.OPENWEBUI_URL = args.openwebui_url
    config.API_TIMEOUT = args.timeout

    print("=== Remote Connection Test ===")
    print(f"Ollama API URL: {config.OLLAMA_API_URL}")
    print(f"OpenWebUI URL: {config.OPENWEBUI_URL}")
    print(f"API Timeout: {config.API_TIMEOUT} seconds")
    print()

    # Test Ollama connection
    print("=== Testing Ollama Connection ===")
    ollama_success = test_ollama_connection()
    print()

    # Test OpenWebUI connection
    print("=== Testing OpenWebUI Connection ===")
    openwebui_success = test_openwebui_connection()
    print()

    # Test generating a response
    if ollama_success:
        print("=== Testing Model Response ===")
        test_generate_response(args.model, args.prompt)

    print("\n=== Test Summary ===")
    print(f"Ollama Connection: {'SUCCESS' if ollama_success else 'FAILED'}")
    print(f"OpenWebUI Connection: {'SUCCESS' if openwebui_success else 'FAILED'}")

if __name__ == "__main__":
    main()