"""
AIPTX Interactive Setup Wizard
==============================

First-run setup wizard that guides users through configuration.
Collects API keys and settings interactively with a beautiful TUI.

Usage:
    aiptx setup              # Run setup wizard
    aiptx scan example.com   # Auto-triggers if not configured
"""

import os
import sys
from pathlib import Path
from typing import Optional, Tuple

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


console = Console()


# ============================================================================
# Configuration File Management
# ============================================================================

def get_config_path() -> Path:
    """Get the path to the .env config file."""
    # Check for existing .env in current directory
    local_env = Path(".env")
    if local_env.exists():
        return local_env

    # Check for global config in home directory
    home_env = Path.home() / ".aiptx" / ".env"
    if home_env.exists():
        return home_env

    # Default to home directory for new installations
    return home_env


def load_existing_config() -> dict:
    """Load existing configuration from .env file."""
    config = {}
    config_path = get_config_path()

    if config_path.exists():
        with open(config_path, "r") as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, value = line.split("=", 1)
                    config[key.strip()] = value.strip().strip('"').strip("'")

    return config


def save_config(config: dict, path: Optional[Path] = None) -> Path:
    """Save configuration to .env file."""
    if path is None:
        path = Path.home() / ".aiptx" / ".env"

    # Create directory if needed
    path.parent.mkdir(parents=True, exist_ok=True)

    # Build config content
    lines = [
        "# AIPTX Configuration",
        "# Generated by 'aiptx setup'",
        "# Edit this file or run 'aiptx setup' again to reconfigure",
        "",
    ]

    # Group settings
    sections = {
        "LLM": ["ANTHROPIC_API_KEY", "OPENAI_API_KEY", "DEEPSEEK_API_KEY", "LLM_API_KEY",
                "AIPT_LLM__PROVIDER", "AIPT_LLM__MODEL"],
        "Acunetix": ["AIPT_SCANNERS__ACUNETIX_URL", "AIPT_SCANNERS__ACUNETIX_API_KEY"],
        "Burp Suite": ["AIPT_SCANNERS__BURP_URL", "AIPT_SCANNERS__BURP_API_KEY"],
        "Nessus": ["AIPT_SCANNERS__NESSUS_URL", "AIPT_SCANNERS__NESSUS_ACCESS_KEY",
                   "AIPT_SCANNERS__NESSUS_SECRET_KEY"],
        "OWASP ZAP": ["AIPT_SCANNERS__ZAP_URL", "AIPT_SCANNERS__ZAP_API_KEY"],
        "VPS": ["AIPT_VPS__HOST", "AIPT_VPS__USER", "AIPT_VPS__KEY_PATH", "AIPT_VPS__PORT"],
    }

    for section, keys in sections.items():
        section_values = [(k, config.get(k)) for k in keys if config.get(k)]
        if section_values:
            lines.append(f"\n# {section} Configuration")
            for key, value in section_values:
                lines.append(f'{key}="{value}"')

    # Write file
    with open(path, "w") as f:
        f.write("\n".join(lines) + "\n")

    # Secure the file (readable only by owner)
    os.chmod(path, 0o600)

    return path


def is_configured() -> bool:
    """Check if AIPTX has been configured with at least an LLM API key."""
    # Check environment variables
    for key in ["ANTHROPIC_API_KEY", "OPENAI_API_KEY", "DEEPSEEK_API_KEY", "LLM_API_KEY"]:
        if os.getenv(key):
            return True

    # Check .env files
    config = load_existing_config()
    for key in ["ANTHROPIC_API_KEY", "OPENAI_API_KEY", "DEEPSEEK_API_KEY", "LLM_API_KEY"]:
        if config.get(key):
            return True

    return False


# ============================================================================
# Interactive Setup Wizard
# ============================================================================

def print_welcome():
    """Print welcome banner."""
    banner = """
[bold cyan]╔═══════════════════════════════════════════════════════════════╗
║                                                               ║
║     █████╗ ██╗██████╗ ████████╗██╗  ██╗                      ║
║    ██╔══██╗██║██╔══██╗╚══██╔══╝╚██╗██╔╝                      ║
║    ███████║██║██████╔╝   ██║    ╚███╔╝                       ║
║    ██╔══██║██║██╔═══╝    ██║    ██╔██╗                       ║
║    ██║  ██║██║██║        ██║   ██╔╝ ██╗                      ║
║    ╚═╝  ╚═╝╚═╝╚═╝        ╚═╝   ╚═╝  ╚═╝                      ║
║                                                               ║
║         AI-Powered Penetration Testing Framework              ║
║                                                               ║
╚═══════════════════════════════════════════════════════════════╝[/bold cyan]
"""
    console.print(banner)
    console.print("[bold green]Welcome to AIPTX Setup![/bold green]")
    console.print("This wizard will help you configure AIPTX for first use.\n")


def setup_llm() -> dict:
    """Configure LLM provider and API key."""
    config = {}

    console.print(Panel(
        "[bold]Step 1: LLM Configuration[/bold]\n\n"
        "AIPTX uses AI to guide penetration testing.\n"
        "You need an API key from one of the supported providers.",
        title="🤖 AI Provider",
        border_style="cyan"
    ))

    # Choose provider
    console.print("\n[bold]Select your LLM provider:[/bold]")
    console.print("  [1] Anthropic (Claude) - [green]Recommended[/green]")
    console.print("  [2] OpenAI (GPT-4)")
    console.print("  [3] DeepSeek")
    console.print("  [4] Other (custom)")

    choice = Prompt.ask(
        "\nEnter choice",
        choices=["1", "2", "3", "4"],
        default="1"
    )

    providers = {
        "1": ("anthropic", "claude-sonnet-4-20250514", "ANTHROPIC_API_KEY"),
        "2": ("openai", "gpt-4o", "OPENAI_API_KEY"),
        "3": ("deepseek", "deepseek-chat", "DEEPSEEK_API_KEY"),
        "4": ("custom", "", "LLM_API_KEY"),
    }

    provider, model, key_name = providers[choice]

    if choice == "4":
        provider = Prompt.ask("Enter provider name")
        model = Prompt.ask("Enter model name")

    config["AIPT_LLM__PROVIDER"] = provider
    config["AIPT_LLM__MODEL"] = model

    # Get API key
    console.print(f"\n[bold]Enter your {provider.title()} API key:[/bold]")

    if provider == "anthropic":
        console.print("[dim]Get one at: https://console.anthropic.com/settings/keys[/dim]")
    elif provider == "openai":
        console.print("[dim]Get one at: https://platform.openai.com/api-keys[/dim]")
    elif provider == "deepseek":
        console.print("[dim]Get one at: https://platform.deepseek.com/api_keys[/dim]")

    api_key = Prompt.ask("API Key", password=True)

    if api_key:
        config[key_name] = api_key

    return config


def setup_scanners() -> dict:
    """Configure enterprise scanners (optional)."""
    config = {}

    console.print(Panel(
        "[bold]Step 2: Enterprise Scanners (Optional)[/bold]\n\n"
        "AIPTX integrates with enterprise DAST scanners for\n"
        "comprehensive vulnerability assessment.",
        title="🔍 Scanners",
        border_style="cyan"
    ))

    if not Confirm.ask("\nDo you want to configure enterprise scanners?", default=False):
        console.print("[dim]Skipping scanner configuration...[/dim]\n")
        return config

    # Acunetix
    if Confirm.ask("\n[bold]Configure Acunetix?[/bold]", default=False):
        url = Prompt.ask("  Acunetix URL (e.g., https://your-instance:3443)")
        api_key = Prompt.ask("  Acunetix API Key", password=True)
        if url:
            config["AIPT_SCANNERS__ACUNETIX_URL"] = url
        if api_key:
            config["AIPT_SCANNERS__ACUNETIX_API_KEY"] = api_key

    # Burp Suite
    if Confirm.ask("\n[bold]Configure Burp Suite Enterprise?[/bold]", default=False):
        url = Prompt.ask("  Burp Suite URL (e.g., https://your-burp:8080)")
        api_key = Prompt.ask("  Burp Suite API Key", password=True)
        if url:
            config["AIPT_SCANNERS__BURP_URL"] = url
        if api_key:
            config["AIPT_SCANNERS__BURP_API_KEY"] = api_key

    # Nessus
    if Confirm.ask("\n[bold]Configure Nessus?[/bold]", default=False):
        url = Prompt.ask("  Nessus URL (e.g., https://your-nessus:8834)")
        access_key = Prompt.ask("  Nessus Access Key", password=True)
        secret_key = Prompt.ask("  Nessus Secret Key", password=True)
        if url:
            config["AIPT_SCANNERS__NESSUS_URL"] = url
        if access_key:
            config["AIPT_SCANNERS__NESSUS_ACCESS_KEY"] = access_key
        if secret_key:
            config["AIPT_SCANNERS__NESSUS_SECRET_KEY"] = secret_key

    # OWASP ZAP
    if Confirm.ask("\n[bold]Configure OWASP ZAP?[/bold]", default=False):
        url = Prompt.ask("  ZAP URL (e.g., http://localhost:8080)")
        api_key = Prompt.ask("  ZAP API Key (leave empty if disabled)", password=True)
        if url:
            config["AIPT_SCANNERS__ZAP_URL"] = url
        if api_key:
            config["AIPT_SCANNERS__ZAP_API_KEY"] = api_key

    return config


def setup_vps() -> dict:
    """Configure VPS for remote execution (optional)."""
    config = {}

    console.print(Panel(
        "[bold]Step 3: VPS Configuration (Optional)[/bold]\n\n"
        "Run security tools on a remote VPS to avoid\n"
        "network restrictions and maintain anonymity.",
        title="🖥️  VPS",
        border_style="cyan"
    ))

    if not Confirm.ask("\nDo you want to configure a VPS for remote execution?", default=False):
        console.print("[dim]Skipping VPS configuration...[/dim]\n")
        return config

    host = Prompt.ask("  VPS IP or hostname")
    user = Prompt.ask("  SSH username", default="ubuntu")
    key_path = Prompt.ask("  Path to SSH private key (e.g., ~/.ssh/id_rsa)")
    port = Prompt.ask("  SSH port", default="22")

    if host:
        config["AIPT_VPS__HOST"] = host
    if user:
        config["AIPT_VPS__USER"] = user
    if key_path:
        # Expand ~ to home directory
        expanded_path = str(Path(key_path).expanduser())
        config["AIPT_VPS__KEY_PATH"] = expanded_path
    if port:
        config["AIPT_VPS__PORT"] = port

    return config


def show_summary(config: dict):
    """Show configuration summary."""
    console.print(Panel(
        "[bold]Configuration Summary[/bold]",
        title="📋 Summary",
        border_style="green"
    ))

    table = Table(box=box.ROUNDED)
    table.add_column("Setting", style="cyan")
    table.add_column("Value", style="green")

    # LLM
    provider = config.get("AIPT_LLM__PROVIDER", "Not set")
    model = config.get("AIPT_LLM__MODEL", "Not set")
    has_key = any(k in config for k in ["ANTHROPIC_API_KEY", "OPENAI_API_KEY", "DEEPSEEK_API_KEY", "LLM_API_KEY"])

    table.add_row("LLM Provider", provider)
    table.add_row("LLM Model", model)
    table.add_row("LLM API Key", "✓ Configured" if has_key else "✗ Not set")

    # Scanners
    table.add_row("─" * 20, "─" * 20)
    table.add_row("Acunetix", "✓ Configured" if config.get("AIPT_SCANNERS__ACUNETIX_URL") else "○ Not configured")
    table.add_row("Burp Suite", "✓ Configured" if config.get("AIPT_SCANNERS__BURP_URL") else "○ Not configured")
    table.add_row("Nessus", "✓ Configured" if config.get("AIPT_SCANNERS__NESSUS_URL") else "○ Not configured")
    table.add_row("OWASP ZAP", "✓ Configured" if config.get("AIPT_SCANNERS__ZAP_URL") else "○ Not configured")

    # VPS
    table.add_row("─" * 20, "─" * 20)
    table.add_row("VPS", "✓ " + config.get("AIPT_VPS__HOST", "") if config.get("AIPT_VPS__HOST") else "○ Not configured")

    console.print(table)


def run_setup_wizard(force: bool = False) -> bool:
    """
    Run the interactive setup wizard.

    Args:
        force: Run even if already configured

    Returns:
        True if setup completed successfully
    """
    try:
        # Check if already configured
        if not force and is_configured():
            console.print("[yellow]AIPTX is already configured.[/yellow]")
            if not Confirm.ask("Do you want to reconfigure?", default=False):
                return False

        print_welcome()

        # Collect configuration
        config = load_existing_config()  # Start with existing config

        # Step 1: LLM (required)
        llm_config = setup_llm()
        config.update(llm_config)

        # Check if we got an API key
        has_key = any(k in config for k in ["ANTHROPIC_API_KEY", "OPENAI_API_KEY", "DEEPSEEK_API_KEY", "LLM_API_KEY"])
        if not has_key:
            console.print("\n[bold red]Error:[/bold red] An LLM API key is required to use AIPTX.")
            console.print("Please run [bold]aiptx setup[/bold] again with a valid API key.")
            return False

        # Step 2: Scanners (optional)
        scanner_config = setup_scanners()
        config.update(scanner_config)

        # Step 3: VPS (optional)
        vps_config = setup_vps()
        config.update(vps_config)

        # Show summary
        console.print()
        show_summary(config)

        # Confirm and save
        console.print()
        if Confirm.ask("[bold]Save this configuration?[/bold]", default=True):
            config_path = save_config(config)

            console.print(Panel(
                f"[bold green]✓ Configuration saved![/bold green]\n\n"
                f"Config file: [cyan]{config_path}[/cyan]\n\n"
                f"You can now run:\n"
                f"  [bold]aiptx scan example.com[/bold]     - Run a security scan\n"
                f"  [bold]aiptx status[/bold]              - Check configuration\n"
                f"  [bold]aiptx setup[/bold]               - Reconfigure AIPTX",
                title="🎉 Setup Complete",
                border_style="green"
            ))

            # Load the config into environment for immediate use
            for key, value in config.items():
                os.environ[key] = value

            return True
        else:
            console.print("[yellow]Setup cancelled. No changes saved.[/yellow]")
            return False

    except KeyboardInterrupt:
        # Handle Ctrl+C gracefully
        console.print("\n[yellow]Setup cancelled.[/yellow]")
        return False


def prompt_first_run_setup() -> bool:
    """
    Prompt for setup on first run when configuration is missing.

    Returns:
        True if setup completed and user can proceed
    """
    try:
        console.print(Panel(
            "[bold yellow]⚠ AIPTX is not configured![/bold yellow]\n\n"
            "This appears to be your first time running AIPTX.\n"
            "You need to configure at least an LLM API key to proceed.",
            title="First Run Setup Required",
            border_style="yellow"
        ))

        if Confirm.ask("\nWould you like to run the setup wizard now?", default=True):
            return run_setup_wizard(force=True)
        else:
            console.print("\n[dim]You can run setup later with: [bold]aiptx setup[/bold][/dim]")
            console.print("[dim]Or set environment variables manually:[/dim]")
            console.print("[dim]  export ANTHROPIC_API_KEY=your-key-here[/dim]\n")
            return False
    except KeyboardInterrupt:
        console.print("\n[yellow]Cancelled.[/yellow]")
        return False


# ============================================================================
# CLI Entry Points
# ============================================================================

def main():
    """Standalone setup wizard entry point."""
    run_setup_wizard(force=True)


if __name__ == "__main__":
    main()
