#!/usr/bin/env python3
"""
Simple test script for team chat functionality.
"""

import requests
import argparse
import sys

def main():
    """Main function."""
    parser = argparse.ArgumentParser(description='Simple test for team chat functionality')
    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('--user-id', type=str, default='test_user', help='User ID for testing')
    parser.add_argument('--chat-id', type=str, help='Existing chat ID to use (optional)')
    parser.add_argument('--message', type=str, default='Hello team!', help='Message to send')
    parser.add_argument('--with-mention', action='store_true', help='Add @ai mention to message')

    args = parser.parse_args()

    print(f"API URL: {args.api_url}")
    print(f"Model: {args.model}")
    print(f"User ID: {args.user_id}")

    # Create a chat if no chat ID is provided
    chat_id = args.chat_id
    if not chat_id:
        print("\nCreating a new team chat...")
        try:
            response = requests.post(
                f"{args.api_url}/chats",
                headers={"Content-Type": "application/json"},
                json={
                    "user_id": args.user_id,
                    "title": "Simple Team Chat Test",
                    "model_id": args.model,
                    "is_team_chat": True
                },
                timeout=30
            )

            response.raise_for_status()
            chat = response.json()
            chat_id = chat.get("id")

            print(f"Created team chat with ID: {chat_id}")
            print(f"OpenWebUI Channel ID: {chat.get('openwebui_channel_id', 'Not available')}")
            print(f"Team Members: {chat.get('team_members', [])}")
        except Exception as e:
            print(f"Error creating team chat: {str(e)}")
            sys.exit(1)
    else:
        print(f"\nUsing existing chat ID: {chat_id}")

    # Prepare message
    message = args.message
    if args.with_mention:
        message = f"@ai {message}"
        print(f"\nSending message WITH @ai mention: {message}")
    else:
        print(f"\nSending message WITHOUT @ai mention: {message}")

    # Send message
    try:
        response = requests.post(
            f"{args.api_url}/chats/{chat_id}/messages",
            headers={"Content-Type": "application/json"},
            json={
                "message": message,
                "user_id": args.user_id
            },
            timeout=10
        )

        response.raise_for_status()
        result = response.json()

        print("\nMessage sent successfully!")
        print(f"AI Response: {result.get('content', 'No content in response')}")

        print("\nNote: When using the API directly, the AI will always respond")
        print("The @ai mention functionality only applies to messages sent through OpenWebUI channels")
        print(f"\nYou can check the OpenWebUI interface at: http://104.225.217.215:8080/")
        print(f"Chat ID: {chat_id}")
    except Exception as e:
        print(f"Error sending message: {str(e)}")
        sys.exit(1)

if __name__ == "__main__":
    main()