wizard.py

"""
Interactive CLI Wizard
======================
Guided configuration wizard for setting up scrape jobs with validation and preview.
"""

import sys
from pathlib import Path
from typing import List, Optional, Dict, Any

import yaml
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.prompt import Prompt, Confirm, IntPrompt
from rich.text import Text
from rich import box

from db import DatabaseManager


console = Console()


class ScrapeWizard:
    """
    Interactive wizard for configuring scrape jobs.

    Features:
    - Step-by-step guided setup
    - Input validation
    - Configuration preview
    - Preset and custom city selection
    - Multi-query support
    - Export previous results
    """

    ZIP_PRESETS = {
        "1": ("top_10",    "Top 10 zip codes by population",     10),
        "2": ("top_100",   "Top 100 zip codes by population",    100),
        "3": ("top_500",   "Top 500 zip codes by population",    500),
        "4": ("top_1000",  "Top 1,000 zip codes by population",  1000),
        "5": ("top_5000",  "Top 5,000 zip codes by population",  5000),
        "6": ("top_10000", "Top 10,000 zip codes by population", 10000),
        "7": ("all",       "All ~33,000 US zip codes",           33000),
        "8": ("city",      "All zips for a specific city",       0),
        "9": ("state",     "All zips for one or more states",    0),
        "0": ("custom",    "Enter zip codes manually",           0),
    }

    def __init__(self, config: dict, db: DatabaseManager):
        self.config = config
        self.db = db
        self.result: Dict[str, Any] = {}

    def run(self) -> Optional[Dict[str, Any]]:
        """Run the wizard and return configuration dict or None if cancelled."""
        self._show_welcome()

        try:
            # Step 1: Choose action
            action = self._step_action()
            if action == "search":
                self._handle_search()
                return None
            elif action == "export":
                self._handle_export()
                return None
            elif action == "resume":
                return self._handle_resume()
            elif action == "cancel":
                return None

            # Step 2: Configure queries
            if not self._step_queries():
                return None

            # Step 3: Select cities
            if not self._step_locations():
                return None

            # Step 4: Configure options
            if not self._step_options():
                return None

            # Step 5: Preview and confirm
            if not self._step_preview():
                return None

            return self.result

        except KeyboardInterrupt:
            console.print("\n[yellow]Wizard cancelled.[/yellow]")
            return None

    def _show_welcome(self):
        """Display welcome banner."""
        welcome = Text()
        welcome.append("\n  GOOGLE MAPS SCRAPER  \n", style="bold white on blue")
        welcome.append("\n  Interactive Setup Wizard  ", style="dim")

        console.print(Panel(welcome, border_style="blue", padding=(1, 2)))
        console.print()

    def _step_action(self) -> str:
        """Step 1: Choose what to do."""
        console.print("[bold cyan]Step 1:[/bold cyan] What would you like to do?\n")

        actions = Table(show_header=False, box=box.SIMPLE, padding=(0, 2))
        actions.add_column("Key", style="bold yellow", width=4)
        actions.add_column("Action")

        actions.add_row("1", "Start a new scrape")
        actions.add_row("2", "Resume a previous job")
        actions.add_row("3", "Search businesses")
        actions.add_row("4", "Export results")
        actions.add_row("5", "Exit")

        console.print(actions)
        console.print()

        choice = Prompt.ask(
            "[bold]Select option",
            choices=["1", "2", "3", "4", "5"],
            default="1"
        )

        if choice == "1":
            return "new"
        elif choice == "2":
            return "resume"
        elif choice == "3":
            return "search"
        elif choice == "4":
            return "export"
        else:
            return "cancel"

    def _step_queries(self) -> bool:
        """Step 2: Configure search queries."""
        console.print()
        console.rule("[bold cyan]Step 2: Search Queries")
        console.print()

        console.print("[dim]Enter the business types you want to search for.[/dim]")
        console.print("[dim]You can enter multiple queries separated by semicolons.[/dim]")
        console.print()

        # Show example
        console.print("[dim]Examples:[/dim]")
        console.print("  [cyan]Pet Cremation[/cyan]")
        console.print("  [cyan]Pet Cremation; Pet Cemetery; Animal Funeral[/cyan]")
        console.print()

        query_input = Prompt.ask("[bold]Search query/queries")

        if not query_input.strip():
            console.print("[red]Error: At least one query is required.[/red]")
            return False

        # Parse queries
        queries = [q.strip() for q in query_input.split(";") if q.strip()]

        if not queries:
            console.print("[red]Error: At least one query is required.[/red]")
            return False

        self.result["queries"] = queries

        # Confirm
        console.print()
        console.print(f"[green]Queries configured:[/green] {len(queries)}")
        for i, q in enumerate(queries, 1):
            console.print(f"  {i}. [cyan]{q}[/cyan]")

        return True

    def _step_locations(self) -> bool:
        """Step 3: Select zip-code-based locations."""
        console.print()
        console.rule("[bold cyan]Step 3: Location Selection")
        console.print()

        zip_count = self.db.get_zip_count()
        if zip_count == 0:
            console.print("[red]Error: No zip codes in database.[/red]")
            console.print("Run [cyan]python setup.py[/cyan] first.")
            return False

        console.print(f"[dim]{zip_count:,} US zip codes available, ranked by 2020 Census population.[/dim]")
        console.print()

        preset_table = Table(show_header=True, box=box.SIMPLE, padding=(0, 2))
        preset_table.add_column("", style="bold yellow", width=4)
        preset_table.add_column("Option", width=38)
        preset_table.add_column("Zips", justify="right", width=8)

        for key, (_, desc, count) in self.ZIP_PRESETS.items():
            if count == 0:
                preset_table.add_row(key, desc, "—")
            else:
                preset_table.add_row(key, desc, f"{min(count, zip_count):,}")

        console.print(preset_table)
        console.print()

        choice = Prompt.ask(
            "[bold]Select option",
            choices=list(self.ZIP_PRESETS.keys()),
            default="2"
        )

        preset_name, desc, _ = self.ZIP_PRESETS[choice]

        if preset_name == "city":
            return self._get_zip_city()
        elif preset_name == "state":
            return self._get_zip_state()
        elif preset_name == "custom":
            return self._get_custom_zips()
        else:
            zips = self.db.get_zips_by_preset(preset_name)
            if not zips:
                console.print("[red]Error: No zip codes found. Run setup.py first.[/red]")
                return False

            self.result["zip_preset"] = preset_name
            self.result["locations"] = zips
            self.result["location_label"] = f"zips_{preset_name}"

            console.print()
            console.print(f"[green]Selected:[/green] {desc}")
            console.print(f"[dim]{len(zips):,} zip codes  •  first: {', '.join(zips[:4])}...[/dim]")
            return True

    def _get_zip_city(self) -> bool:
        """Decompose a city into its zip codes."""
        console.print()
        console.print("[dim]Format: City, State  (e.g. Dallas, TX  or  Chicago, IL)[/dim]")
        console.print()

        raw = Prompt.ask("[bold]City, State").strip()
        if not raw:
            console.print("[red]Error: input required.[/red]")
            return False

        if ',' in raw:
            city_part, state_part = raw.split(',', 1)
        else:
            parts = raw.rsplit(' ', 1)
            city_part, state_part = (parts[0], parts[1]) if len(parts) == 2 else (raw, '')

        city_part = city_part.strip().title()
        state_part = state_part.strip().upper()

        zips = self.db.get_zips_by_city(city_part, state_part)
        if not zips:
            console.print(f"[red]No zip codes found for '{city_part}, {state_part}'. Check spelling or run setup.py.[/red]")
            return False

        self.result["zip_city"] = f"{city_part}, {state_part}"
        self.result["locations"] = zips
        safe = "".join(c for c in f"{city_part}{state_part}" if c.isalnum())[:20]
        self.result["location_label"] = f"zips_{safe}"

        console.print()
        console.print(f"[green]Found {len(zips)} zip codes[/green] for {city_part}, {state_part}")
        console.print(f"[dim]{', '.join(zips[:6])}{'...' if len(zips) > 6 else ''}[/dim]")
        return True

    def _get_zip_state(self) -> bool:
        """Get all zip codes for one or more states."""
        console.print()
        console.print("[dim]Enter state abbreviations separated by commas  (e.g. TX  or  TX,CA,FL)[/dim]")
        console.print()

        raw = Prompt.ask("[bold]State(s)").strip().upper()
        if not raw:
            console.print("[red]Error: input required.[/red]")
            return False

        states = [s.strip() for s in raw.split(',') if s.strip()]
        all_zips = []
        for state in states:
            state_zips = self.db.get_zips_by_state(state)
            if not state_zips:
                console.print(f"[yellow]Warning: no zip codes found for '{state}'[/yellow]")
            else:
                console.print(f"  [cyan]{state}[/cyan]: {len(state_zips):,} zip codes")
                all_zips.extend(state_zips)

        if not all_zips:
            console.print("[red]No zip codes found for the given state(s).[/red]")
            return False

        self.result["zip_state"] = raw
        self.result["locations"] = all_zips
        self.result["location_label"] = f"zips_{raw.replace(',', '_')}"

        console.print()
        console.print(f"[green]Total: {len(all_zips):,} zip codes[/green]")
        return True

    def _get_custom_zips(self) -> bool:
        """Enter zip codes manually."""
        console.print()
        console.print("[dim]Enter 5-digit zip codes separated by semicolons  (e.g. 75201; 90011; 60629)[/dim]")
        console.print()

        raw = Prompt.ask("[bold]Zip codes").strip()
        if not raw:
            console.print("[red]Error: input required.[/red]")
            return False

        zips = [z.strip().zfill(5) for z in raw.replace(',', ';').split(';') if z.strip().isdigit()]
        if not zips:
            console.print("[red]Error: no valid zip codes entered.[/red]")
            return False

        self.result["locations"] = zips
        self.result["location_label"] = "zips_custom"

        console.print()
        console.print(f"[green]{len(zips)} zip codes entered:[/green] {', '.join(zips[:8])}{'...' if len(zips) > 8 else ''}")
        return True

    def _step_options(self) -> bool:
        """Step 4: Configure additional options."""
        console.print()
        console.rule("[bold cyan]Step 4: Options")
        console.print()

        # Parallel workers (new option)
        console.print("[dim]Parallel workers run multiple browser instances simultaneously.[/dim]")
        console.print("[dim]More workers = faster scraping, but uses more system resources.[/dim]")
        console.print()
        self.result["parallel"] = IntPrompt.ask(
            "[bold]Number of parallel workers[/bold] (1-4)",
            default=3
        )
        # Clamp to valid range
        self.result["parallel"] = max(1, min(4, self.result["parallel"]))

        console.print()

        # Deep search option
        console.print("[dim]Deep search clicks into each listing to extract phone numbers and websites.[/dim]")
        console.print("[dim]This is slower but provides complete contact data.[/dim]")
        console.print()
        self.result["deep"] = Confirm.ask(
            "[bold]Enable deep search?[/bold] (extract phones & websites)",
            default=False
        )

        console.print()

        # Headless mode (default: True)
        self.result["headless"] = Confirm.ask(
            "[bold]Run in headless mode?[/bold] (no browser window)",
            default=True
        )

        # Skip test (default: True)
        self.result["skip_test"] = Confirm.ask(
            "[bold]Skip test phase?[/bold] (go straight to full scrape)",
            default=True
        )

        # Live monitor
        self.result["live"] = Confirm.ask(
            "[bold]Enable live business feed monitor?[/bold]",
            default=True
        )

        # Proxy (default: disabled)
        proxy_configured = "USERNAME" not in self.config.get("proxy", {}).get("server", "USERNAME")

        if proxy_configured:
            self.result["no_proxy"] = not Confirm.ask(
                "[bold]Use proxy?[/bold] (configured in config.yaml)",
                default=False
            )
        else:
            console.print("[dim]Proxy not configured - running without proxy.[/dim]")
            self.result["no_proxy"] = True

        # Bandwidth budget
        budget_default = self.config.get("bandwidth", {}).get("budget_mb", 900)
        console.print()
        console.print(f"[dim]Bandwidth budget (default: {budget_default} MB, 0 = unlimited)[/dim]")
        self.result["budget"] = IntPrompt.ask(
            "[bold]Bandwidth budget (MB)",
            default=budget_default
        )

        # Cross-scrape dedup
        existing_tables = self.db.list_result_tables()
        if existing_tables:
            console.print()
            console.print(f"[dim]Found {len(existing_tables)} existing result tables.[/dim]")
            exclude = Confirm.ask(
                "[bold]Exclude businesses from previous scrapes?[/bold]",
                default=False
            )
            if exclude:
                console.print()
                console.print("[dim]Enter table pattern (e.g., results_Pet*, or leave blank for all):[/dim]")
                pattern = Prompt.ask("[bold]Exclude pattern", default="*")
                self.result["exclude_tables"] = pattern
            else:
                self.result["exclude_tables"] = None
        else:
            self.result["exclude_tables"] = None

        return True

    def _step_preview(self) -> bool:
        """Step 5: Preview configuration and confirm."""
        console.print()
        console.rule("[bold cyan]Step 5: Review Configuration")
        console.print()

        # Build preview table
        preview = Table(show_header=False, box=box.ROUNDED, padding=(0, 2))
        preview.add_column("Setting", style="bold", width=20)
        preview.add_column("Value", style="cyan")

        # Queries
        queries_str = "; ".join(self.result["queries"][:3])
        if len(self.result["queries"]) > 3:
            queries_str += f" (+{len(self.result['queries']) - 3} more)"
        preview.add_row("Queries", queries_str)

        # Locations
        locations = self.result.get("locations", [])
        label = self.result.get("location_label", "custom")
        if self.result.get("zip_city"):
            loc_str = f"{len(locations)} zip codes in {self.result['zip_city']}"
        elif self.result.get("zip_state"):
            loc_str = f"{len(locations):,} zip codes  ({self.result['zip_state']})"
        elif self.result.get("zip_preset"):
            loc_str = f"{len(locations):,} zip codes  (preset: {self.result['zip_preset']})"
        else:
            loc_str = f"{len(locations)} custom zip codes"
        preview.add_row("Locations", loc_str)

        # Options
        preview.add_row("Parallel Workers", str(self.result.get("parallel", 3)))
        preview.add_row("Deep Search", "Yes (phones & websites)" if self.result.get("deep") else "No")
        preview.add_row("Headless", "Yes" if self.result.get("headless") else "No")
        preview.add_row("Skip Test", "Yes" if self.result.get("skip_test") else "No")
        preview.add_row("Live Monitor", "Yes" if self.result.get("live") else "No")
        preview.add_row("Proxy", "Disabled" if self.result.get("no_proxy") else "Enabled")
        preview.add_row("Budget", f"{self.result.get('budget', 0)} MB" if self.result.get('budget') else "Unlimited")

        if self.result.get("exclude_tables"):
            preview.add_row("Exclude Pattern", self.result["exclude_tables"])

        # Estimate
        total_searches = len(self.result["queries"]) * len(locations)
        preview.add_row("", "")
        preview.add_row("Total Searches", f"{total_searches:,}")

        console.print(preview)
        console.print()

        return Confirm.ask("[bold green]Start scraping?[/bold green]", default=True)

    def _handle_search(self):
        """Handle search action."""
        console.print()
        console.rule("[bold cyan]Search Businesses")
        console.print()

        tables = self.db.list_result_tables()
        if not tables:
            console.print("[yellow]No result tables found. Run a scrape first.[/yellow]")
            return

        # Choose search type
        console.print("[bold]Search by:[/bold]")
        console.print("  1. Any field (name, phone, address, city)")
        console.print("  2. Business name")
        console.print("  3. Phone number")
        console.print("  4. City")
        console.print()

        search_type = Prompt.ask("[bold]Select search type", choices=["1", "2", "3", "4"], default="1")

        console.print()
        search_query = Prompt.ask("[bold]Enter search term")

        if not search_query.strip():
            console.print("[red]Search term cannot be empty.[/red]")
            return

        # Perform search
        console.print()
        console.print("[dim]Searching...[/dim]")

        if search_type == "1":
            results = self.db.search(search_query, limit=100)
        elif search_type == "2":
            results = self.db.search_by_name(search_query, limit=100)
        elif search_type == "3":
            results = self.db.search_by_phone(search_query, limit=100)
        else:
            results = self.db.search_by_city(search_query, limit=100)

        if not results:
            console.print(f"[yellow]No results found for '{search_query}'[/yellow]")
            return

        # Display results
        console.print()
        result_table = Table(
            title=f"Search Results ({len(results)} found)",
            show_header=True,
            header_style="bold magenta"
        )
        result_table.add_column("#", style="dim", width=4)
        result_table.add_column("Name", style="cyan", width=28, no_wrap=True)
        result_table.add_column("Phone", style="yellow", width=14)
        result_table.add_column("City", style="green", width=14)
        result_table.add_column("Rating", justify="right", width=6)

        for i, r in enumerate(results[:30], 1):
            name = (r.get('name') or '-')[:26]
            phone = (r.get('phone_number') or '-')[:14]
            city = (r.get('search_city') or '-')[:12]
            rating = f"{r.get('rating', '-')}" if r.get('rating') else '-'
            result_table.add_row(str(i), name, phone, city, rating)

        console.print(result_table)

        if len(results) > 30:
            console.print(f"[dim]Showing 30 of {len(results)} results.[/dim]")

        # Offer export
        console.print()
        if Confirm.ask("[bold]Export search results to file?[/bold]", default=False):
            self._export_search_results(results, search_query)

    def _export_search_results(self, results: list, search_query: str):
        """Export search results to file."""
        import pandas as pd

        # Choose format
        console.print()
        console.print("[bold]Export format:[/bold]")
        console.print("  1. CSV")
        console.print("  2. JSON")
        console.print("  3. Excel (xlsx)")
        console.print()

        format_choice = Prompt.ask("[bold]Select format", choices=["1", "2", "3"], default="1")

        export_dir = Path(self.config.get("output_directory", "exports"))
        export_dir.mkdir(exist_ok=True)

        # Generate filename
        safe_query = "".join(c for c in search_query if c.isalnum() or c == "_")[:20]
        base_name = f"search_{safe_query}"

        df = pd.DataFrame(results)

        if format_choice == "1":
            filename = export_dir / f"{base_name}.csv"
            df.to_csv(filename, index=False)
        elif format_choice == "2":
            filename = export_dir / f"{base_name}.json"
            df.to_json(filename, orient='records', indent=2)
        else:
            filename = export_dir / f"{base_name}.xlsx"
            df.to_excel(filename, index=False)

        console.print(Panel(
            f"[green]Exported {len(results)} results to:[/green]\n[cyan]{filename}[/cyan]",
            title="Export Complete"
        ))

    def _show_table_page(self, tables: list, page: int, page_size: int = 20):
        """Render one page of the tables list."""
        total_pages = (len(tables) + page_size - 1) // page_size
        start = page * page_size
        end = min(start + page_size, len(tables))

        table_list = Table(show_header=True, box=box.SIMPLE)
        table_list.add_column("#", style="bold yellow", width=5)
        table_list.add_column("Table Name", style="cyan")
        table_list.add_column("Rows", justify="right", width=7)

        for i, tbl in enumerate(tables[start:end], start + 1):
            count = self.db.get_table_row_count(tbl)
            table_list.add_row(str(i), tbl, str(count))

        table_list.add_row("A", "[bold]Export ALL tables[/bold]", "")

        console.print(table_list)
        console.print(
            f"[dim]Page {page + 1}/{total_pages}  "
            f"({start + 1}–{end} of {len(tables)} tables)  "
            f"  [N] next  [P] prev  [#] select[/dim]"
        )

    def _handle_export(self):
        """Handle export action."""
        console.print()
        console.rule("[bold cyan]Export Results")
        console.print()

        tables = self.db.list_result_tables()
        if not tables:
            console.print("[yellow]No result tables found.[/yellow]")
            return

        PAGE_SIZE = 20
        page = 0
        total_pages = (len(tables) + PAGE_SIZE - 1) // PAGE_SIZE

        while True:
            console.print()
            self._show_table_page(tables, page, PAGE_SIZE)
            console.print()

            choice = Prompt.ask(
                "[bold]# to select, N next page, P prev page, A for all",
                default="1"
            )

            if choice.upper() == "N":
                page = min(page + 1, total_pages - 1)
                continue
            elif choice.upper() == "P":
                page = max(page - 1, 0)
                continue
            else:
                break

        # Choose format
        console.print()
        console.print("[bold]Export format:[/bold]")
        console.print("  1. CSV")
        console.print("  2. JSON")
        console.print("  3. Excel (xlsx)")
        console.print()

        format_choice = Prompt.ask("[bold]Select format", choices=["1", "2", "3"], default="1")
        format_ext = {
            "1": ("csv", lambda df, f: df.to_csv(f, index=False)),
            "2": ("json", lambda df, f: df.to_json(f, orient='records', indent=2)),
            "3": ("xlsx", lambda df, f: df.to_excel(f, index=False))
        }
        ext, export_func = format_ext[format_choice]

        export_dir = Path(self.config.get("output_directory", "exports"))
        export_dir.mkdir(exist_ok=True)

        # Handle export
        if choice.upper() == "A":
            # Export all tables
            console.print(f"\n[cyan]Exporting {len(tables)} tables...[/cyan]")
            for tbl in tables:
                df = self.db.export_table_to_dataframe(tbl)
                filename = export_dir / f"{tbl}.{ext}"
                export_func(df, filename)
                console.print(f"  [green]✓[/green] {tbl} ({len(df)} rows)")

            console.print(Panel(
                f"[green]Exported {len(tables)} tables to:[/green]\n[cyan]{export_dir}[/cyan]",
                title="Export Complete"
            ))
        else:
            # Export single table
            if choice.isdigit() and 1 <= int(choice) <= len(tables):
                table_name = tables[int(choice) - 1]
            elif choice in tables:
                table_name = choice
            else:
                console.print("[red]Invalid selection.[/red]")
                return

            df = self.db.export_table_to_dataframe(table_name)
            filename = export_dir / f"{table_name}.{ext}"
            export_func(df, filename)

            console.print(Panel(
                f"[green]Exported {len(df)} rows to:[/green]\n[cyan]{filename}[/cyan]",
                title="Export Complete"
            ))

    def _handle_resume(self) -> Optional[Dict[str, Any]]:
        """Handle resume action."""
        from job_manager import JobManager

        console.print()
        console.rule("[bold cyan]Resume Job")
        console.print()

        job_manager = JobManager(self.db)
        jobs = job_manager.get_resumable_jobs()

        if not jobs:
            console.print("[yellow]No resumable jobs found.[/yellow]")
            return None

        # Show jobs
        job_table = Table(show_header=True, box=box.SIMPLE)
        job_table.add_column("#", style="bold yellow", width=4)
        job_table.add_column("Job ID", style="cyan", width=35)
        job_table.add_column("Progress")
        job_table.add_column("Status")

        for i, job in enumerate(jobs[:10], 1):
            status_color = "yellow" if job["status"] == "running" else "blue"
            job_table.add_row(
                str(i),
                job["job_id"][:33],
                job["progress"],
                f"[{status_color}]{job['status']}[/{status_color}]"
            )

        console.print(job_table)
        console.print()

        choice = Prompt.ask("[bold]Enter job number", default="1")

        if not choice.isdigit() or not 1 <= int(choice) <= len(jobs):
            console.print("[red]Invalid selection.[/red]")
            return None

        job_id = jobs[int(choice) - 1]["job_id"]
        return {"resume": job_id}


def run_wizard(config: dict, db: DatabaseManager) -> Optional[Dict[str, Any]]:
    """Run the wizard and return configuration or None."""
    wizard = ScrapeWizard(config, db)
    return wizard.run()