#!/usr/bin/env python3
"""
Advanced Conversation Model for MemoryAI

This module provides enhanced conversation capabilities with:
- Multi-turn dialog management
- Context-aware response generation
- Personality and style control
- Emotion detection and response
- Topic tracking and continuity
"""

import os
import re
import random
from typing import List, Dict, Optional, Tuple
from datetime import datetime
import numpy as np
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import torch

# Check for GPU availability
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class ConversationModel:
    """
    Advanced conversation model with memory and context awareness.
    
    Features:
    - Multi-turn conversation handling
    - Context-aware responses
    - Emotion detection
    - Topic tracking
    - Personality control
    """
    
    def __init__(self, model_name: str = "facebook/blenderbot-400M-distill", 
                 embedding_model: str = "all-MiniLM-L6-v2"):
        """
        Initialize the conversation model.
        
        Args:
            model_name: Hugging Face model name for conversation
            embedding_model: Model for semantic embeddings
        """
        self.model_name = model_name
        self.embedding_model_name = embedding_model
        
        # Load models
        self.tokenizer = None
        self.model = None
        self.embedding_model = None
        self.conversation_pipeline = None
        
        self.load_models()
        
        # Conversation state
        self.conversation_history = []
        self.current_topic = "general"
        self.user_emotion = "neutral"
        self.conversation_length = 0
        
        # Personality settings
        self.personality = {
            "friendliness": 0.8,
            "humor": 0.6,
            "formality": 0.3,
            "verbosity": 0.7,
            "curiosity": 0.9
        }
        
        # Response enhancements
        self.response_enhancers = {
            "greetings": ["Hello!", "Hi there!", "Hey!", "Greetings!", "Nice to see you!"],
            "goodbyes": ["Goodbye!", "See you later!", "Take care!", "Bye!", "Have a great day!"],
            "agreements": ["Yes!", "Absolutely!", "I agree!", "Exactly!", "You're right!"],
            "disagreements": ["I see your point, but...", "That's interesting, however...", 
                            "I understand, but I think...", "That's a good perspective, but..."],
            "questions": ["What do you think about that?", "Does that make sense?", 
                         "How does that sound?", "What's your opinion?"]
        }
    
    def load_models(self):
        """Load the conversation and embedding models."""
        try:
            print(f"Loading conversation model: {self.model_name}")
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name).to(device)
            
            # Create conversation pipeline
            self.conversation_pipeline = pipeline(
                "conversational",
                model=self.model,
                tokenizer=self.tokenizer,
                device=0 if torch.cuda.is_available() else -1
            )
            
            print(f"Loading embedding model: {self.embedding_model_name}")
            self.embedding_model = SentenceTransformer(self.embedding_model_name)
            
            print("✅ Models loaded successfully!")
            
        except Exception as e:
            print(f"❌ Error loading models: {e}")
            # Fallback to simpler model
            print("Falling back to basic conversation model...")
            self.model_name = "microsoft/DialoGPT-small"
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name).to(device)
            self.conversation_pipeline = pipeline(
                "conversational",
                model=self.model,
                tokenizer=self.tokenizer,
                device=0 if torch.cuda.is_available() else -1
            )
    
    def detect_emotion(self, text: str) -> str:
        """Detect emotion in user input."""
        # Simple emotion detection based on keywords
        text_lower = text.lower()
        
        happy_keywords = ["happy", "joy", "excited", "great", "awesome", "wonderful", "love"]
        sad_keywords = ["sad", "unhappy", "depressed", "terrible", "awful", "hate"]
        angry_keywords = ["angry", "mad", "furious", "annoyed", "frustrated"]
        
        if any(keyword in text_lower for keyword in happy_keywords):
            return "happy"
        elif any(keyword in text_lower for keyword in sad_keywords):
            return "sad"
        elif any(keyword in text_lower for keyword in angry_keywords):
            return "angry"
        else:
            return "neutral"
    
    def detect_topic(self, text: str) -> str:
        """Detect the topic of conversation."""
        text_lower = text.lower()
        
        topic_keywords = {
            "technology": ["tech", "computer", "software", "hardware", "ai", "machine learning"],
            "sports": ["sports", "game", "football", "basketball", "soccer", "tennis"],
            "movies": ["movie", "film", "cinema", "actor", "actress", "director"],
            "music": ["music", "song", "band", "artist", "concert", "album"],
            "travel": ["travel", "vacation", "trip", "hotel", "flight", "destination"],
            "food": ["food", "restaurant", "cooking", "recipe", "cuisine", "dish"],
            "work": ["work", "job", "career", "office", "meeting", "project"],
            "personal": ["life", "family", "friend", "relationship", "feeling", "emotion"]
        }
        
        for topic, keywords in topic_keywords.items():
            if any(keyword in text_lower for keyword in keywords):
                return topic
                
        return "general"
    
    def generate_response(self, user_input: str, conversation_history: List[Dict] = None) -> str:
        """
        Generate a response to user input with full conversation context.
        
        Args:
            user_input: The user's message
            conversation_history: Previous conversation turns
            
        Returns:
            Generated response string
        """
        if conversation_history is None:
            conversation_history = []
            
        # Update conversation state
        self.user_emotion = self.detect_emotion(user_input)
        self.current_topic = self.detect_topic(user_input)
        self.conversation_length += 1
        
        # Add current input to history
        conversation_history.append({"role": "user", "content": user_input})
        
        try:
            # Generate response using the conversation model
            response = self.conversation_pipeline(
                conversation_history,
                max_length=150,
                temperature=0.7,
                top_p=0.9,
                repetition_penalty=1.2,
                num_return_sequences=1,
                do_sample=True  # Enable sampling for temperature/top_p to work
            )
            
            # Handle different response formats
            if isinstance(response, list) and len(response) > 0:
                if 'generated_text' in response[0]:
                    generated_text = response[0]['generated_text']
                elif 'text' in response[0]:
                    generated_text = response[0]['text']
                else:
                    # Try to get the first available text
                    generated_text = str(response[0].get('generated_response', response[0].get('response', '')))
            else:
                generated_text = str(response)
            
            # Clean and enhance the response
            enhanced_response = self.enhance_response(generated_text, user_input)
            
            # Add to conversation history
            conversation_history.append({"role": "assistant", "content": enhanced_response})
            
            return enhanced_response
            
        except Exception as e:
            print(f"Error generating response: {e}")
            return self.get_fallback_response(user_input)
    
    def enhance_response(self, response: str, user_input: str) -> str:
        """Enhance the generated response based on context and personality."""
        # Clean up the response
        response = self.clean_response(response)
        
        # Add personality traits
        response = self.add_personality(response)
        
        # Make it more conversational
        response = self.make_conversational(response, user_input)
        
        return response
    
    def clean_response(self, response: str) -> str:
        """Clean up the generated response text."""
        # Remove special tokens and cleanup
        response = response.strip()
        response = re.sub(r'\s+', ' ', response)
        response = re.sub(r'[""\'\']', '', response)
        
        # Capitalize first letter and add period if missing
        if response and response[0].islower():
            response = response[0].upper() + response[1:]
            
        if response and response[-1] not in ['.', '!', '?']:
            response += '.'
            
        return response
    
    def add_personality(self, response: str) -> str:
        """Add personality traits to the response."""
        # Add friendliness
        if self.personality["friendliness"] > 0.7:
            friendly_phrases = ["by the way", "I think", "in my opinion", 
                              "that's interesting", "I'd say"]
            if random.random() < 0.3:  # 30% chance to add friendly phrase
                phrase = random.choice(friendly_phrases)
                response = f"{phrase}, {response}"
        
        # Add humor if appropriate
        if self.personality["humor"] > 0.5 and self.user_emotion in ["happy", "neutral"]:
            if random.random() < 0.2:  # 20% chance to add humor
                humor_tags = ["😄", "😊", "🤣", "😆"]
                response += " " + random.choice(humor_tags)
        
        return response
    
    def make_conversational(self, response: str, user_input: str) -> str:
        """Make the response more conversational and context-aware."""
        # Add context references
        if self.conversation_length > 1:
            context_phrases = [
                "As we were discussing",
                "Regarding what you mentioned",
                "Building on that idea",
                "That reminds me"
            ]
            if random.random() < 0.25:
                response = f"{random.choice(context_phrases)}, {response}"
        
        # Add follow-up questions
        if random.random() < 0.4:  # 40% chance to add a follow-up
            follow_ups = [
                "What do you think about that?",
                "Does that make sense?",
                "How does that sound to you?",
                "Would you like me to elaborate?"
            ]
            response += " " + random.choice(follow_ups)
        
        return response
    
    def get_fallback_response(self, user_input: str) -> str:
        """Get a fallback response when model generation fails."""
        fallback_responses = [
            "That's an interesting question! Let me think about that...",
            "I'm not sure I understand completely. Could you elaborate?",
            "That's a complex topic. What specifically would you like to know?",
            "I'd love to help with that. Can you provide more details?",
            "That's fascinating! Tell me more about what you're thinking."
        ]
        
        return random.choice(fallback_responses)
    
    def get_conversation_summary(self) -> str:
        """Get a summary of the current conversation."""
        if not self.conversation_history:
            return "No conversation history yet."
            
        summary = f"Conversation Summary:\n"
        summary += f"- Topic: {self.current_topic}\n"
        summary += f"- User Emotion: {self.user_emotion}\n"
        summary += f"- Duration: {self.conversation_length} turns\n"
        summary += f"- Main Points:\n"
        
        # Extract key points from conversation
        for i, turn in enumerate(self.conversation_history):
            role = "You" if turn["role"] == "user" else "AI"
            summary += f"  {i+1}. {role}: {turn['content'][:50]}...\n"
            
        return summary
    
    def find_similar_conversations(self, query: str, top_k: int = 3) -> List[Tuple[str, float]]:
        """Find similar conversations from history using semantic search."""
        if not self.conversation_history or not self.embedding_model:
            return []
            
        try:
            # Get embedding for the query
            query_embedding = self.embedding_model.encode([query])
            
            # Get embeddings for conversation history
            history_texts = [turn["content"] for turn in self.conversation_history if turn["role"] == "user"]
            history_embeddings = self.embedding_model.encode(history_texts)
            
            # Calculate similarities
            similarities = cosine_similarity(query_embedding, history_embeddings)[0]
            
            # Get top k similar conversations
            top_indices = np.argsort(similarities)[-top_k:][::-1]
            
            similar_conversations = []
            for idx in top_indices:
                similar_conversations.append((history_texts[idx], similarities[idx]))
                
            return similar_conversations
            
        except Exception as e:
            print(f"Error in semantic search: {e}")
            return []
    
    def reset_conversation(self):
        """Reset the conversation state."""
        self.conversation_history = []
        self.current_topic = "general"
        self.user_emotion = "neutral"
        self.conversation_length = 0
        print("Conversation reset successfully!")
    
    def get_conversation_stats(self) -> Dict:
        """Get statistics about the current conversation."""
        return {
            "length": self.conversation_length,
            "current_topic": self.current_topic,
            "user_emotion": self.user_emotion,
            "personality": self.personality,
            "model": self.model_name
        }

# Example usage and testing
if __name__ == "__main__":
    print("🤖 Advanced Conversation Model - Testing")
    print("=" * 50)
    
    # Initialize the conversation model
    conv_model = ConversationModel()
    
    # Test conversation
    print("Starting test conversation...")
    
    conversation = []
    
    # Test inputs
    test_inputs = [
        "Hello! How are you doing today?",
        "I'm really excited about the new AI technologies!",
        "What do you think about machine learning?",
        "Can you tell me more about neural networks?",
        "That was very helpful, thank you!"
    ]
    
    for user_input in test_inputs:
        print(f"\n👤 User: {user_input}")
        response = conv_model.generate_response(user_input, conversation)
        print(f"🤖 AI: {response}")
        
        # Show conversation stats
        stats = conv_model.get_conversation_stats()
        print(f"📊 Topic: {stats['current_topic']} | Emotion: {stats['user_emotion']}")
    
    # Show conversation summary
    print(f"\n{conv_model.get_conversation_summary()}")
    
    # Test semantic search
    print("\n🔍 Testing semantic search...")
    similar = conv_model.find_similar_conversations("AI technologies", top_k=2)
    print("Similar conversations found:")
    for text, score in similar:
        print(f"  - '{text[:30]}...' (similarity: {score:.3f})")
    
    print("\n✅ Conversation model testing complete!")