main.py

"""
SQL Query Writer Agent - Main Entry Point

This file provides an interactive interface for testing your QueryWriter agent.
Your agent implementation should be in agent.py.

Usage:
    python main.py
    python main.py --ui

LLM Configuration:
    Set these environment variables to configure Ollama:
    - OLLAMA_HOST: Ollama server URL (Carleton RCS or local Ollama)
    - OLLAMA_MODEL: Optional explicit model override
    - OLLAMA_API_KEY: Carleton RCS API key (sent as x-api-key header)

UI Configuration:
    - UI_HOST: Web UI host (default: 127.0.0.1)
    - UI_PORT: Web UI port (default: 7860)
    - UI_AUTO_OPEN: Auto-open browser (true/false)
"""

import os
import sys
import time
import argparse
import threading
import duckdb
from db.bike_store import BikeStoreDb
from agent import QueryWriter
from src.conversation import ConversationManager
from src.model_fallback import ensure_model_configured


# ============================================================================
# ANSI Color Codes (works on Windows 10+ and Unix terminals)
# ============================================================================
class Colors:
    # Enable ANSI on Windows
    if sys.platform == 'win32':
        os.system('')  # Enables ANSI escape sequences on Windows
    
    RESET = '\033[0m'
    BOLD = '\033[1m'
    DIM = '\033[2m'
    
    # Colors
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    MAGENTA = '\033[95m'
    CYAN = '\033[96m'
    WHITE = '\033[97m'
    
    # Backgrounds
    BG_BLUE = '\033[44m'
    BG_GREEN = '\033[42m'


def c(text, color):
    """Colorize text."""
    return f"{color}{text}{Colors.RESET}"


def env_bool(name: str, default: bool = False) -> bool:
    """Parse boolean environment variables."""
    raw = os.getenv(name)
    if raw is None:
        return default
    return raw.strip().lower() in ("1", "true", "yes", "on")


# ============================================================================
# ASCII Art & Branding
# ============================================================================
LOGO = f"""
{Colors.CYAN}{Colors.BOLD}
   ____                        ______                    
  / __ \\__  _____  _______  __/ ____/___  _________ ____ 
 / / / / / / / _ \\/ ___/ / / / /_  / __ \\/ ___/ __ `/ _ \\
/ /_/ / /_/ /  __/ /  / /_/ / __/ / /_/ / /  / /_/ /  __/
\\___\\_\\__,_/\\___/_/   \\__, /_/    \\____/_/   \\__, /\\___/ 
                     /____/                 /____/       
{Colors.RESET}
  {Colors.YELLOW}~ Ask questions. Get SQL. ~{Colors.RESET}
{Colors.CYAN}╔════════════════════════════════════════════════════════════╗
║  {Colors.WHITE}Transform natural language into powerful SQL queries{Colors.CYAN}      ║
║  {Colors.DIM}Powered by Local LLM + DuckDB{Colors.CYAN}                             ║
╚════════════════════════════════════════════════════════════╝{Colors.RESET}
"""

def get_goodbye():
    """Generate goodbye message with proper alignment."""
    # Box inner width = 60 chars
    line1 = "   Thanks for using QueryForge!"  # 32 chars
    line2 = "   May your queries always return results."  # 42 chars
    
    pad1 = " " * (60 - len(line1))  # 28 spaces
    pad2 = " " * (60 - len(line2))  # 18 spaces
    empty = " " * 60
    
    return f"""
{Colors.CYAN}╔════════════════════════════════════════════════════════════╗
║{empty}║
║{Colors.YELLOW}{line1}{Colors.CYAN}{pad1}║
║{Colors.DIM}{line2}{Colors.CYAN}{pad2}║
║{empty}║
╚════════════════════════════════════════════════════════════╝{Colors.RESET}
"""

GOODBYE = None  # Will be set at runtime


# ============================================================================
# Progress & Animation Utilities  
# ============================================================================
class Spinner:
    """Animated spinner for long operations."""
    
    def __init__(self, message="Processing"):
        self.message = message
        self.running = False
        self.thread = None
        self.frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
        self.current_step = ""
    
    def _spin(self):
        idx = 0
        while self.running:
            frame = self.frames[idx % len(self.frames)]
            step_text = f" {Colors.DIM}({self.current_step}){Colors.RESET}" if self.current_step else ""
            sys.stdout.write(f"\r  {Colors.CYAN}{frame}{Colors.RESET} {self.message}{step_text}   ")
            sys.stdout.flush()
            time.sleep(0.1)
            idx += 1
        sys.stdout.write("\r" + " " * 80 + "\r")
        sys.stdout.flush()
    
    def start(self):
        self.running = True
        self.thread = threading.Thread(target=self._spin)
        self.thread.start()
    
    def update_step(self, step):
        self.current_step = step
    
    def stop(self, success=True):
        self.running = False
        if self.thread:
            self.thread.join()
        icon = c("✓", Colors.GREEN) if success else c("✗", Colors.RED)
        print(f"  {icon} {self.message}")


def show_progress_steps(steps, delay=0.3):
    """Display a series of progress steps."""
    for i, step in enumerate(steps, 1):
        time.sleep(delay)
        print(f"  {Colors.GREEN}[{i}/{len(steps)}]{Colors.RESET} {step}")


def print_box(title, content, color=Colors.CYAN, footer=None):
    """Print content in a nice box that expands to fit content."""
    min_width = 40
    max_width = 100
    
    # Split content into lines and wrap long lines
    raw_lines = content.split('\n')
    wrapped_lines = []
    
    for line in raw_lines:
        if len(line) <= max_width - 4:
            wrapped_lines.append(line)
        else:
            # Wrap long lines
            while len(line) > max_width - 4:
                # Try to break at a space
                break_point = line[:max_width - 4].rfind(' ')
                if break_point == -1:
                    break_point = max_width - 4
                wrapped_lines.append(line[:break_point])
                line = line[break_point:].lstrip()
            if line:
                wrapped_lines.append(line)
    
    # Calculate width based on content
    max_line_len = max(len(line) for line in wrapped_lines) if wrapped_lines else 0
    inner_width = max(min_width, min(max_width, max_line_len + 4))
    
    # Ensure title and footer fit
    inner_width = max(inner_width, len(title) + 4)
    if footer:
        inner_width = max(inner_width, len(footer) + 4)

    print(f"\n{color}┌{'─' * inner_width}┐{Colors.RESET}")

    # Title line
    title_padding = inner_width - len(title) - 2
    print(f"{color}│{Colors.RESET} {Colors.BOLD}{title}{Colors.RESET}{' ' * title_padding} {color}│{Colors.RESET}")

    print(f"{color}├{'─' * inner_width}┤{Colors.RESET}")

    # Content lines
    for line in wrapped_lines:
        content_padding = inner_width - len(line) - 2
        print(f"{color}│{Colors.RESET} {line}{' ' * content_padding} {color}│{Colors.RESET}")

    # Footer (e.g., confidence score)
    if footer:
        print(f"{color}├{'─' * inner_width}┤{Colors.RESET}")
        footer_padding = inner_width - len(footer) - 2
        print(f"{color}│{Colors.RESET} {footer}{' ' * footer_padding} {color}│{Colors.RESET}")

    print(f"{color}└{'─' * inner_width}┘{Colors.RESET}")


def print_results_table(results, max_rows=10):
    """Print results in a formatted table."""
    if not results:
        print(f"  {Colors.YELLOW}No results found.{Colors.RESET}")
        return
    
    # Calculate column widths
    num_cols = len(results[0])
    col_widths = [0] * num_cols
    
    display_results = results[:max_rows]
    for row in display_results:
        for i, val in enumerate(row):
            col_widths[i] = max(col_widths[i], len(str(val)[:30]))
    
    # Print header separator
    separator = f"  {Colors.DIM}+" + "+".join("-" * (w + 2) for w in col_widths) + f"+{Colors.RESET}"
    
    print(separator)
    for row in display_results:
        row_str = f"  {Colors.DIM}|{Colors.RESET}"
        for i, val in enumerate(row):
            val_str = str(val)[:30]
            row_str += f" {val_str:<{col_widths[i]}} {Colors.DIM}|{Colors.RESET}"
        print(row_str)
    print(separator)
    
    if len(results) > max_rows:
        print(f"  {Colors.DIM}... and {len(results) - max_rows} more rows{Colors.RESET}")


# ============================================================================
# Database & Query Functions
# ============================================================================
def initialize_database(db_path: str = 'bike_store.db'):
    """Initialize the bike store database with fancy output."""
    spinner = Spinner("Loading bike store database")
    spinner.start()

    needs_init = False
    if not os.path.exists(db_path):
        needs_init = True
    else:
        try:
            con = duckdb.connect(database=db_path, read_only=True)
            tables = con.execute("SHOW TABLES").fetchall()
            con.close()
            needs_init = len(tables) == 0
        except Exception:
            needs_init = True

    if needs_init:
        spinner.update_step("downloading data")
        db = BikeStoreDb(db_path=db_path)
    else:
        spinner.update_step("using cached database")
        db = db_path
    
    spinner.stop(success=True)
    return db


def execute_query(sql: str, db_path: str = 'bike_store.db'):
    """Execute a SQL query against the DuckDB database."""
    con = duckdb.connect(database=db_path, read_only=True)
    try:
        result = con.execute(sql).fetchall()
        return result
    finally:
        con.close()


class ProgressBar:
    """Animated progress bar with step descriptions."""
    
    def __init__(self, steps, width=40):
        self.steps = steps
        self.width = width
        self.current_step = 0
        self.running = False
        self.thread = None
        self.complete = False
    
    def _render(self):
        """Render the progress bar."""
        while self.running and self.current_step < len(self.steps):
            step_name = self.steps[self.current_step]
            progress = (self.current_step + 1) / len(self.steps)
            filled = int(self.width * progress)
            
            # Progress bar characters
            bar = f"{Colors.GREEN}{'█' * filled}{Colors.DIM}{'░' * (self.width - filled)}{Colors.RESET}"
            percent = int(progress * 100)
            
            # Clear line and print
            sys.stdout.write(f"\r  {bar} {percent:3d}%  {Colors.CYAN}{step_name}{Colors.RESET}    ")
            sys.stdout.flush()
            
            self.current_step += 1
            if self.current_step < len(self.steps):
                time.sleep(0.4)
        
        # Keep showing animation while waiting for completion
        frames = ["◐", "◓", "◑", "◒"]
        idx = 0
        while self.running and not self.complete:
            frame = frames[idx % len(frames)]
            bar = f"{Colors.GREEN}{'█' * self.width}{Colors.RESET}"
            sys.stdout.write(f"\r  {bar} 100%  {Colors.YELLOW}{frame} Finalizing...{Colors.RESET}    ")
            sys.stdout.flush()
            time.sleep(0.15)
            idx += 1
    
    def start(self):
        """Start the progress bar animation."""
        self.running = True
        self.thread = threading.Thread(target=self._render)
        self.thread.start()
    
    def finish(self, success=True, cancelled=False):
        """Stop the progress bar and show completion."""
        self.complete = True
        self.running = False
        if self.thread:
            self.thread.join()
        
        # Final state
        if cancelled:
            bar = f"{Colors.YELLOW}{'█' * (self.current_step * self.width // len(self.steps))}{Colors.DIM}{'░' * (self.width - self.current_step * self.width // len(self.steps))}{Colors.RESET}"
            sys.stdout.write(f"\r  {bar}       {Colors.YELLOW}⊘ Cancelled{Colors.RESET}                    \n")
        else:
            bar = f"{Colors.GREEN}{'█' * self.width}{Colors.RESET}"
            icon = f"{Colors.GREEN}✓{Colors.RESET}" if success else f"{Colors.RED}✗{Colors.RESET}"
            status = "Complete!" if success else "Failed"
            sys.stdout.write(f"\r  {bar} 100%  {icon} {status}                    \n")
        sys.stdout.flush()


def generate_query_with_progress(agent, user_query, show_insights=False):
    """Generate SQL query with visual progress bar and optional insights."""
    
    # Show query insights if enabled
    if show_insights and hasattr(agent, 'schema_ctx') and hasattr(agent, 'prompt_builder'):
        query_type = agent.prompt_builder.classify_query(user_query)
        tables = agent.schema_ctx.get_relevant_tables(user_query)
        complexity = agent._classify_complexity(user_query, query_type)
        
        print(f"\n  {Colors.BOLD}Query Analysis:{Colors.RESET}")
        print(f"  {Colors.DIM}├{Colors.RESET} Type: {Colors.CYAN}{query_type}{Colors.RESET}")
        print(f"  {Colors.DIM}├{Colors.RESET} Tables: {Colors.CYAN}{' → '.join(tables[:4])}{Colors.RESET}")
        print(f"  {Colors.DIM}└{Colors.RESET} Strategy: {Colors.CYAN}{complexity} path{Colors.RESET}")
        print()
    
    print(f"  {Colors.DIM}Press Ctrl+C to cancel{Colors.RESET}")
    
    steps = [
        "Analyzing question",
        "Identifying tables", 
        "Building query",
        "Validating SQL",
        "Running checks",
    ]
    
    progress = ProgressBar(steps)
    progress.start()
    
    try:
        result = agent.generate_query_with_metadata(user_query)
        progress.finish(success=True)
        return result
    except KeyboardInterrupt:
        progress.finish(cancelled=True)
        return None
    except Exception as e:
        progress.finish(success=False)
        raise e


# ============================================================================
# Help & Commands
# ============================================================================
HELP_TEXT = f"""
{Colors.BOLD}Available Commands:{Colors.RESET}
  {Colors.YELLOW}help{Colors.RESET}      - Show this help message
  {Colors.YELLOW}tables{Colors.RESET}    - List all available tables
  {Colors.YELLOW}schema{Colors.RESET}    - Show database schema
  {Colors.YELLOW}examples{Colors.RESET}  - Show example queries
  {Colors.YELLOW}context{Colors.RESET}   - Show conversation memory (if enabled)
  {Colors.YELLOW}reset context{Colors.RESET} - Clear conversation memory
  {Colors.YELLOW}clear{Colors.RESET}     - Clear the screen
  {Colors.YELLOW}quit{Colors.RESET}      - Exit QueryForge

{Colors.BOLD}Keyboard Shortcuts:{Colors.RESET}
  {Colors.YELLOW}Ctrl+C{Colors.RESET}    - Cancel current query generation/execution

{Colors.BOLD}Tips:{Colors.RESET}
  - Ask questions in natural language
  - Be specific about what you want
  - Use table/column names for precision
"""

EXAMPLES = f"""
{Colors.BOLD}Example Questions:{Colors.RESET}
  {Colors.DIM}>{Colors.RESET} How many customers are there?
  {Colors.DIM}>{Colors.RESET} Which customer has the longest name?
  {Colors.DIM}>{Colors.RESET} What did he buy?   {Colors.DIM}(conversation mode){Colors.RESET}
  {Colors.DIM}>{Colors.RESET} What are the top 5 most expensive products?
  {Colors.DIM}>{Colors.RESET} Show total revenue by store
  {Colors.DIM}>{Colors.RESET} Find customers from New York
  {Colors.DIM}>{Colors.RESET} Which staff member processed the most orders?
  {Colors.DIM}>{Colors.RESET} Monthly revenue trend for 2018
  {Colors.DIM}>{Colors.RESET} Products that have never been ordered
"""


def show_tables(agent):
    """Display available tables."""
    print(f"\n{Colors.BOLD}Available Tables:{Colors.RESET}")
    for table_name in agent.schema.keys():
        print(f"  {Colors.CYAN}◆{Colors.RESET} {table_name}")


def show_schema(agent):
    """Display database schema."""
    print(f"\n{Colors.BOLD}Database Schema:{Colors.RESET}")
    for table_name, columns in agent.schema.items():
        print(f"\n  {Colors.CYAN}┌─ {Colors.BOLD}{table_name}{Colors.RESET}")
        for col in columns:
            print(f"  {Colors.CYAN}│{Colors.RESET}   {Colors.DIM}•{Colors.RESET} {col}")


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


def is_valid_query(text: str) -> bool:
    """Check if input looks like a database question vs greeting/nonsense."""
    text_lower = text.lower().strip()
    
    # Common greetings and non-queries
    non_queries = [
        'hello', 'hi', 'hey', 'hola', 'sup', 'yo', 'greetings',
        'good morning', 'good afternoon', 'good evening', 'good night',
        'thanks', 'thank you', 'thx', 'bye', 'goodbye', 'see you',
        'ok', 'okay', 'cool', 'nice', 'great', 'awesome', 'yes', 'no',
        'what', 'why', 'who are you', 'what are you', 'how are you',
        'test', 'testing', 'asdf', 'asd', 'foo', 'bar', 'abc', '123',
    ]
    
    if text_lower in non_queries:
        return False
    
    # Too short to be a real query (less than 3 words)
    if len(text_lower.split()) < 2:
        return False
    
    # Check for database-related keywords
    query_keywords = [
        'how many', 'count', 'total', 'sum', 'average', 'avg', 'list', 'show',
        'find', 'get', 'select', 'which', 'what is', 'what are', 'where',
        'top', 'best', 'most', 'least', 'highest', 'lowest', 'first', 'last',
        'customer', 'product', 'order', 'store', 'staff', 'brand', 'category',
        'revenue', 'sales', 'inventory', 'stock', 'price', 'quantity',
        'buy', 'bought', 'purchase', 'purchased', 'spend', 'spent', 'sold', 'shop',
        'from', 'in', 'by', 'per', 'between', 'during', 'year', 'month',
        'all', 'each', 'every', 'never', 'not', 'without',
    ]
    
    return any(kw in text_lower for kw in query_keywords)


# ============================================================================
# Main Application
# ============================================================================
def main(conversation_override: bool | None = None):
    """Main function to run BikeQL interactively."""
    db_path = 'bike_store.db'
    
    # Clear screen and show logo
    clear_screen()
    print(LOGO)
    
    # Initialize database
    initialize_database(db_path)

    # Ensure model is configured, with interactive local fallback.
    model_ok, model_msg = ensure_model_configured(interactive=True)
    if model_ok:
        print(f"  {Colors.GREEN}✓{Colors.RESET} {model_msg}")
    else:
        print(f"\n  {Colors.RED}Model setup required:{Colors.RESET}")
        for line in model_msg.splitlines():
            print(f"  {Colors.DIM}{line}{Colors.RESET}")
        return
    
    # Initialize agent
    spinner = Spinner("Initializing AI agent")
    spinner.start()
    agent = QueryWriter(db_path=db_path)
    spinner.stop(success=True)

    conversation = ConversationManager(
        enabled=conversation_override if conversation_override is not None else env_bool("ENABLE_CONVERSATION", False),
        max_turns=int(os.getenv("CONVERSATION_MAX_TURNS", "8")),
        use_llm_rewrite=env_bool("CONVERSATION_LLM_REWRITE", True),
        preview_rows=int(os.getenv("CONVERSATION_PREVIEW_ROWS", "3")),
    )
    
    # Show configuration
    print(f"\n  {Colors.BOLD}Configuration:{Colors.RESET}")
    print(f"  {Colors.DIM}├{Colors.RESET} Model: {Colors.GREEN}{os.getenv('OLLAMA_MODEL', 'unset')}{Colors.RESET}")
    print(f"  {Colors.DIM}├{Colors.RESET} Host:  {Colors.GREEN}{os.getenv('OLLAMA_HOST', 'unset')}{Colors.RESET}")
    api_key_set = bool(os.getenv('OLLAMA_API_KEY') or os.getenv('RCS_API_KEY'))
    auth_color = Colors.GREEN if api_key_set else Colors.YELLOW
    auth_text = "x-api-key set" if api_key_set else "none"
    print(f"  {Colors.DIM}├{Colors.RESET} Auth:  {auth_color}{auth_text}{Colors.RESET}")
    print(f"  {Colors.DIM}├{Colors.RESET} Tables: {Colors.GREEN}{len(agent.schema)} loaded{Colors.RESET}")
    conv_status = "enabled" if conversation.enabled else "disabled"
    conv_color = Colors.GREEN if conversation.enabled else Colors.DIM
    print(f"  {Colors.DIM}└{Colors.RESET} Conversation: {conv_color}{conv_status}{Colors.RESET}")
    
    print(f"\n  {Colors.DIM}Type 'help' for commands or ask a question in natural language.{Colors.RESET}")
    print(f"  {Colors.DIM}Type 'quit' to exit.{Colors.RESET}")
    if conversation.enabled:
        print(f"  {Colors.DIM}Follow-ups are enabled (e.g., 'What did he buy?'). Use 'reset context' to clear memory.{Colors.RESET}")
    
    # Query counter
    query_count = 0
    
    # Main interaction loop
    while True:
        try:
            # Fancy prompt
            print()
            user_query = input(f"  {Colors.CYAN}QueryForge{Colors.RESET} {Colors.DIM}>{Colors.RESET} ").strip()
            
            if not user_query:
                continue
            
            # Handle commands
            cmd = user_query.lower()
            
            if cmd in ['quit', 'exit', 'q']:
                print(get_goodbye())
                break
            
            if cmd == 'help':
                print(HELP_TEXT)
                continue
            
            if cmd == 'tables':
                show_tables(agent)
                continue
            
            if cmd == 'schema':
                show_schema(agent)
                continue
            
            if cmd == 'examples':
                print(EXAMPLES)
                continue
            
            if cmd == 'clear':
                clear_screen()
                print(LOGO)
                continue

            if cmd == 'context':
                if conversation.enabled:
                    print(f"\n{Colors.BOLD}Conversation Context:{Colors.RESET}")
                    print(f"  {conversation.summarize()}")
                else:
                    print(f"\n  {Colors.DIM}Conversation mode is disabled. Set ENABLE_CONVERSATION=true in .env.{Colors.RESET}")
                continue

            if cmd in ['reset context', 'clear context', 'forget']:
                conversation.clear()
                print(f"\n  {Colors.GREEN}Conversation context cleared.{Colors.RESET}")
                continue

            effective_query = user_query
            rewrite_meta = {}
            if conversation.enabled:
                effective_query, rewrite_meta = conversation.rewrite_query(
                    user_query,
                    llm_client=agent.client,
                    model=agent.model,
                )
                if rewrite_meta.get("rewritten"):
                    print(
                        f"\n  {Colors.DIM}Resolved follow-up:{Colors.RESET} "
                        f"{Colors.CYAN}{effective_query}{Colors.RESET}"
                    )
                elif rewrite_meta.get("follow_up_detected") and rewrite_meta.get("reason") == "could_not_resolve":
                    print(
                        f"\n  {Colors.YELLOW}I detected a follow-up but could not resolve who/what it refers to.{Colors.RESET}"
                    )
                    print(
                        f"  {Colors.DIM}Try a standalone question, e.g. 'What products did Christopher Richardson buy?'{Colors.RESET}"
                    )
                    continue
            
            # Check if input looks like a valid database question
            if not is_valid_query(effective_query):
                print(f"\n  {Colors.YELLOW}I'm designed to answer questions about the database.{Colors.RESET}")
                print(f"  {Colors.DIM}Try asking something like:{Colors.RESET}")
                print(f"    {Colors.CYAN}>{Colors.RESET} How many customers are there?")
                print(f"    {Colors.CYAN}>{Colors.RESET} Show the top 5 products by price")
                print(f"    {Colors.CYAN}>{Colors.RESET} Total revenue by store")
                print(f"\n  {Colors.DIM}Type 'help' for more options or 'examples' for sample queries.{Colors.RESET}")
                continue
            
            # Generate SQL from natural language
            query_count += 1
            start_time = time.time()
            
            # Check if insights mode is enabled
            show_insights = env_bool('SHOW_INSIGHTS', False)
            result = generate_query_with_progress(agent, effective_query, show_insights=show_insights)
            
            # Check if query was cancelled
            if result is None:
                print(f"\n  {Colors.YELLOW}Query generation cancelled.{Colors.RESET}")
                continue

            # Extract SQL and metadata from result
            if hasattr(result, 'sql'):
                sql = result.sql
                confidence = result.confidence
                strategy = result.strategy
                voting = result.voting_agreement
                
                # Build confidence footer
                conf_color = Colors.GREEN if confidence == 'high' else (Colors.YELLOW if confidence == 'medium' else Colors.RED)
                if voting:
                    footer = f"Confidence: {conf_color}{confidence.upper()}{Colors.RESET} ({voting}/3 agreement) | Strategy: {strategy}"
                else:
                    footer = f"Confidence: {conf_color}{confidence.upper()}{Colors.RESET} | Strategy: {strategy}"
            else:
                # Fallback for backward compatibility
                sql = result
                footer = None

            # Display generated SQL with confidence
            print_box("Generated SQL", sql, Colors.MAGENTA, footer=footer)

            # Execute query
            print(f"  {Colors.DIM}Press Ctrl+C to cancel{Colors.RESET}")
            spinner = Spinner("Executing query on database")
            spinner.start()
            try:
                results = execute_query(sql, db_path)
                elapsed = time.time() - start_time
                spinner.stop(success=True)
                
                # Display results
                print(f"\n  {Colors.BOLD}Results{Colors.RESET} {Colors.DIM}({len(results)} rows in {elapsed:.2f}s){Colors.RESET}")
                print_results_table(results)

                if conversation.enabled:
                    query_type = getattr(result, 'query_type', "")
                    tables_used = getattr(result, 'tables_used', [])
                    conversation.record_turn(
                        original_question=user_query,
                        effective_question=effective_query,
                        sql=sql,
                        results=results,
                        query_type=query_type,
                        tables=tables_used,
                    )
            except KeyboardInterrupt:
                spinner.stop(success=False)
                print(f"\n  {Colors.YELLOW}Query execution cancelled.{Colors.RESET}")
                continue

        except KeyboardInterrupt:
            print(f"\n\n  {Colors.YELLOW}Interrupted. Type 'quit' to exit.{Colors.RESET}")
        except NotImplementedError as e:
            print(f"\n  {Colors.RED}Error:{Colors.RESET} {e}")
            print(f"  {Colors.DIM}Please implement the generate_query method in agent.py!{Colors.RESET}")
        except Exception as e:
            print(f"\n  {Colors.RED}Error:{Colors.RESET} {e}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="QueryForge CLI/UI")
    parser.add_argument("--ui", action="store_true", help="Run web UI instead of terminal chat")
    parser.add_argument("--no-browser", action="store_true", help="Do not auto-open browser in UI mode")
    parser.add_argument(
        "--conversation",
        choices=["on", "off"],
        help="Force conversation memory mode on/off for this run",
    )
    args = parser.parse_args()

    conversation_override = None
    if args.conversation is not None:
        conversation_override = args.conversation == "on"
        os.environ["ENABLE_CONVERSATION"] = "true" if conversation_override else "false"

    if args.ui:
        from src.ui_server import run_ui

        open_browser = not args.no_browser
        run_ui(open_browser=open_browser)
    else:
        main(conversation_override=conversation_override)