main.py

"""AI Shell Agent - Main Application"""

import sys
import os
import logging
from typing import Dict, Any
import click
from rich.console import Console
from rich.panel import Panel
from rich.syntax import Syntax
from rich.table import Table
from rich.prompt import Prompt, Confirm
from rich.text import Text
from rich.logging import RichHandler
from prompt_toolkit import prompt
from prompt_toolkit.history import InMemoryHistory
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from config.settings import settings
from graph.shell_graph import shell_agent

logging.basicConfig(
    level=logging.INFO if settings.DEBUG else logging.WARNING,
    format="%(message)s",
    datefmt="[%X]",
    handlers=[RichHandler(rich_tracebacks=True)]
)

logger = logging.getLogger(__name__)

class AIShellAgent:
    """Main application class for the AI Shell Agent."""
    
    def __init__(self):
        """Initialize the AI Shell Agent."""
        self.console = Console()
        self.history = InMemoryHistory()
        self.running = True
        
        self._check_configuration()
        self._setup_os_preference()
    
    def _check_configuration(self):
        """Check if the application is properly configured."""
        if not settings.GROQ_API_KEY or settings.GROQ_API_KEY == "gsk_your_groq_api_key_here":
            self.console.print(
                Panel.fit(
                    "[red]API key not found![/red]\n\n"
                    "Please set your Groq API key in config/settings.py\n"
                    "Get free API key at: https://console.groq.com/",
                    title="Configuration Error",
                    border_style="red"
                )
            )
            sys.exit(1)
    
    def _setup_os_preference(self):
        """Ask user for their operating system preference."""
        if not settings.USER_OS:
            self.console.print("\n[cyan]OS Detection:[/cyan]")
            self.console.print(f"Detected OS: [yellow]{settings.DETECTED_OS}[/yellow]")
            
            os_options = {
                "1": ("windows", "Windows (PowerShell/CMD)"),
                "2": ("linux", "Linux (Bash/Shell)"), 
                "3": ("macos", "macOS (Bash/Shell)"),
                "4": ("auto", f"Auto-detect ({settings.DETECTED_OS})")
            }
            
            self.console.print("\n[green]Please select your preferred operating system:[/green]")
            for key, (_, description) in os_options.items():
                self.console.print(f"  {key}. {description}")
            
            while True:
                choice = Prompt.ask("\nSelect option", choices=["1", "2", "3", "4"], default="4")
                
                if choice in os_options:
                    selected_os, description = os_options[choice]
                    if selected_os == "auto":
                        selected_os = settings.DETECTED_OS.lower()
                    
                    settings.set_user_os(selected_os)
                    self.console.print(f"\n[green]✓[/green] Using [cyan]{description}[/cyan] commands")
                    break
    
    def display_welcome(self):
        """Display welcome message and instructions."""
        welcome_text = Text()
        welcome_text.append("AI Shell Agent", style="bold cyan")
        welcome_text.append(" - Convert natural language to shell commands", style="white")
        
        instructions = """
[bold green]How to use:[/bold green]
• Type natural language descriptions of what you want to do
• The AI will generate safe shell commands
• Review and confirm before execution
• Type 'help' for more commands
• Type 'quit' or 'exit' to leave

[bold yellow]Examples:[/bold yellow]
• "find all python files in current directory"
• "show disk usage of folders here"
• "list files modified today"
• "count lines in all .txt files"

[bold red]Safety Features:[/bold red]
• Command validation and safety checks
• Whitelist of allowed commands
• User confirmation for risky operations
• Protection against dangerous patterns
        """
        
        self.console.print(Panel(welcome_text, expand=False))
        self.console.print(Panel(instructions, title="Welcome to AI Shell Agent", border_style="cyan"))
    
    def display_help(self):
        """Display help information."""
        help_text = """
[bold green]Available Commands:[/bold green]

[cyan]help[/cyan] - Show this help message
[cyan]quit, exit[/cyan] - Exit the application
[cyan]history[/cyan] - Show command history
[cyan]clear[/cyan] - Clear the screen
[cyan]status[/cyan] - Show system status
[cyan]os[/cyan] - Change operating system preference

[bold green]Natural Language Examples:[/bold green]

• "list all files in current directory"
• "find files larger than 10MB"
• "show system information"
• "display running processes"
• "check disk space"
• "search for text in files"

[bold yellow]Tips:[/bold yellow]
• Commands are generated for your selected OS
• Be specific about what you want to do
• The AI will explain commands before execution
• You can always decline to run a command
        """
        
        self.console.print(Panel(help_text, title="Help", border_style="green"))
    
    def display_status(self):
        """Display system status."""
        table = Table(title="System Status")
        table.add_column("Setting", style="cyan")
        table.add_column("Value", style="green")
        
        table.add_row("API Provider", settings.API_PROVIDER)
        table.add_row("Model", settings.GROQ_MODEL)
        table.add_row("Command Timeout", f"{settings.COMMAND_TIMEOUT}s")
        table.add_row("Debug Mode", "Enabled" if settings.DEBUG else "Disabled")
        table.add_row("Detected OS", settings.DETECTED_OS)
        table.add_row("Target OS", settings.get_effective_os().upper())
        
        self.console.print(table)
    
    def display_command_result(self, result: Dict[str, Any]):
        """Display the result of command processing."""
        
        if result.get("errors"):
            for error in result["errors"]:
                self.console.print(f"[red]{error}[/red]")
        
        if result.get("generated_command"):
            self.console.print("\n[green]AI Analysis:[/green]")
            if result.get("explanation"):
                self.console.print(f"   {result['explanation']}")
            
            command_panel = Panel(
                Syntax(result["generated_command"], "bash", theme="monokai"),
                title="Generated Command",
                border_style="yellow"
            )
            self.console.print(command_panel)
            
            safety_check = result.get("safety_check", {})
            if safety_check:
                risk_level = safety_check.get("risk_level", "unknown")
                risk_colors = {
                    "low": "green",
                    "medium": "yellow", 
                    "high": "orange",
                    "critical": "red"
                }
                risk_color = risk_colors.get(risk_level, "white")
                
                safety_text = f"Risk Level: [{risk_color}]{risk_level.upper()}[/{risk_color}]"
                if safety_check.get("recommendation"):
                    safety_text += f"\n   {safety_check['recommendation']}"
                
                self.console.print(safety_text)
        
        if result.get("warnings"):
            for warning in result["warnings"]:
                self.console.print(f"[yellow]{warning}[/yellow]")
        
        if result.get("execution_result"):
            exec_result = result["execution_result"]
            
            if exec_result.get("success"):
                self.console.print("\n[green]Command executed successfully![/green]")
                
                if exec_result.get("stdout"):
                    output_panel = Panel(
                        exec_result["stdout"],
                        title="Output",
                        border_style="green"
                    )
                    self.console.print(output_panel)
                
                if exec_result.get("execution_time"):
                    self.console.print(f"Execution time: {exec_result['execution_time']:.2f}s")
                
            else:
                self.console.print("\n[red]Command execution failed![/red]")
                if exec_result.get("stderr"):
                    error_panel = Panel(
                        exec_result["stderr"],
                        title="Error Output",
                        border_style="red"
                    )
                    self.console.print(error_panel)
        
        if result.get("alternatives"):
            self.console.print("\n[cyan]Alternative suggestions:[/cyan]")
            for i, alt in enumerate(result["alternatives"], 1):
                if isinstance(alt, dict):
                    cmd = alt.get("command", str(alt))
                    explanation = alt.get("explanation", "")
                    self.console.print(f"   {i}. [cyan]{cmd}[/cyan]")
                    if explanation:
                        self.console.print(f"      {explanation}")
                else:
                    self.console.print(f"   {i}. [cyan]{alt}[/cyan]")
    
    def process_user_input(self, user_input: str) -> bool:
        """Process user input and handle commands."""
        user_input = user_input.strip()
        
        if not user_input:
            return True
        
        if user_input.lower() in ["quit", "exit", "q"]:
            self.console.print("[cyan]Goodbye![/cyan]")
            return False
        
        elif user_input.lower() == "help":
            self.display_help()
            return True
        
        elif user_input.lower() == "clear":
            os.system('cls' if os.name == 'nt' else 'clear')
            return True
        
        elif user_input.lower() == "status":
            self.display_status()
            return True
        
        elif user_input.lower() == "history":
            self.console.print("[cyan]Command History:[/cyan]")
            for i, cmd in enumerate(self.history.get_strings()[-10:], 1):
                self.console.print(f"   {i}. {cmd}")
            return True
        
        elif user_input.lower() in ["os", "change-os", "set-os"]:
            self._setup_os_preference()
            return True
        
        return self.process_natural_language(user_input)
    
    def process_natural_language(self, user_input: str) -> bool:
        """Process natural language input through the AI workflow."""
        try:
            self.console.print(f"\n[cyan]Processing:[/cyan] {user_input}")
            
            with self.console.status("[bold green]Thinking...") as status:
                result = shell_agent.run(user_input, user_confirmed=False)
            
            self.display_command_result(result)
            
            if (result.get("require_confirmation") and 
                result.get("generated_command") and 
                not result.get("errors")):
                
                confirmed = Confirm.ask(
                    "\nDo you want to execute this command?",
                    default=False
                )
                
                if confirmed:
                    with self.console.status("[bold green]Executing...") as status:
                        result = shell_agent.run(user_input, user_confirmed=True)
                    
                    if result.get("execution_result"):
                        exec_result = result["execution_result"]
                        if exec_result.get("success"):
                            self.console.print("\n[green]Command executed successfully![/green]")
                            if exec_result.get("stdout"):
                                output_panel = Panel(
                                    exec_result["stdout"],
                                    title="Output",
                                    border_style="green"
                                )
                                self.console.print(output_panel)
                        else:
                            self.console.print("\n[red]Command execution failed![/red]")
                            if exec_result.get("stderr"):
                                error_panel = Panel(
                                    exec_result["stderr"],
                                    title="Error Output", 
                                    border_style="red"
                                )
                                self.console.print(error_panel)
                else:
                    self.console.print("[yellow]Command execution cancelled.[/yellow]")
            
        except KeyboardInterrupt:
            self.console.print("\n[yellow]Operation cancelled.[/yellow]")
        except Exception as e:
            self.console.print(f"\n[red]Error: {str(e)}[/red]")
            if settings.DEBUG:
                self.console.print_exception()
        
        return True
    
    def run_interactive(self):
        """Run the interactive shell interface."""
        self.display_welcome()
        
        while self.running:
            try:
                user_input = prompt(
                    "AI Shell Agent > ",
                    history=self.history,
                    auto_suggest=AutoSuggestFromHistory()
                )
                
                should_continue = self.process_user_input(user_input)
                
                if not should_continue:
                    break
                    
            except KeyboardInterrupt:
                self.console.print("\n\n[cyan]Goodbye![/cyan]")
                break
            except EOFError:
                self.console.print("\n\n[cyan]Goodbye![/cyan]")
                break
            except Exception as e:
                self.console.print(f"\n[red]Unexpected error: {str(e)}[/red]")
                if settings.DEBUG:
                    self.console.print_exception()

@click.command()
@click.option("--debug", is_flag=True, help="Enable debug mode")
@click.option("--command", "-c", help="Execute a single command and exit")
@click.version_option(version="1.0.0", prog_name="AI Shell Agent")
def main(debug: bool, command: str):
    """AI Shell Agent - Convert natural language to shell commands."""
    
    if debug:
        settings.DEBUG = True
        logging.getLogger().setLevel(logging.DEBUG)
    
    agent = AIShellAgent()
    
    if command:
        agent.process_user_input(command)
    else:
        agent.run_interactive()

if __name__ == "__main__":
    main()