chat_cli.py

# MovieMindAI - CLI Chat Interface
# Interactive command-line interface for movie recommendations

import sys
import re
from typing import Optional, Tuple, Dict, List
import config

try:
    from rich.console import Console
    from rich.panel import Panel
    from rich.markdown import Markdown
    from rich.table import Table
    from rich.prompt import Prompt
    from rich import print as rprint
    RICH_AVAILABLE = True
except ImportError:
    RICH_AVAILABLE = False
    print("Note: Install 'rich' for better formatting: pip install rich")

from conversation import Conversation, ConversationManager
from rag_retriever import RAGRetriever
from response_generator import ResponseGenerator
from ollama_client import OllamaClient


class ChatCLI:
    """
    Command-line interface for MovieMindAI chat.
    """
    
    # Command patterns
    COMMAND_PATTERN = re.compile(r"^/(\w+)(?:\s+(.*))?$")
    
    # Available commands
    COMMANDS = {
        "help": "Show available commands",
        "clear": "Clear conversation history",
        "filter": "Show/set filters (usage: /filter [genre|year|type|rating|country|actor|director] [value])",
        "filters": "Show current active filters",
        "clearfilters": "Clear all filters",
        "info": "Get info about a title (usage: /info <title>)",
        "compare": "Compare titles (usage: /compare <title1> vs <title2>)",
        "similar": "Find similar titles (usage: /similar <title>)",
        "exit": "Exit the chat",
        "quit": "Exit the chat",
    }
    
    def __init__(self):
        """Initialize the chat interface."""
        self.console = Console() if RICH_AVAILABLE else None
        self.conversation = Conversation()
        self.retriever: Optional[RAGRetriever] = None
        self.generator: Optional[ResponseGenerator] = None
        self.running = False
    
    def print(self, text: str, style: str = None) -> None:
        """Print text with optional styling."""
        if self.console and config.USE_COLORS:
            self.console.print(text, style=style)
        else:
            print(text)
    
    def print_panel(self, content: str, title: str = None, style: str = "blue") -> None:
        """Print a styled panel."""
        if self.console and config.USE_COLORS:
            self.console.print(Panel(content, title=title, border_style=style))
        else:
            if title:
                print(f"\n=== {title} ===")
            print(content)
            print()
    
    def print_markdown(self, text: str) -> None:
        """Print markdown-formatted text."""
        if self.console and config.USE_COLORS:
            self.console.print(Markdown(text))
        else:
            print(text)
    
    def initialize(self) -> bool:
        """Initialize the retriever and generator."""
        self.print("\n🎬 Initializing MovieMindAI...\n", style="bold cyan")
        
        try:
            # Load retriever
            self.print("📚 Loading vector store and retriever...", style="dim")
            self.retriever = RAGRetriever()
            self.print("   ✓ Retriever loaded", style="green")
            
            # Initialize Ollama client
            self.print("🤖 Connecting to Ollama...", style="dim")
            ollama = OllamaClient()
            
            if not ollama.is_available():
                self.print(f"   ⚠ Ollama model '{config.OLLAMA_MODEL}' not found.", style="yellow")
                self.print(f"   Run: ollama pull {config.OLLAMA_MODEL}", style="yellow")
                return False
            
            self.print("   ✓ Ollama connected", style="green")
            
            # Initialize generator
            self.generator = ResponseGenerator(self.retriever, ollama)
            self.print("   ✓ Response generator ready\n", style="green")
            
            return True
            
        except Exception as e:
            self.print(f"\n❌ Error initializing: {e}", style="red")
            return False
    
    def show_welcome(self) -> None:
        """Show welcome message."""
        welcome = """
# 🎬 MovieMindAI

Welcome! I'm your AI-powered movie and TV show recommendation assistant.

**How to use:**
- Just type what you're looking for (e.g., "I want a thrilling sci-fi movie")
- Use filters to narrow down results
- Type `/help` for all available commands

**Example queries:**
- "Recommend me a comedy from the 2010s"
- "What are some good crime TV series?"
- "I loved Inception, what should I watch next?"
"""
        self.print_markdown(welcome)
    
    def show_help(self) -> None:
        """Show available commands."""
        if self.console and config.USE_COLORS:
            table = Table(title="Available Commands")
            table.add_column("Command", style="cyan")
            table.add_column("Description")
            
            for cmd, desc in self.COMMANDS.items():
                table.add_row(f"/{cmd}", desc)
            
            self.console.print(table)
        else:
            print("\nAvailable Commands:")
            for cmd, desc in self.COMMANDS.items():
                print(f"  /{cmd}: {desc}")
            print()
    
    def show_filters(self) -> None:
        """Show current active filters."""
        filters = self.conversation.get_filters()
        
        if not filters:
            self.print("No active filters.", style="dim")
            return
        
        if self.console and config.USE_COLORS:
            table = Table(title="Active Filters")
            table.add_column("Filter", style="cyan")
            table.add_column("Value")
            
            for key, value in filters.items():
                table.add_row(key, str(value))
            
            self.console.print(table)
        else:
            print("\nActive Filters:")
            for key, value in filters.items():
                print(f"  {key}: {value}")
            print()
    
    def set_filter(self, args: str) -> None:
        """Set a filter value."""
        if not args:
            self.show_filters()
            return
        
        parts = args.split(maxsplit=1)
        if len(parts) < 2:
            self.print("Usage: /filter <type> <value>", style="yellow")
            self.print("Types: genre, year, type, rating, country, actor, director", style="dim")
            return
        
        filter_type, value = parts
        filter_type = filter_type.lower()
        
        # Parse filter value based on type
        try:
            if filter_type == "genre":
                self.conversation.set_filter("genres", [value])
            elif filter_type == "year":
                # Support ranges like "2000-2010"
                if "-" in value:
                    min_year, max_year = map(int, value.split("-"))
                    self.conversation.set_filter("min_year", min_year)
                    self.conversation.set_filter("max_year", max_year)
                else:
                    year = int(value)
                    self.conversation.set_filter("min_year", year)
                    self.conversation.set_filter("max_year", year)
            elif filter_type == "type":
                self.conversation.set_filter("title_type", value)
            elif filter_type == "rating":
                self.conversation.set_filter("min_rating", float(value))
            elif filter_type == "country":
                self.conversation.set_filter("country", value)
            elif filter_type == "actor":
                self.conversation.set_filter("actors", [value])
            elif filter_type == "director":
                self.conversation.set_filter("directors", [value])
            else:
                self.print(f"Unknown filter type: {filter_type}", style="yellow")
                return
            
            self.print(f"Filter set: {filter_type} = {value}", style="green")
            
        except ValueError as e:
            self.print(f"Invalid value: {e}", style="red")
    
    def handle_command(self, command: str, args: str) -> bool:
        """
        Handle a command.
        
        Returns:
            True if should continue, False if should exit
        """
        command = command.lower()
        
        if command == "help":
            self.show_help()
        
        elif command == "clear":
            self.conversation.clear()
            self.print("Conversation cleared.", style="green")
        
        elif command == "filter":
            self.set_filter(args)
        
        elif command == "filters":
            self.show_filters()
        
        elif command == "clearfilters":
            self.conversation.clear_filters()
            self.print("Filters cleared.", style="green")
        
        elif command == "info":
            if not args:
                self.print("Usage: /info <title>", style="yellow")
            else:
                self.handle_info(args)
        
        elif command == "compare":
            if not args or " vs " not in args.lower():
                self.print("Usage: /compare <title1> vs <title2>", style="yellow")
            else:
                self.handle_compare(args)
        
        elif command == "similar":
            if not args:
                self.print("Usage: /similar <title>", style="yellow")
            else:
                self.handle_similar(args)
        
        elif command in ["exit", "quit"]:
            self.print("\n👋 Goodbye! Happy watching!", style="bold cyan")
            return False
        
        else:
            self.print(f"Unknown command: /{command}", style="yellow")
            self.print("Type /help for available commands.", style="dim")
        
        return True
    
    def handle_info(self, title: str) -> None:
        """Handle /info command."""
        self.print(f"\n🔍 Looking up '{title}'...\n", style="dim")
        
        response = self.generator.generate_info(
            title,
            conversation_history=self.conversation.get_history()
        )
        
        self.print_markdown(response)
    
    def handle_compare(self, args: str) -> None:
        """Handle /compare command."""
        # Split by "vs" (case-insensitive)
        parts = re.split(r"\s+vs\s+", args, flags=re.IGNORECASE)
        if len(parts) != 2:
            self.print("Usage: /compare <title1> vs <title2>", style="yellow")
            return
        
        title1, title2 = parts
        self.print(f"\n⚖️ Comparing '{title1}' and '{title2}'...\n", style="dim")
        
        response = self.generator.generate_comparison(
            title1, title2,
            conversation_history=self.conversation.get_history()
        )
        
        self.print_markdown(response)
    
    def handle_similar(self, title: str) -> None:
        """Handle /similar command."""
        self.print(f"\n🔍 Finding titles similar to '{title}'...\n", style="dim")
        
        # Find the title first
        results = self.retriever.get_by_title(title, fuzzy=True)
        
        if len(results) == 0:
            self.print(f"Couldn't find '{title}' in the database.", style="yellow")
            return
        
        match = results.iloc[0]
        self.print(f"Found: {match['primaryTitle']} ({int(match['startYear'])})\n", style="dim")
        
        # Get similar titles
        similar = self.retriever.get_similar(match["tconst"], top_k=5)
        
        if len(similar) == 0:
            self.print("No similar titles found.", style="yellow")
            return
        
        # Display similar titles
        if self.console and config.USE_COLORS:
            table = Table(title=f"Similar to {match['primaryTitle']}")
            table.add_column("Title", style="cyan")
            table.add_column("Year")
            table.add_column("Rating")
            table.add_column("Genres")
            
            for _, row in similar.iterrows():
                genres = ", ".join(row["genres"]) if isinstance(row["genres"], list) else ""
                table.add_row(
                    row["primaryTitle"],
                    str(int(row["startYear"])),
                    f"{row['averageRating']}/10",
                    genres
                )
            
            self.console.print(table)
        else:
            print(f"\nSimilar to {match['primaryTitle']}:")
            for _, row in similar.iterrows():
                print(f"  - {row['primaryTitle']} ({int(row['startYear'])}) - {row['averageRating']}/10")
    
    def handle_query(self, query: str) -> None:
        """Handle a natural language query."""
        self.print("\n💭 Thinking...\n", style="dim")
        
        # Add to conversation
        self.conversation.add_user_message(query)
        
        # Get current filters
        filters = self.conversation.get_filters()
        
        # Generate response
        response, titles = self.generator.generate_recommendation(
            query,
            filters=filters,
            conversation_history=self.conversation.get_history()
        )
        
        # Track mentioned titles
        if len(titles) > 0:
            self.conversation.add_mentioned_titles(titles["tconst"].tolist())
        
        # Add response to history
        self.conversation.add_assistant_message(response)
        
        # Display response
        self.print_markdown(response)
    
    def run(self) -> None:
        """Run the chat loop."""
        # Initialize
        if not self.initialize():
            self.print("\n❌ Failed to initialize. Please check the errors above.", style="red")
            return
        
        # Show welcome
        self.show_welcome()
        
        self.running = True
        
        while self.running:
            try:
                # Get input
                if self.console and config.USE_COLORS:
                    user_input = Prompt.ask("\n[bold cyan]You[/bold cyan]")
                else:
                    user_input = input("\nYou: ")
                
                user_input = user_input.strip()
                
                if not user_input:
                    continue
                
                # Check for command
                match = self.COMMAND_PATTERN.match(user_input)
                if match:
                    command, args = match.groups()
                    args = args or ""
                    self.running = self.handle_command(command, args)
                else:
                    # Regular query
                    self.handle_query(user_input)
                    
            except KeyboardInterrupt:
                self.print("\n\n👋 Goodbye! Happy watching!", style="bold cyan")
                break
            except Exception as e:
                self.print(f"\n❌ Error: {e}", style="red")
                if config.DEBUG_MODE:
                    import traceback
                    traceback.print_exc()


def main():
    """Entry point for CLI."""
    cli = ChatCLI()
    cli.run()


if __name__ == "__main__":
    main()