#!/usr/bin/env python3
"""
Test script for AI_Talk.py core functionality
Tests real API calls to Groq and Gemini models
"""

import os
import sys
from dotenv import load_dotenv
from openai import OpenAI
import anthropic
from groq import Groq

# Load environment variables
load_dotenv()

class AIFriendsTalkTest:
    def __init__(self):
        self.setup_apis()
        self.setup_characters()
        
    def setup_apis(self):
        """Setup API clients following day1.ipynb structure"""
        print("Setting up APIs...")
        
        # Groq client for Alex character
        self.groq_client = Groq()
        print("OK Groq client initialized")
        
        # Gemini clients via OpenAI interface for Blake and Charlie
        google_api_key = os.getenv('GOOGLE_API_KEY')
        if google_api_key:
            self.gemini_client = OpenAI(
                api_key=google_api_key,
                base_url="https://generativelanguage.googleapis.com/v1beta/openai/"
            )
            print("OK Gemini client initialized")
        else:
            print("ERROR Google API key not found")
            
    def setup_characters(self):
        """Define AI characters with enhanced personalities"""
        self.characters = {
            "Alex": {
                "model": "groq",
                "model_name": "llama3-70b-8192",
                "personality": "You are Alex, a witty and charismatic AI debater who thrives on intellectual challenges. You have a sharp sense of humor and love to play devil's advocate. Respond with engaging messages that spark curiosity and debate.",
                "color": "#FF6B6B"
            },
            "Blake": {
                "model": "gemini2",
                "model_name": "gemini-2.0-flash",
                "personality": "You are Blake, an imaginative and boundlessly optimistic AI who sees magic and possibility in everything. You're a natural storyteller who loves to paint vivid pictures with words. Write poetic responses that inspire and delight.",
                "color": "#4ECDC4"
            },
            "Charlie": {
                "model": "gemini1.5",
                "model_name": "gemini-1.5-flash",
                "personality": "You are Charlie, a thoughtful and systematic AI analyst who approaches every topic with scientific curiosity and methodical thinking. You love to break down complex ideas and examine evidence. Provide detailed, structured responses.",
                "color": "#45B7D1"
            }
        }
        print("OK Characters configured")
        
    def get_ai_response(self, character_name, topic, language="en"):
        """Get response from specific AI character"""
        character = self.characters[character_name]
        
        # Build prompt
        lang_instruction = {
            "en": "Please respond in English with an engaging message.",
            "vi": "Vui lòng trả lời bằng tiếng Việt với tin nhắn hấp dẫn.",
            "de": "Bitte antworten Sie auf Deutsch mit einer ansprechenden Nachricht."
        }
        
        prompt = f"{character['personality']}\n\n{lang_instruction[language]}\n\nTopic: {topic}\n\nRespond as {character_name} with your thoughts on this topic:"
        
        try:
            # Make REAL API calls to actual AI models
            if character["model"] == "groq":
                # Alex character using Groq llama3-70b-8192
                print(f"Making API call to Groq for {character_name}...")
                response = self.groq_client.chat.completions.create(
                    model=character["model_name"],
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=200,
                    temperature=0.8
                )
                return response.choices[0].message.content
                
            elif character["model"] == "gemini2":
                # Blake character using Gemini 2.0 Flash
                if not hasattr(self, 'gemini_client'):
                    return "Sorry, Blake is unavailable (Google API key not configured)."
                    
                print(f"Making API call to Gemini 2.0 for {character_name}...")
                response = self.gemini_client.chat.completions.create(
                    model=character["model_name"],
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=200,
                    temperature=0.8
                )
                return response.choices[0].message.content
                
            elif character["model"] == "gemini1.5":
                # Charlie character using Gemini 1.5 Flash
                if not hasattr(self, 'gemini_client'):
                    return "Sorry, Charlie is unavailable (Google API key not configured)."
                    
                print(f"Making API call to Gemini 1.5 for {character_name}...")
                response = self.gemini_client.chat.completions.create(
                    model=character["model_name"],
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=200,
                    temperature=0.8
                )
                return response.choices[0].message.content
                
        except Exception as e:
            error_msg = str(e)[:100] if str(e) else "Unknown error"
            return f"[{character_name}] API Error: {error_msg}..."

def main():
    print("=== AI Friends Talk Test ===")
    print("Testing real API integration with Groq and Gemini models\n")
    
    # Initialize the AI system
    ai_friends = AIFriendsTalkTest()
    
    # Test topic
    test_topic = "Should AI have emotions like humans?"
    print(f"Test Topic: {test_topic}\n")
    
    # Test each character
    characters = ["Alex", "Blake", "Charlie"]
    
    for character in characters:
        print(f"--- Testing {character} ---")
        try:
            response = ai_friends.get_ai_response(character, test_topic)
            print(f"{character}: {response}\n")
        except Exception as e:
            print(f"Error testing {character}: {e}\n")
    
    print("=== Test Complete ===")

if __name__ == "__main__":
    main()