perplexity

#!/usr/bin/env python3
"""
Perplexity CLI - A beautiful CLI for Perplexity AI API
Cross-platform CLI tool for Perplexity AI API
Works on Linux, macOS, and Windows

Usage:
  perplexity "your question here"           # Single query
  perplexity --chat                         # Interactive chat mode
  perplexity --help                         # Show all options
"""

import argparse
import json
import os
import platform
import re
import sys
import tty
import termios
from datetime import datetime
from pathlib import Path
from textwrap import fill

# Windows: Enable ANSI color support for older Windows versions
if platform.system() == "Windows":
    try:
        import ctypes
        kernel32 = ctypes.windll.kernel32
        # Enable ANSI color codes for Windows 10+
        kernel32.SetConsoleMode(kernel32.GetStdHandle(-11), 7)
    except:
        pass

# Load .env file if available
try:
    from dotenv import load_dotenv
    env_path = Path.home() / ".perplexity" / ".env"
    load_dotenv(env_path)
except ImportError:
    pass

# Default API key (your personal key)
DEFAULT_API_KEY = os.environ.get("PERPLEXITY_API_KEY")

# ANSI color codes
class Colors:
    RESET = "\033[0m"
    BOLD = "\033[1m"
    DIM = "\033[2m"
    UNDERLINE = "\033[4m"

    # Foreground colors
    BLACK = "\033[30m"
    RED = "\033[31m"
    GREEN = "\033[32m"
    YELLOW = "\033[33m"
    BLUE = "\033[34m"
    MAGENTA = "\033[35m"
    CYAN = "\033[36m"
    WHITE = "\033[37m"
    BRIGHT_BLUE = "\033[94m"
    BRIGHT_GREEN = "\033[92m"
    BRIGHT_YELLOW = "\033[93m"
    BRIGHT_MAGENTA = "\033[95m"

    # Background colors
    BG_BLACK = "\033[40m"
    BG_RED = "\033[41m"
    BG_GREEN = "\033[42m"
    BG_YELLOW = "\033[43m"
    BG_BLUE = "\033[44m"
    BG_MAGENTA = "\033[45m"
    BG_CYAN = "\033[46m"
    BG_BRIGHT_BLUE = "\033[104m"


def clear_screen():
    """Clear terminal screen"""
    os.system('clear' if os.name != 'nt' else 'cls')


def get_key():
    """Get a single keypress - cross-platform"""
    if platform.system() == "Windows":
        import msvcrt
        key = msvcrt.getch()
        if key == b'\xe0':  # Special key prefix
            key = msvcrt.getch()
            if key == b'H':
                return 'UP'
            elif key == b'P':
                return 'DOWN'
            elif key == b'K':
                return 'LEFT'
            elif key == b'M':
                return 'RIGHT'
        elif key == b'\r':
            return 'ENTER'
        elif key == b'\x03':
            return 'CTRL_C'
        elif key == b'\x1b':
            return 'ESC'
        return key.decode('utf-8', errors='ignore')
    else:
        fd = sys.stdin.fileno()
        old_settings = termios.tcgetattr(fd)
        try:
            tty.setraw(sys.stdin.fileno())
            ch = sys.stdin.read(1)
            if ch == '\x1b':
                # Arrow key sequence
                ch2 = sys.stdin.read(1)
                if ch2 == '[':
                    ch3 = sys.stdin.read(1)
                    if ch3 == 'A':
                        return 'UP'
                    elif ch3 == 'B':
                        return 'DOWN'
                    elif ch3 == 'C':
                        return 'RIGHT'
                    elif ch3 == 'D':
                        return 'LEFT'
            elif ch == '\r' or ch == '\n':
                return 'ENTER'
            elif ch == '\x03':
                return 'CTRL_C'
            elif ch == '\x1b':
                return 'ESC'
            return ch
        finally:
            termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)


def show_menu(title, options, default=0):
    """Show an interactive menu with arrow navigation

    Returns: selected index or None if cancelled
    """
    selected = default

    while True:
        clear_screen()

        # Title
        print(f"\n{Colors.BOLD}{Colors.CYAN}╔{'═' * 68}╗{Colors.RESET}")
        print(f"{Colors.BOLD}{Colors.CYAN}║  {title.center(64)}  ║{Colors.RESET}")
        print(f"{Colors.BOLD}{Colors.CYAN}╚{'═' * 68}╝{Colors.RESET}\n")

        # Options
        for i, option in enumerate(options):
            prefix = " "
            suffix = ""

            if i == selected:
                prefix = f"{Colors.BG_BRIGHT_BLUE}{Colors.WHITE}► {Colors.RESET}"
                suffix = " " + Colors.RESET
            else:
                prefix = "  "

            print(f"{prefix}{option}{suffix}")

        print(f"\n  {Colors.DIM}↑↓ Navigate • ENTER Select • ESC Cancel{Colors.RESET}")

        # Get key
        key = get_key()

        if key == 'UP':
            selected = (selected - 1) % len(options)
        elif key == 'DOWN':
            selected = (selected + 1) % len(options)
        elif key == 'ENTER':
            return selected
        elif key == 'ESC' or key == 'CTRL_C':
            return None


def print_header(text, show_cli_branding=True):
    """Print a styled header with Perplexity CLI branding"""
    branding = f" {Colors.DIM}|{Colors.RESET} Perplexity CLI" if show_cli_branding else ""
    print(f"\n{Colors.BOLD}{Colors.CYAN}{'─' * 70}{Colors.RESET}")
    print(f"{Colors.BOLD}{Colors.CYAN}  {text}{branding}{Colors.RESET}")
    print(f"{Colors.BOLD}{Colors.CYAN}{'─' * 70}{Colors.RESET}\n")


def print_box(text, width=60, padding=1):
    """Print text in a styled box"""
    border = Colors.CYAN + "═" * width + Colors.RESET
    space = " " * (width - 2)

    print(f"\n{Colors.CYAN}╔{border[1:-1]}╗{Colors.RESET}")

    # Top padding
    for _ in range(padding):
        print(f"{Colors.CYAN}║{space}║{Colors.RESET}")

    # Content
    lines = text.split('\n')
    for line in lines:
        line = line.strip()
        if line:
            content = " " + line + " "
            content = content.ljust(width - 2)
            print(f"{Colors.CYAN}║{Colors.BOLD}{Colors.CYAN}{content}{Colors.CYAN}║{Colors.RESET}")

    # Bottom padding
    for _ in range(padding):
        print(f"{Colors.CYAN}║{space}║{Colors.RESET}")

    print(f"{Colors.CYAN}╚{border[1:-1]}╝{Colors.RESET}\n")


def print_welcome(model="sonar-pro"):
    """Print welcome screen with Perplexity CLI branding"""
    clear_screen()

    # Logo
    print(f"""
{Colors.BOLD}{Colors.CYAN}
   ██████╗ ███████╗██████╗ ██████╗ ███████╗
  ██╔════╝ ██╔════╝██╔══██╗██╔══██╗██╔════╝
  ██║  ███╗█████╗  ██████╔╝██████╔╝███████╗
  ██║   ██║██╔══╝  ██╔══██╗██╔══██╗╚════██║
  ╚██████╔╝███████╗██║  ██║██║  ██║███████║
   ╚═════╝ ╚══════╝╚═╝  ╚═╝╚═╝  ╚═╝╚══════╝
                  {Colors.BRIGHT_BLUE}CLI{Colors.YELLOW} v2.0{Colors.RESET}
{Colors.RESET}""")

    # Status bar
    model_display = model.replace('-', ' ').upper()
    print(f"{Colors.BG_BRIGHT_BLUE}{Colors.WHITE}  ● Model: {Colors.BOLD}{model_display}{Colors.RESET}", end="")

    # Calculate padding to push date to right
    date_str = f"  {datetime.now().strftime('%Y-%m-%d %I:%M %p')}  "
    total_width = 70
    current_width = len(f"  ● Model: {model_display}")
    padding = total_width - current_width - len(date_str)

    print(" " * padding + f"{date_str}{Colors.RESET}\n")

    # Quick help
    print(f"{Colors.BG_BLACK}{Colors.GREEN}  MENU OPTIONS:{Colors.RESET}")
    print(f"{Colors.BG_BLACK}{Colors.WHITE}  Type '{Colors.YELLOW}/menu{Colors.WHITE}' to open settings{Colors.RESET}")
    print(f"{Colors.BG_BLACK}{Colors.WHITE}  Type '{Colors.YELLOW}/help{Colors.WHITE}' for all commands{Colors.RESET}{Colors.RESET}\n")

    # Help box
    help_text = f"""{Colors.GREEN}Quick Commands:{Colors.RESET}
  {Colors.YELLOW}/menu{Colors.RESET}        Open settings menu (arrow navigation)
  {Colors.YELLOW}/model{Colors.RESET}       Switch AI model
  {Colors.YELLOW}/clear{Colors.RESET}       Clear conversation
  {Colors.YELLOW}/save{Colors.RESET}        Save conversation
  {Colors.YELLOW}/export{Colors.RESET}      Export to Markdown
  {Colors.YELLOW}/history{Colors.RESET}     List saved chats
  {Colors.YELLOW}/load{Colors.RESET}        Load saved chat
  {Colors.YELLOW}/quit{Colors.RESET}        Exit

{Colors.DIM}Type your message and press Enter to chat.{Colors.RESET}
"""
    print(help_text)


def show_settings_menu():
    """Show interactive settings menu"""
    options = [
        "🤖 Change Model",
        "🔧 API Settings",
        "💾 Conversation History",
        "📤 Export Conversation",
        "ℹ️  About",
        "← Back to Chat"
    ]

    choice = show_menu("PERPLEXITY CLI • SETTINGS", options, default=0)

    settings_map = {
        0: 'model',
        1: 'api',
        2: 'history',
        3: 'export',
        4: 'about',
        5: 'back'
    }

    return settings_map.get(choice, 'back')


def show_model_menu(current_model):
    """Show model selection menu"""
    models = [
        ("Sonar Pro (Best)", "sonar-pro", "Recommended for complex queries"),
        ("Sonar Medium", "sonar-medium-online", "Faster responses"),
        ("Sonar Small", "sonar-small-online", "Fastest responses")
    ]

    options = []
    for name, model_id, desc in models:
        marker = f"{Colors.GREEN}●{Colors.RESET}" if model_id == current_model else "○"
        options.append(f"{marker} {name:20} {Colors.DIM}—{desc}")

    options.append("← Cancel")

    choice = show_menu("PERPLEXITY CLI • SELECT MODEL", options)

    if choice is not None and choice < len(models):
        return models[choice][1]
    return None


def show_about():
    """Show about information"""
    clear_screen()

    print(f"""
{Colors.BOLD}{Colors.CYAN}
╔════════════════════════════════════════════════════════════════════════╗
║                        {Colors.BRIGHT_CYAN}PERPLEXITY CLI{Colors.CYAN}                          ║
╠════════════════════════════════════════════════════════════════════════╣
║  {Colors.WHITE}A beautiful CLI for the Perplexity AI API{Colors.CYAN}                        ║
╚════════════════════════════════════════════════════════════════════════╝
{Colors.RESET}

{Colors.BOLD}{Colors.GREEN}Features:{Colors.RESET}
  • Interactive chat with conversation history
  • Real-time streaming responses
  • Markdown formatting with syntax highlighting
  • Save/load/export conversations
  • Research mode with citations
  • Custom AI personas
  • Beautiful terminal UI with arrow navigation

{Colors.BOLD}{Colors.GREEN}Version:{Colors.RESET} 2.0
{Colors.BOLD}{Colors.GREEN}Repository:{Colors.RESET} https://github.com/Dodothereal/perplexity-cli
{Colors.BOLD}{Colors.GREEN}Author:{Colors.RESET} Dodothereal

{Colors.BOLD}{Colors.YELLOW}Press ENTER to continue...{Colors.RESET}
""")

    # Wait for ENTER
    if platform.system() == "Windows":
        import msvcrt
        while True:
            key = msvcrt.getch()
            if key == b'\r':
                break
    else:
        input()


def print_status(client, session):
    """Print current session status"""
    msgs = len(session.messages) - 1  # Minus system prompt
    print(f"\n{Colors.BOLD}{Colors.CYAN}┌─ Perplexity CLI • Session ─────────────────────────{Colors.RESET}")
    print(f"{Colors.BOLD}{Colors.CYAN}│{Colors.RESET}")
    print(f"{Colors.BOLD}{Colors.CYAN}│{Colors.RESET}  Model:    {Colors.YELLOW}{session.model.upper()}{Colors.RESET}")
    print(f"{Colors.BOLD}{Colors.CYAN}│{Colors.RESET}  Messages: {Colors.YELLOW}{msgs}{Colors.RESET}")
    print(f"{Colors.BOLD}{Colors.CYAN}│{Colors.RESET}  History:  {Colors.DIM}{session.history_dir}{Colors.RESET}")
    print(f"{Colors.BOLD}{Colors.CYAN}└───────────────────────────────────────────────────────{Colors.RESET}\n")


def print_help():
    """Print help information"""
    print_box("""
Interactive Commands (slash commands)

/menu         Open settings menu with arrow navigation
/model        Switch AI model (interactive)
/clear        Clear conversation history
/save         Save current conversation
/export       Export conversation to Markdown
/history      List all saved conversations
/load         Load a saved conversation
/status       Show current session info
/help         Show this help

Keyboard Shortcuts
Ctrl+C        Interrupt current response
Ctrl+D        Exit chat mode
ESC           Cancel menu navigation

Tip: Use /menu for an interactive settings experience!
    """)


def format_response(text):
    """Format response with markdown-like styling"""
    # Bold text
    text = re.sub(r'\*\*(.*?)\*\*', f'{Colors.BOLD}\\1{Colors.RESET}', text)
    # Italic text
    text = re.sub(r'\*(.*?)\*', f'{Colors.UNDERLINE}\\1{Colors.RESET}', text)
    # Code blocks
    text = re.sub(r'```(\w*)\n(.*?)```', f'\n{Colors.BG_BLACK}{Colors.GREEN}\\2{Colors.RESET}\n', text, flags=re.DOTALL)
    # Inline code
    text = re.sub(r'`(.*?)`', f'{Colors.BG_BLACK}{Colors.GREEN}\\1{Colors.RESET}', text)
    # Headers
    text = re.sub(r'^### (.*?)$', f'\n{Colors.BOLD}{Colors.BRIGHT_BLUE}\\1{Colors.RESET}\n', text, flags=re.MULTILINE)
    text = re.sub(r'^## (.*?)$', f'\n{Colors.BOLD}{Colors.BRIGHT_MAGENTA}\\1{Colors.RESET}\n', text, flags=re.MULTILINE)
    text = re.sub(r'^# (.*?)$', f'\n{Colors.BOLD}{Colors.CYAN}═ {Colors.RESET}{Colors.BOLD}{Colors.CYAN}\\1{Colors.RESET}\n', text, flags=re.MULTILINE)
    # Lists
    text = re.sub(r'^- (.*?)$', f'{Colors.YELLOW}•{Colors.RESET} \\1', text, flags=re.MULTILINE)
    text = re.sub(r'^(\d+)\. (.*?)$', f'{Colors.YELLOW}\\1.{Colors.RESET} \\2', text, flags=re.MULTILINE)
    # Links
    text = re.sub(r'\[([^\]]+)\]\(([^)]+)\)', f'{Colors.UNDERLINE}{Colors.BLUE}\\1{Colors.RESET}', text)

    return text


def print_user(text):
    """Print user message with styling"""
    print(f"\n{Colors.BOLD}{Colors.BRIGHT_GREEN}┌─ You ─────────────────────────────────────────────────{Colors.RESET}")
    print(f"{Colors.BOLD}{Colors.BRIGHT_GREEN}│{Colors.RESET} {text}")
    print(f"{Colors.BOLD}{Colors.BRIGHT_GREEN}└───────────────────────────────────────────────────────{Colors.RESET}")


def print_assistant(text, streamed=False):
    """Print assistant message with styling"""
    formatted = format_response(text) if not streamed else text

    print(f"\n{Colors.BOLD}{Colors.BRIGHT_BLUE}┌─ Perplexity ────────────────────────────────────────{Colors.RESET}")
    print(f"{Colors.BOLD}{Colors.BRIGHT_BLUE}│{Colors.RESET}")

    for line in formatted.split('\n'):
        print(f"{Colors.BOLD}{Colors.BRIGHT_BLUE}│{Colors.RESET} {line}")

    print(f"{Colors.BOLD}{Colors.BRIGHT_BLUE}└───────────────────────────────────────────────────────{Colors.RESET}")


def print_citations(citations):
    """Print citations with styling"""
    if citations:
        print(f"\n{Colors.BOLD}{Colors.YELLOW}📚 Sources{Colors.RESET}")
        for i, citation in enumerate(citations, 1):
            print(f"  {Colors.DIM}[{i}] {citation[:80]}{'...' if len(citation) > 80 else ''}{Colors.RESET}")


class ChatSession:
    """Manage chat session with conversation history"""

    def __init__(self, model="sonar-pro", system_prompt=None):
        self.model = model
        self.system_prompt = system_prompt or "You are a helpful assistant."
        self.messages = [{"role": "system", "content": self.system_prompt}]
        self.history_dir = Path.home() / ".cache" / "perplexity"
        self.history_dir.mkdir(parents=True, exist_ok=True)
        self.created_at = datetime.now()

    def add_message(self, role, content):
        self.messages.append({"role": role, "content": content})

    def get_messages(self):
        return self.messages

    def save_session(self, name=None):
        if not name:
            name = f"chat_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        filepath = self.history_dir / name

        with open(filepath, 'w') as f:
            json.dump({
                "timestamp": datetime.now().isoformat(),
                "created_at": self.created_at.isoformat(),
                "model": self.model,
                "messages": self.messages
            }, f, indent=2)

        return filepath

    def list_sessions(self):
        sessions = []
        for file in sorted(self.history_dir.glob("chat_*.json")):
            try:
                with open(file) as f:
                    data = json.load(f)
                    sessions.append({
                        "file": file,
                        "timestamp": data.get("timestamp", ""),
                        "message_count": len(data.get("messages", [])) - 1
                    })
            except:
                pass
        return sessions

    def load_session(self, filepath):
        with open(filepath) as f:
            data = json.load(f)
        self.messages = data.get("messages", [])
        self.model = data.get("model", self.model)
        return len(self.messages) - 1

    def export_markdown(self, output_file):
        """Export conversation to markdown file"""
        with open(output_file, 'w') as f:
            f.write(f"# Perplexity CLI Chat Export\n\n")
            f.write(f"**Date:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"**Model:** {self.model}\n\n")
            f.write("---\n\n")

            for msg in self.messages:
                if msg["role"] == "system":
                    continue
                elif msg["role"] == "user":
                    f.write(f"## You\n\n{msg['content']}\n\n")
                elif msg["role"] == "assistant":
                    f.write(f"## Perplexity\n\n{msg['content']}\n\n")


def query_perplexity(client, prompt, session, stream=False):
    """Send query to Perplexity API"""
    session.add_message("user", prompt)

    if stream:
        full_response = ""

        print(f"\n{Colors.BOLD}{Colors.BRIGHT_BLUE}┌─ Perplexity ────────────────────────────────────────{Colors.RESET}")
        print(f"{Colors.BOLD}{Colors.BRIGHT_BLUE}│{Colors.RESET} ", end="", flush=True)

        response = client.chat.completions.create(
            model=session.model,
            messages=session.get_messages(),
            stream=True
        )

        col = 0
        for chunk in response:
            if chunk.choices and chunk.choices[0].delta.content:
                content = chunk.choices[0].delta.content
                full_response += content

                for char in content:
                    print(char, end="", flush=True)
                    col += 1
                    if col >= 68:
                        print(f"\n{Colors.BOLD}{Colors.BRIGHT_BLUE}│{Colors.RESET} ", end="", flush=True)
                        col = 0

        print(f"{Colors.BOLD}{Colors.BRIGHT_BLUE}\n└───────────────────────────────────────────────────────{Colors.RESET}")
        session.add_message("assistant", full_response)
        return full_response, None
    else:
        response = client.chat.completions.create(
            model=session.model,
            messages=session.get_messages()
        )

        answer = response.choices[0].message.content
        citations = []

        if hasattr(response.choices[0].message, "citations") and response.choices[0].message.citations:
            citations = response.choices[0].message.citations

        session.add_message("assistant", answer)
        return answer, citations


def interactive_chat(client, session, args):
    """Interactive chat mode (REPL)"""
    print_welcome(args.model)
    conversation_count = 0

    while True:
        try:
            # Get user input
            user_input = input(f"{Colors.BOLD}{Colors.GREEN}You:{Colors.RESET} ").strip()

            # Handle slash commands
            if user_input.startswith('/'):
                cmd = user_input.lower()

                if cmd in ('/quit', '/exit', '/q'):
                    if conversation_count > 0:
                        save = input(f"\n{Colors.YELLOW}Save this conversation? (y/n): {Colors.RESET}")
                        if save.lower() == 'y':
                            filepath = session.save_session()
                            print(f"{Colors.GREEN}✓ Saved to: {filepath}{Colors.RESET}")
                    print(f"\n{Colors.BOLD}{Colors.CYAN}Goodbye! 👋{Colors.RESET}\n")
                    break

                elif cmd == '/clear':
                    session.messages = [{"role": "system", "content": session.system_prompt}]
                    conversation_count = 0
                    print(f"{Colors.YELLOW}✓ Conversation cleared.{Colors.RESET}\n")
                    continue

                elif cmd == '/save':
                    filepath = session.save_session()
                    print(f"{Colors.GREEN}✓ Saved to: {filepath}{Colors.RESET}\n")
                    continue

                elif cmd == '/export':
                    output_file = input(f"{Colors.YELLOW}Output filename (e.g. chat.md): {Colors.RESET}")
                    if not output_file.endswith('.md'):
                        output_file += '.md'
                    output_path = Path(output_file).expanduser()
                    session.export_markdown(output_path)
                    print(f"{Colors.GREEN}✓ Exported to: {output_path}{Colors.RESET}\n")
                    continue

                elif cmd == '/history':
                    sessions = session.list_sessions()
                    if sessions:
                        print(f"\n{Colors.BOLD}📂 Saved Conversations{Colors.RESET}")
                        for s in sessions:
                            print(f"  • {Colors.CYAN}{s['file'].name}{Colors.RESET} ({s['message_count']} msgs) - {s['timestamp']}")
                    else:
                        print(f"{Colors.DIM}No saved conversations.{Colors.RESET}")
                    print()
                    continue

                elif cmd == '/load':
                    sessions = session.list_sessions()
                    if sessions:
                        print(f"\n{Colors.BOLD}📂 Saved Conversations{Colors.RESET}")
                        for i, s in enumerate(sessions, 1):
                            print(f"  [{i}] {Colors.CYAN}{s['file'].name}{Colors.RESET} ({s['message_count']} msgs)")
                        choice = input(f"\n{Colors.YELLOW}Load which one? (number or 'cancel'): {Colors.RESET}")
                        if choice.lower() != 'cancel':
                            try:
                                idx = int(choice) - 1
                                if 0 <= idx < len(sessions):
                                    count = session.load_session(sessions[idx]['file'])
                                    print(f"{Colors.GREEN}✓ Loaded {count} messages.{Colors.RESET}\n")
                                    continue
                            except ValueError:
                                pass
                    else:
                        print(f"{Colors.DIM}No saved conversations.{Colors.RESET}\n")
                    continue

                elif cmd == '/status':
                    print_status(client, session)
                    continue

                elif cmd == '/menu':
                    # Show settings menu
                    while True:
                        action = show_settings_menu()

                        if action == 'model':
                            new_model = show_model_menu(session.model)
                            if new_model:
                                session.model = new_model
                                print(f"{Colors.GREEN}✓ Model changed to {new_model}{Colors.RESET}\n")
                                print_status(client, session)
                                input(f"{Colors.DIM}Press ENTER to continue...{Colors.RESET}")
                                print_welcome(session.model)
                                break
                        elif action == 'api':
                            print(f"\n{Colors.DIM}API key is configured via PERPLEXITY_API_KEY or ~/.perplexity/.env{Colors.RESET}\n")
                            input(f"{Colors.DIM}Press ENTER to continue...{Colors.RESET}")
                            print_welcome(session.model)
                        elif action == 'history':
                            sessions = session.list_sessions()
                            print(f"\n{Colors.BOLD}📂 {len(sessions)} Saved Conversations{Colors.RESET}\n")
                            for s in sessions[-10:]:  # Last 10
                                print(f"  • {Colors.CYAN}{s['file'].name}{Colors.RESET} ({s['message_count']} msgs)")
                            if len(sessions) > 10:
                                print(f"  {Colors.DIM}... and {len(sessions) - 10} more{Colors.RESET}")
                            input(f"\n{Colors.DIM}Press ENTER to continue...{Colors.RESET}")
                            print_welcome(session.model)
                        elif action == 'export':
                            output_file = input(f"{Colors.YELLOW}Output filename (e.g. chat.md): {Colors.RESET}")
                            if not output_file.endswith('.md'):
                                output_file += '.md'
                            output_path = Path(output_file).expanduser()
                            session.export_markdown(output_path)
                            print(f"{Colors.GREEN}✓ Exported to: {output_path}{Colors.RESET}\n")
                            input(f"{Colors.DIM}Press ENTER to continue...{Colors.RESET}")
                            print_welcome(session.model)
                        elif action == 'about':
                            show_about()
                            print_welcome(session.model)
                        elif action == 'back':
                            print_welcome(session.model)
                            break
                    continue

                elif cmd == '/model':
                    new_model = show_model_menu(session.model)
                    if new_model:
                        session.model = new_model
                        print(f"{Colors.GREEN}✓ Model changed to {new_model}{Colors.RESET}\n")
                    continue

                elif cmd == '/help':
                    print_help()
                    continue

                elif cmd in ('/?', '--help'):
                    print_help()
                    continue

                else:
                    print(f"{Colors.RED}Unknown command: {user_input}{Colors.RESET}")
                    print(f"{Colors.DIM}Type '/help' or '/menu' for options.{Colors.RESET}\n")
                    continue

            elif not user_input:
                continue

            # Process the query
            try:
                answer, citations = query_perplexity(
                    client, user_input, session, stream=args.stream
                )

                if not args.stream:
                    print_assistant(answer)
                    if citations:
                        print_citations(citations)

                print()
                conversation_count += 1

            except KeyboardInterrupt:
                print(f"\n\n{Colors.YELLOW}⚠ Interrupted. Type '/quit' to exit.{Colors.RESET}\n")
            except Exception as e:
                print(f"\n{Colors.RED}✗ Error: {e}{Colors.RESET}\n")

        except KeyboardInterrupt:
            print(f"\n\n{Colors.YELLOW}⚠ Use '/quit' to exit.{Colors.RESET}\n")
        except EOFError:
            print(f"\n\n{Colors.BOLD}{Colors.CYAN}Goodbye! 👋{Colors.RESET}\n")
            break


def main():
    parser = argparse.ArgumentParser(
        description="Perplexity CLI - A beautiful CLI for Perplexity AI API",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=f"""{Colors.BOLD}Examples:{Colors.RESET}
  {Colors.GREEN}perplexity{Colors.RESET} "What is the capital of France?"
  {Colors.GREEN}perplexity{Colors.RESET} "Latest AI developments" --search
  {Colors.GREEN}perplexity{Colors.RESET} --chat
  {Colors.GREEN}perplexity{Colors.RESET} --chat --stream
  {Colors.GREEN}perplexity{Colors.RESET} "Explain quantum computing" --system "You are a physics teacher"

{Colors.BOLD}Interactive Commands:{Colors.RESET}
  /menu        Open settings menu (arrow navigation)
  /model       Switch AI model
  /help        Show all commands
  /quit        Exit chat mode
"""
    )
    parser.add_argument(
        "query",
        nargs="?",
        help="Your question or search query (use --chat for interactive mode)"
    )
    parser.add_argument(
        "--model",
        default="sonar-pro",
        choices=["sonar-pro", "sonar-medium-online", "sonar-small-online"],
        help="Model to use (default: sonar-pro)"
    )
    parser.add_argument(
        "--api-key",
        default=DEFAULT_API_KEY,
        help="Perplexity API key (default: uses PERPLEXITY_API_KEY env var)"
    )
    parser.add_argument(
        "--search",
        action="store_true",
        help="Use research-focused mode"
    )
    parser.add_argument(
        "--stream",
        action="store_true",
        help="Stream the response as it comes in"
    )
    parser.add_argument(
        "--chat",
        "-c",
        action="store_true",
        help="Enter interactive chat mode"
    )
    parser.add_argument(
        "--system",
        "-s",
        help="Custom system prompt"
    )
    parser.add_argument(
        "--no-color",
        action="store_true",
        help="Disable colored output"
    )

    args = parser.parse_args()

    # Disable colors if requested or not a TTY
    if args.no_color or not sys.stdout.isatty():
        for attr in dir(Colors):
            if attr.isupper():
                setattr(Colors, attr, "")

    try:
        from openai import OpenAI
    except ImportError:
        print(f"{Colors.RED}Error: OpenAI library not installed.{Colors.RESET}")
        print("Install it with: pip install openai")
        sys.exit(1)

    try:
        client = OpenAI(api_key=args.api_key, base_url="https://api.perplexity.ai")

        # Determine system prompt
        if args.system:
            system_prompt = args.system
        elif args.search:
            system_prompt = "You are a research assistant. Provide detailed, well-sourced answers with citations."
        else:
            system_prompt = "You are a helpful assistant."

        # Create session
        session = ChatSession(model=args.model, system_prompt=system_prompt)

        # Interactive chat mode
        if args.chat or not args.query:
            interactive_chat(client, session, args)
        else:
            # Single query mode
            print_header("Perplexity CLI • Query", show_cli_branding=True)

            try:
                answer, citations = query_perplexity(
                    client, args.query, session, stream=args.stream
                )

                if not args.stream:
                    print_assistant(answer)
                    if citations:
                        print_citations(citations)
                print()

            except KeyboardInterrupt:
                print(f"\n\n{Colors.YELLOW}Interrupted.{Colors.RESET}\n")
            except Exception as e:
                print(f"{Colors.RED}Error: {e}{Colors.RESET}\n", file=sys.stderr)
                sys.exit(1)

    except Exception as e:
        print(f"{Colors.RED}Error: {e}{Colors.RESET}\n", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()