chat_interface.py

"""
Chat Interface for MCP-powered conversational migration
"""

import tkinter as tk
from tkinter import ttk, scrolledtext
from datetime import datetime
import threading
import json


class ChatInterface(ttk.Frame):
    """Chat interface for MCP interactions"""
    
    def __init__(self, parent, mcp_manager, theme_manager):
        super().__init__(parent)
        
        self.mcp = mcp_manager
        self.theme = theme_manager
        self.conversation_history = []
        self.migration_context = {
            "source_db": None,
            "target_db": "MySQL",
            "target_database_name": None,
            "source_connection": {},
            "schema_analyzed": False,
            "schema_data": None,
            "migration_completed": False,
            "migration_status": "Not started",
            "tables_created": [],
            "relationships_found": [],
            "ai_recommendations": [],
            "pattern_relationships": []
        }
        
        self._create_ui()
    
    def _create_ui(self):
        """Create chat UI"""
        # Header with close button
        header = ttk.Frame(self)
        header.pack(fill=tk.X, padx=10, pady=10)
        
        ttk.Label(header, text="💬 AI Assistant", 
                 font=("Arial", 14, "bold")).pack(side=tk.LEFT)
        
        self.status_label = ttk.Label(header, text="●", 
                                     foreground=self.theme.get_color("disabled"))
        self.status_label.pack(side=tk.RIGHT, padx=(5, 0))
        
        self.status_text = ttk.Label(header, text="Offline", 
                                    style="Secondary.TLabel")
        self.status_text.pack(side=tk.RIGHT, padx=(0, 5))
        
        # Separator
        ttk.Separator(self, orient=tk.HORIZONTAL).pack(fill=tk.X, padx=10)
        
        # Chat display
        chat_frame = ttk.Frame(self)
        chat_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        self.chat_display = scrolledtext.ScrolledText(
            chat_frame,
            wrap=tk.WORD,
            font=("Arial", 10),
            state=tk.DISABLED,
            relief=tk.FLAT,
            padx=10,
            pady=10
        )
        self.chat_display.pack(fill=tk.BOTH, expand=True)
        
        # Apply theme
        self.theme.apply_to_text_widget(self.chat_display)
        
        # Configure tags for styling
        self.chat_display.tag_config("user", 
                                    foreground=self.theme.get_color("accent"),
                                    font=("Arial", 10, "bold"))
        self.chat_display.tag_config("assistant", 
                                    foreground=self.theme.get_color("success"),
                                    font=("Arial", 10, "bold"))
        self.chat_display.tag_config("system", 
                                    foreground=self.theme.get_color("fg_secondary"),
                                    font=("Arial", 9, "italic"))
        self.chat_display.tag_config("error", 
                                    foreground=self.theme.get_color("error"))
        self.chat_display.tag_config("timestamp", 
                                    foreground=self.theme.get_color("fg_secondary"),
                                    font=("Arial", 8))
        
        # Input area
        input_frame = ttk.Frame(self)
        input_frame.pack(fill=tk.X, padx=10, pady=(0, 10))
        
        self.input_field = tk.Text(
            input_frame,
            height=3,
            font=("Arial", 10),
            relief=tk.FLAT,
            padx=5,
            pady=5
        )
        self.input_field.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        self.theme.apply_to_text_widget(self.input_field)
        
        # Bind Enter key
        self.input_field.bind("<Return>", self._on_enter)
        self.input_field.bind("<Shift-Return>", lambda e: None)  # Allow line breaks with Shift+Enter
        
        # Send button
        send_btn = ttk.Button(input_frame, text="Send", 
                             command=self._send_message,
                             style="Accent.TButton",
                             width=10)
        send_btn.pack(side=tk.RIGHT, padx=(10, 0))
        
        # Quick actions
        actions_frame = ttk.Frame(self)
        actions_frame.pack(fill=tk.X, padx=10, pady=(0, 10))
        
        ttk.Label(actions_frame, text="Quick Actions:", 
                 style="Secondary.TLabel").pack(side=tk.LEFT, padx=(0, 10))
        
        quick_actions = [
            ("Analyze Schema", "Analyze the current database schema with all available context"),
            ("Suggest Relationships", "Recommend relationships between tables based on analyzed schema"),
            ("Migration Status", "What is the current status of my migration?"),
            ("Clear Chat", self._clear_chat)
        ]
        
        for i, action in enumerate(quick_actions):
            if isinstance(action[1], str):
                btn = ttk.Button(actions_frame, text=action[0],
                               command=lambda msg=action[1]: self._send_quick_action(msg))
            else:
                btn = ttk.Button(actions_frame, text=action[0],
                               command=action[1])
            btn.pack(side=tk.LEFT, padx=2)
        
        # Welcome message
        self._add_system_message("Welcome! I'm your MCP-powered migration assistant.")
        self._add_system_message("Ask me anything about database migration, schema design, or use the quick actions above.")
    
    def set_mcp_status(self, is_active):
        """Update MCP status indicator"""
        if is_active:
            self.status_label.config(foreground=self.theme.get_color("success"))
            self.status_text.config(text="Online")
        else:
            self.status_label.config(foreground=self.theme.get_color("disabled"))
            self.status_text.config(text="Offline")
    
    def _on_enter(self, event):
        """Handle Enter key press"""
        if not event.state & 0x1:  # Not Shift+Enter
            self._send_message()
            return "break"  # Prevent newline
    
    def _send_message(self):
        """Send user message"""
        message = self.input_field.get("1.0", tk.END).strip()
        
        if not message:
            return
        
        # Clear input
        self.input_field.delete("1.0", tk.END)
        
        # Add to display
        self._add_user_message(message)
        
        # Process with MCP
        self._process_message(message)
    
    def _send_quick_action(self, message):
        """Send a quick action message"""
        self._add_user_message(message)
        self._process_message(message)
    
    def _process_message(self, message):
        """Process message with MCP"""
        if not self.mcp or not self.mcp.is_initialized:
            self._add_error_message("MCP is not initialized. Please enable MCP in settings and configure your API key.")
            return
        
        # Show thinking indicator
        self._add_system_message("Thinking...")
        
        # Process in background thread
        def process():
            try:
                # Build comprehensive context message
                context_parts = [
                    "=== CURRENT APPLICATION STATE ===",
                    f"Source Database: {self.migration_context.get('source_db', 'Not selected')}",
                    f"Target Database: {self.migration_context.get('target_db', 'MySQL')}",
                    f"Target DB Name: {self.migration_context.get('target_database_name', 'Not set')}",
                    f"Migration Status: {self.migration_context.get('migration_status', 'Not started')}",
                    f"Schema Analyzed: {self.migration_context.get('schema_analyzed', False)}",
                ]
                
                # Add schema data if available
                if self.migration_context.get('schema_data'):
                    context_parts.append(f"\n=== SOURCE SCHEMA DATA ===")
                    context_parts.append(json.dumps(self.migration_context['schema_data'], indent=2))
                
                # Add AI recommendations if available
                if self.migration_context.get('ai_recommendations'):
                    context_parts.append(f"\n=== AI RECOMMENDATIONS ({len(self.migration_context['ai_recommendations'])}) ===")
                    for i, rec in enumerate(self.migration_context['ai_recommendations'][:5], 1):
                        context_parts.append(f"{i}. {rec.get('from_table')}.{rec.get('from_field')} → {rec.get('to_table')}.{rec.get('to_field')} ({rec.get('type')})")
                
                # Add pattern relationships if available
                if self.migration_context.get('pattern_relationships'):
                    context_parts.append(f"\n=== PATTERN-BASED RELATIONSHIPS ({len(self.migration_context['pattern_relationships'])}) ===")
                    for i, rel in enumerate(self.migration_context['pattern_relationships'][:5], 1):
                        context_parts.append(f"{i}. {rel.get('from_table')}.{rel.get('from_field')} → {rel.get('to_table')}.{rel.get('to_field')} ({rel.get('type')})")
                
                # Add tables created during migration
                if self.migration_context.get('tables_created'):
                    context_parts.append(f"\n=== TABLES CREATED ({len(self.migration_context['tables_created'])}) ===")
                    context_parts.append(", ".join(self.migration_context['tables_created']))
                
                context_parts.append(f"\n=== USER QUERY ===")
                context_parts.append(message)
                
                context_message = "\n".join(context_parts)
                
                response = self.mcp.chat_query(context_message, self.conversation_history)
                self.after(0, lambda: self._add_assistant_message(response))
            except Exception as e:
                self.after(0, lambda: self._add_error_message(f"Error: {str(e)}"))
        
        thread = threading.Thread(target=process, daemon=True)
        thread.start()
    
    def update_migration_context(self, source_db=None, target_db_name=None, schema_analyzed=False, 
                                 schema_data=None, migration_completed=False, tables_created=None,
                                 ai_recommendations=None, pattern_relationships=None, migration_status=None):
        """Update the migration context for better AI responses"""
        if source_db:
            self.migration_context['source_db'] = source_db
            self._add_system_message(f"📊 Context updated: Now migrating from {source_db} to MySQL")
        
        if target_db_name:
            self.migration_context['target_database_name'] = target_db_name
            self._add_system_message(f"🎯 Target database: {target_db_name}")
        
        if schema_analyzed:
            self.migration_context['schema_analyzed'] = True
            self._add_system_message("✅ Schema analysis complete. I can now help with specific recommendations.")
        
        if schema_data:
            self.migration_context['schema_data'] = schema_data
            db_type = schema_data.get('database_type', 'Unknown')
            self._add_system_message(f"📥 Loaded {db_type} schema data")
        
        if ai_recommendations is not None:
            self.migration_context['ai_recommendations'] = ai_recommendations
            if len(ai_recommendations) > 0:
                self._add_system_message(f"🤖 Received {len(ai_recommendations)} AI recommendations")
        
        if pattern_relationships is not None:
            self.migration_context['pattern_relationships'] = pattern_relationships
            if len(pattern_relationships) > 0:
                self._add_system_message(f"🔗 Generated {len(pattern_relationships)} pattern-based relationships")
        
        if migration_completed:
            self.migration_context['migration_completed'] = True
            self.migration_context['migration_status'] = "Completed"
            self._add_system_message("✅ Migration completed successfully!")
        
        if tables_created:
            self.migration_context['tables_created'] = tables_created
            self._add_system_message(f"📊 Created {len(tables_created)} tables")
        
        if migration_status:
            self.migration_context['migration_status'] = migration_status
    
    def _add_user_message(self, message):
        """Add user message to chat"""
        self._add_message("You", message, "user")
        self.conversation_history.append({
            "role": "user",
            "content": message,
            "timestamp": datetime.now().isoformat()
        })
    
    def _add_assistant_message(self, message):
        """Add assistant message to chat"""
        # Remove thinking indicator
        self.chat_display.config(state=tk.NORMAL)
        content = self.chat_display.get("1.0", tk.END)
        if "Thinking..." in content:
            # Find and remove last "Thinking..." line
            lines = content.split('\n')
            for i in range(len(lines) - 1, -1, -1):
                if "Thinking..." in lines[i]:
                    del lines[i]
                    break
            self.chat_display.delete("1.0", tk.END)
            self.chat_display.insert("1.0", '\n'.join(lines))
        self.chat_display.config(state=tk.DISABLED)
        
        self._add_message("Assistant", message, "assistant")
        self.conversation_history.append({
            "role": "assistant",
            "content": message,
            "timestamp": datetime.now().isoformat()
        })
    
    def _add_system_message(self, message):
        """Add system message to chat"""
        self.chat_display.config(state=tk.NORMAL)
        
        timestamp = datetime.now().strftime("%H:%M")
        self.chat_display.insert(tk.END, f"[{timestamp}] ", "timestamp")
        self.chat_display.insert(tk.END, f"ℹ️ {message}\n\n", "system")
        
        self.chat_display.see(tk.END)
        self.chat_display.config(state=tk.DISABLED)
    
    def _add_error_message(self, message):
        """Add error message to chat"""
        self.chat_display.config(state=tk.NORMAL)
        
        timestamp = datetime.now().strftime("%H:%M")
        self.chat_display.insert(tk.END, f"[{timestamp}] ", "timestamp")
        self.chat_display.insert(tk.END, f"❌ {message}\n\n", "error")
        
        self.chat_display.see(tk.END)
        self.chat_display.config(state=tk.DISABLED)
    
    def _add_message(self, sender, message, tag):
        """Add a message to the chat display"""
        self.chat_display.config(state=tk.NORMAL)
        
        timestamp = datetime.now().strftime("%H:%M")
        self.chat_display.insert(tk.END, f"[{timestamp}] ", "timestamp")
        self.chat_display.insert(tk.END, f"{sender}:\n", tag)
        self.chat_display.insert(tk.END, f"{message}\n\n")
        
        self.chat_display.see(tk.END)
        self.chat_display.config(state=tk.DISABLED)
    
    def _clear_chat(self):
        """Clear chat history"""
        self.chat_display.config(state=tk.NORMAL)
        self.chat_display.delete("1.0", tk.END)
        self.chat_display.config(state=tk.DISABLED)
        
        self.conversation_history.clear()
        
        self._add_system_message("Chat cleared")
        self._add_system_message("How can I help you with your database migration?")
    
    def add_migration_update(self, message, status="info"):
        """Add a migration progress update"""
        if status == "success":
            self._add_system_message(f"✅ {message}")
        elif status == "error":
            self._add_error_message(message)
        elif status == "warning":
            self._add_system_message(f"⚠️ {message}")
        else:
            self._add_system_message(message)


if __name__ == "__main__":
    # Test chat interface
    from theme_manager import ThemeManager
    from config_manager import SettingsManager
    from mcp_manager import MCPManager
    
    root = tk.Tk()
    root.title("Chat Interface Test")
    root.geometry("500x600")
    
    config = SettingsManager()
    theme = ThemeManager(root, "light")
    mcp = MCPManager(config)
    
    chat = ChatInterface(root, mcp, theme)
    chat.pack(fill=tk.BOTH, expand=True)
    
    # Simulate MCP status
    chat.set_mcp_status(False)
    
    root.mainloop()