#!/usr/bin/env python3
"""
Test script for OpenWebUI bot.

This script tests the connection to OpenWebUI and the bot's functionality.
"""

import sys
import os
import asyncio
import aiohttp
import json
import logging

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('bot_test')

# Configuration
OPENWEBUI_URL = "http://104.225.217.215:8080"
API_KEY = "GdCU4ieYDqHsLfH2"
MODEL_ID = "llama3.1"

async def test_openwebui_api():
    """Test the OpenWebUI API connection."""
    logger.info("Testing OpenWebUI API connection...")

    try:
        async with aiohttp.ClientSession() as session:
            # Test 1: Check if the API is accessible
            logger.info("Testing API access...")
            headers = {"Authorization": f"Bearer {API_KEY}"}
            async with session.get(
                f"{OPENWEBUI_URL}/api/models",
                headers=headers
            ) as response:
                if response.status == 200:
                    models = await response.json()
                    logger.info(f"API accessible. Available models: {json.dumps(models, indent=2)}")
                else:
                    logger.error(f"API error: {response.status} - {await response.text()}")
                    return False

            # Test 2: Test authentication
            logger.info("Testing API authentication...")
            headers = {"Authorization": f"Bearer {API_KEY}"}
            async with session.get(
                f"{OPENWEBUI_URL}/api/channels",
                headers=headers
            ) as response:
                if response.status == 200:
                    channels = await response.json()
                    logger.info(f"Authentication successful. Channels: {json.dumps(channels, indent=2)}")
                else:
                    logger.error(f"Authentication error: {response.status} - {await response.text()}")
                    return False

            # Test 3: Test chat completion
            logger.info("Testing chat completion...")
            payload = {
                "model": MODEL_ID,
                "messages": [
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": "Say hello in one short sentence."}
                ],
                "stream": False
            }

            async with session.post(
                f"{OPENWEBUI_URL}/api/chat/completions",
                headers=headers,
                json=payload
            ) as response:
                if response.status == 200:
                    result = await response.json()
                    if "choices" in result and len(result["choices"]) > 0:
                        content = result["choices"][0]["message"]["content"]
                        logger.info(f"Chat completion successful. Response: {content}")
                    else:
                        logger.error(f"Chat completion error: Invalid response format - {json.dumps(result, indent=2)}")
                        return False
                else:
                    logger.error(f"Chat completion error: {response.status} - {await response.text()}")
                    return False

        logger.info("All API tests passed!")
        return True

    except Exception as e:
        logger.error(f"API test error: {str(e)}")
        return False

async def test_channel_message():
    """Test sending a message to a channel."""
    logger.info("Testing channel message...")

    try:
        async with aiohttp.ClientSession() as session:
            # Get available channels
            headers = {"Authorization": f"Bearer {API_KEY}"}
            async with session.get(
                f"{OPENWEBUI_URL}/api/channels",
                headers=headers
            ) as response:
                if response.status == 200:
                    channels = await response.json()
                    if not channels or len(channels) == 0:
                        logger.warning("No channels available. Creating a test channel...")

                        # Create a test channel
                        channel_data = {
                            "name": "Bot Test Channel",
                            "description": "Channel for testing the bot"
                        }

                        async with session.post(
                            f"{OPENWEBUI_URL}/api/channels",
                            headers=headers,
                            json=channel_data
                        ) as create_response:
                            if create_response.status == 200:
                                channel = await create_response.json()
                                channel_id = channel["id"]
                                logger.info(f"Created test channel with ID: {channel_id}")
                            else:
                                logger.error(f"Failed to create channel: {create_response.status} - {await create_response.text()}")
                                return False
                    else:
                        # Use the first available channel
                        channel_id = channels[0]["id"]
                        logger.info(f"Using existing channel with ID: {channel_id}")

                    # Send a test message to the channel
                    message_data = {
                        "content": "@ai Hello, are you working?"
                    }

                    async with session.post(
                        f"{OPENWEBUI_URL}/api/v1/channels/{channel_id}/messages/post",
                        headers=headers,
                        json=message_data
                    ) as message_response:
                        if message_response.status == 200:
                            logger.info("Test message sent successfully!")
                            logger.info("Check the channel to see if the bot responds.")
                            logger.info(f"Channel URL: {OPENWEBUI_URL}/channels/{channel_id}")
                            return True
                        else:
                            logger.error(f"Failed to send message: {message_response.status} - {await message_response.text()}")
                            return False
                else:
                    logger.error(f"Failed to get channels: {response.status} - {await response.text()}")
                    return False

    except Exception as e:
        logger.error(f"Channel test error: {str(e)}")
        return False

async def main():
    """Run all tests."""
    logger.info("Starting OpenWebUI bot tests...")

    # Test 1: OpenWebUI API
    api_result = await test_openwebui_api()
    if not api_result:
        logger.error("API tests failed!")
        return False

    # Test 2: Channel message
    channel_result = await test_channel_message()
    if not channel_result:
        logger.error("Channel tests failed!")
        return False

    logger.info("All tests completed successfully!")
    logger.info("Note: The bot should respond to the test message in the channel.")
    logger.info("Check the bot logs to see if it received and processed the message.")

    return True

if __name__ == "__main__":
    try:
        result = asyncio.run(main())
        if result:
            logger.info("Tests passed!")
            sys.exit(0)
        else:
            logger.error("Tests failed!")
            sys.exit(1)
    except KeyboardInterrupt:
        logger.info("Tests interrupted by user.")
        sys.exit(130)
    except Exception as e:
        logger.error(f"Unexpected error: {str(e)}")
        sys.exit(1)