src/cli.py

"""
Interactive CLI for querying personas about urban planning topics
"""

import sys
from typing import Optional
from rich.console import Console
from rich.prompt import Prompt, Confirm
from rich.table import Table
from rich.panel import Panel

from src.pipeline.query_engine import QueryEngine

console = Console()


class PersonaCLI:
    """Interactive command-line interface for persona queries"""

    def __init__(self):
        """Initialize CLI"""
        self.engine: Optional[QueryEngine] = None

    def initialize(self) -> bool:
        """Initialize the query engine"""
        try:
            console.print("\n[cyan]Initializing AI Personas system...[/cyan]\n")
            self.engine = QueryEngine()

            if not self.engine.test_system():
                console.print("[red]System initialization failed.[/red]")
                return False

            console.print("[green]✓ System ready![/green]\n")
            return True

        except Exception as e:
            console.print(f"[red]Error initializing system: {e}[/red]")
            return False

    def show_personas(self):
        """Display available personas"""
        personas = self.engine.list_available_personas()

        table = Table(title="Available Personas", show_header=True)
        table.add_column("ID", style="cyan")
        table.add_column("Name", style="green")
        table.add_column("Role", style="yellow")

        for persona_id, name, role in personas:
            table.add_row(persona_id, name, role)

        console.print(table)
        console.print()

    def show_contexts(self):
        """Display available contexts"""
        contexts = self.engine.list_available_contexts()

        if not contexts:
            console.print("[yellow]No environmental contexts loaded.[/yellow]\n")
            return

        console.print("[cyan]Available Contexts:[/cyan]")
        for context_id in contexts:
            console.print(f"  • {context_id}")
        console.print()

    def query_single_persona(self):
        """Interactive single persona query"""
        # Show available personas
        self.show_personas()

        # Get persona selection
        personas = self.engine.list_available_personas()
        persona_ids = [p[0] for p in personas]

        persona_id = Prompt.ask(
            "Select a persona",
            choices=persona_ids,
        )

        # Get question
        console.print()
        question = Prompt.ask("[cyan]Your question[/cyan]")

        # Optional context
        console.print()
        use_context = Confirm.ask("Use environmental context?", default=False)
        context_id = None

        if use_context:
            contexts = self.engine.list_available_contexts()
            if contexts:
                context_id = Prompt.ask(
                    "Select context",
                    choices=contexts,
                    default=contexts[0] if contexts else None,
                )

        # Query
        console.print("\n[dim]Generating response...[/dim]\n")

        try:
            response = self.engine.query(
                persona_id=persona_id,
                question=question,
                context_id=context_id,
            )

            # Display response
            console.print(Panel(
                f"[bold]{response.persona_name}[/bold] [dim]({response.persona_role})[/dim]\n\n"
                f"{response.response}",
                title="Response",
                border_style="green",
            ))
            console.print()

        except Exception as e:
            console.print(f"[red]Error: {e}[/red]\n")

    def query_multiple_personas(self):
        """Query multiple personas with the same question"""
        # Show available personas
        self.show_personas()

        # Get question
        question = Prompt.ask("[cyan]Your question (will be asked to all personas)[/cyan]")

        # Optional context
        console.print()
        use_context = Confirm.ask("Use environmental context?", default=False)
        context_id = None

        if use_context:
            contexts = self.engine.list_available_contexts()
            if contexts:
                context_id = Prompt.ask(
                    "Select context",
                    choices=contexts,
                    default=contexts[0] if contexts else None,
                )

        # Query all personas
        personas = self.engine.list_available_personas()
        persona_ids = [p[0] for p in personas]

        console.print(f"\n[dim]Querying {len(persona_ids)} personas...[/dim]\n")

        try:
            responses = self.engine.query_multiple(
                persona_ids=persona_ids,
                question=question,
                context_id=context_id,
            )

            # Display all responses
            for i, response in enumerate(responses, 1):
                console.print(f"[bold cyan]{i}. {response.persona_name}[/bold cyan] [dim]({response.persona_role})[/dim]")
                console.print("-" * 70)
                console.print(response.response)
                console.print()

            console.print(f"[green]✓ Received {len(responses)} responses[/green]\n")

        except Exception as e:
            console.print(f"[red]Error: {e}[/red]\n")

    def show_help(self):
        """Show help information"""
        help_text = """
[bold cyan]AI Personas for Urban Planning - Interactive CLI[/bold cyan]

[bold]Commands:[/bold]
  1 - Query a single persona
  2 - Query all personas with the same question
  3 - List available personas
  4 - List available contexts
  h - Show this help
  q - Quit

[bold]About:[/bold]
This system allows you to query synthetic personas representing different
urban planning stakeholders. Each persona has unique values, experiences,
and perspectives that shape their responses.
        """
        console.print(Panel(help_text.strip(), border_style="cyan"))
        console.print()

    def run(self):
        """Run the interactive CLI"""
        if not self.initialize():
            return

        self.show_help()

        while True:
            console.print("[cyan]Options:[/cyan] [1]Single [2]Multiple [3]List Personas [4]List Contexts [h]Help [q]Quit")
            choice = Prompt.ask("Select", default="1")

            if choice.lower() in ["q", "quit", "exit"]:
                console.print("\n[cyan]Goodbye![/cyan]\n")
                break

            elif choice == "1":
                console.print()
                self.query_single_persona()

            elif choice == "2":
                console.print()
                self.query_multiple_personas()

            elif choice == "3":
                console.print()
                self.show_personas()

            elif choice == "4":
                console.print()
                self.show_contexts()

            elif choice.lower() in ["h", "help"]:
                console.print()
                self.show_help()

            else:
                console.print("[yellow]Invalid choice. Press 'h' for help.[/yellow]\n")


def main():
    """Main entry point"""
    try:
        cli = PersonaCLI()
        cli.run()
    except KeyboardInterrupt:
        console.print("\n\n[cyan]Interrupted. Goodbye![/cyan]\n")
        sys.exit(0)
    except Exception as e:
        console.print(f"\n[red]Error: {e}[/red]\n")
        sys.exit(1)


if __name__ == "__main__":
    main()