#!/usr/bin/env python3
"""
souleyez.ui.tool_setup - Tool installation wizard for Ubuntu/Debian systems

Helps users install pentesting tools that aren't available via apt on Ubuntu.
Handles PATH configuration for go and pipx installed tools.
"""
import click
import shutil
import subprocess
import sys
import os
from pathlib import Path
from typing import List, Dict, Optional

from souleyez.utils.tool_checker import (
    detect_distro,
    get_missing_tools,
    get_tools_by_category,
    get_category_name,
    check_tool,
    get_tool_stats,
    EXTERNAL_TOOLS,
)
from souleyez.ui.design_system import DesignSystem

# Prerequisites needed for various install methods
PREREQUISITES = {
    'build-deps': {
        'check': None,  # Always install to ensure all deps present
        'install': 'sudo apt install -y build-essential python3-dev libxml2-dev libxslt1-dev libuv1-dev libffi-dev libssl-dev rustc cargo',
        'description': 'Build dependencies for Python packages with native extensions',
        'always_install': True  # Flag to always run this
    },
    'pipx': {
        'check': 'pipx',
        'install': 'sudo apt install -y pipx && pipx ensurepath',
        'description': 'Python application installer (for theHarvester, NetExec, etc.)',
        'path_additions': ['~/.local/bin']
    },
    'golang': {
        'check': 'go',
        'install': 'sudo apt install -y golang-go',
        'description': 'Go programming language (for nuclei, ffuf)',
        'path_additions': ['~/go/bin']
    },
    'ruby': {
        'check': 'gem',
        'install': 'sudo apt install -y ruby-full ruby-dev build-essential',
        'description': 'Ruby programming language (for wpscan)'
    },
    'snap': {
        'check': 'snap',
        'install': 'sudo apt install -y snapd',
        'description': 'Snap package manager (for enum4linux)'
    },
    'git': {
        'check': 'git',
        'install': 'sudo apt install -y git',
        'description': 'Git version control (for exploitdb, Responder)'
    },
}


def _ensure_path_configured():
    """Ensure go and pipx paths are in current session PATH."""
    home = os.path.expanduser("~")
    paths_to_add = [
        os.path.join(home, "go", "bin"),
        os.path.join(home, ".local", "bin"),
        "/snap/bin",
    ]

    current_path = os.environ.get("PATH", "")
    new_paths = []

    for p in paths_to_add:
        if os.path.exists(p) and p not in current_path:
            new_paths.append(p)

    if new_paths:
        os.environ["PATH"] = ":".join(new_paths) + ":" + current_path


def _add_paths_to_bashrc():
    """Add tool paths to ~/.bashrc if not already present."""
    bashrc = Path.home() / ".bashrc"
    paths_to_add = [
        ('export PATH="$HOME/go/bin:$PATH"', "go/bin"),
        ('export PATH="$HOME/.local/bin:$PATH"', ".local/bin"),
    ]

    if not bashrc.exists():
        return

    try:
        content = bashrc.read_text()
        additions = []

        for line, marker in paths_to_add:
            if marker not in content:
                additions.append(line)

        if additions:
            with open(bashrc, "a") as f:
                f.write("\n# Added by souleyez setup\n")
                for line in additions:
                    f.write(line + "\n")
    except Exception:
        pass  # Don't fail on PATH configuration issues


def _run_command(cmd: str, console, description: str = "", capture: bool = False) -> tuple:
    """Run a command with proper error handling."""
    try:
        if capture:
            result = subprocess.run(
                cmd,
                shell=True,
                capture_output=True,
                text=True,
                timeout=600  # 10 minute timeout
            )
            return result.returncode == 0, result.stdout, result.stderr
        else:
            result = subprocess.run(
                cmd,
                shell=True,
                timeout=600
            )
            return result.returncode == 0, "", ""
    except subprocess.TimeoutExpired:
        return False, "", "Command timed out"
    except Exception as e:
        return False, "", str(e)


def run_tool_setup(check_only: bool = False, install_all: bool = False):
    """Run the tool setup wizard."""
    console = DesignSystem.get_console()
    distro = detect_distro()

    # Ensure paths are configured for this session
    _ensure_path_configured()

    # Header
    DesignSystem.clear_screen()
    console.print()
    console.print("[bold cyan]╔══════════════════════════════════════════════════════════════╗[/bold cyan]")
    console.print("[bold cyan]║              SOULEYEZ TOOL SETUP WIZARD                      ║[/bold cyan]")
    console.print("[bold cyan]╚══════════════════════════════════════════════════════════════╝[/bold cyan]")
    console.print()

    # Distro detection
    distro_names = {
        'kali': 'Kali Linux',
        'parrot': 'Parrot OS',
        'ubuntu': 'Ubuntu',
        'debian': 'Debian',
        'unknown': 'Unknown Linux'
    }
    console.print(f"  Detected OS: [bold]{distro_names.get(distro, distro)}[/bold]")
    console.print()

    if distro in ('kali', 'parrot'):
        console.print("  [green]✓ You're on a pentesting distro![/green]")
        console.print("    All tools should be available via apt install.")
        console.print()
        _show_tool_status(console)

        if check_only:
            return

        missing = get_missing_tools(distro)
        if missing:
            console.print()
            if install_all or click.confirm("  Install missing tools via apt?", default=True):
                _install_apt_tools(console, missing)
        else:
            console.print("  [green]✓ All tools are installed![/green]")
        return

    # Ubuntu/Debian path
    if distro in ('ubuntu', 'debian'):
        console.print("  [yellow]Note:[/yellow] Some pentesting tools aren't in Ubuntu/Debian repos.")
        console.print("  This wizard will install them using pipx, go, snap, or from source.")
        console.print()

    _show_tool_status(console)

    if check_only:
        return

    missing = get_missing_tools(distro)
    if not missing:
        console.print()
        console.print("  [green]✓ All tools are installed![/green]")
        return

    console.print()
    console.print(f"  [yellow]⚠ {len(missing)} tools need installation[/yellow]")
    console.print()

    if not install_all:
        if not click.confirm("  Proceed with installation?", default=True):
            console.print("  Setup cancelled.")
            return

    # Check and install prerequisites first
    _check_prerequisites(console, missing, distro)

    # Group tools by install method for smarter ordering
    apt_tools = [t for t in missing if t['install_method'] == 'apt']
    pipx_tools = [t for t in missing if 'pipx' in t['install']]
    go_tools = [t for t in missing if 'go install' in t['install']]
    gem_tools = [t for t in missing if 'gem install' in t['install']]
    snap_tools = [t for t in missing if 'snap install' in t['install']]
    git_tools = [t for t in missing if 'git clone' in t['install']]
    other_tools = [t for t in missing if t not in apt_tools + pipx_tools + go_tools + gem_tools + snap_tools + git_tools]

    total_to_install = len(missing)
    installed_count = 0
    failed_tools = []

    # Install apt tools first (batch)
    if apt_tools:
        console.print()
        console.print("[bold cyan]  Installing apt packages...[/bold cyan]")
        success = _install_apt_tools(console, apt_tools)
        if success:
            installed_count += len(apt_tools)
        else:
            failed_tools.extend([t['name'] for t in apt_tools])

    # Install pipx tools
    if pipx_tools:
        console.print()
        console.print("[bold cyan]  Installing pipx packages...[/bold cyan]")
        for tool in pipx_tools:
            if _install_pipx_tool(console, tool):
                installed_count += 1
            else:
                failed_tools.append(tool['name'])

    # Install go tools
    if go_tools:
        console.print()
        console.print("[bold cyan]  Installing Go tools...[/bold cyan]")
        for tool in go_tools:
            if _install_go_tool(console, tool):
                installed_count += 1
            else:
                failed_tools.append(tool['name'])

    # Install gem tools
    if gem_tools:
        console.print()
        console.print("[bold cyan]  Installing Ruby gems...[/bold cyan]")
        for tool in gem_tools:
            if _install_gem_tool(console, tool):
                installed_count += 1
            else:
                failed_tools.append(tool['name'])

    # Install snap tools
    if snap_tools:
        console.print()
        console.print("[bold cyan]  Installing snap packages...[/bold cyan]")
        for tool in snap_tools:
            if _install_snap_tool(console, tool):
                installed_count += 1
            else:
                failed_tools.append(tool['name'])

    # Install git-based tools
    if git_tools:
        console.print()
        console.print("[bold cyan]  Installing from git...[/bold cyan]")
        for tool in git_tools:
            if _install_git_tool(console, tool):
                installed_count += 1
            else:
                failed_tools.append(tool['name'])

    # Install other tools (like metasploit)
    if other_tools:
        console.print()
        console.print("[bold cyan]  Installing additional tools...[/bold cyan]")
        for tool in other_tools:
            if _install_other_tool(console, tool):
                installed_count += 1
            else:
                failed_tools.append(tool['name'])

    # Configure PATH in bashrc
    _add_paths_to_bashrc()

    # Final status
    console.print()
    console.print("  " + "═" * 60)
    console.print("[bold cyan]  SETUP COMPLETE[/bold cyan]")
    console.print()

    if failed_tools:
        console.print(f"  [yellow]⚠ {len(failed_tools)} tools failed to install:[/yellow]")
        for name in failed_tools:
            console.print(f"    • {name}")
        console.print()

    # Refresh PATH and show final status
    _ensure_path_configured()
    _show_tool_status(console)

    # Remind about PATH for pipx/go tools
    if distro in ('ubuntu', 'debian'):
        console.print()
        console.print("  [yellow]Important:[/yellow] To use newly installed tools, either:")
        console.print("    1. Restart your terminal, OR")
        console.print("    2. Run: [cyan]source ~/.bashrc[/cyan]")
        console.print()


def _show_tool_status(console):
    """Display current tool installation status."""
    tools_by_cat = get_tools_by_category()
    installed, total = get_tool_stats()

    status_color = "green" if installed == total else "yellow" if installed > 0 else "red"
    console.print(f"  [bold]Tool Status:[/bold] [{status_color}]{installed}/{total} installed[/{status_color}]")
    console.print()

    for category, tools in tools_by_cat.items():
        cat_name = get_category_name(category)
        console.print(f"  [bold]{cat_name}[/bold]")

        for tool in tools:
            if tool['installed']:
                status = "[green]✓[/green]"
            else:
                status = "[red]✗[/red]"

            console.print(f"    {status} {tool['name']:<18} - {tool['description'][:40]}")

        console.print()


def _check_prerequisites(console, missing_tools: List[Dict], distro: str):
    """Check and install prerequisites needed for tool installation."""
    needed_prereqs = set()

    for tool in missing_tools:
        install_cmd = tool['install']
        if 'pipx' in install_cmd:
            needed_prereqs.add('pipx')
            # pipx tools with native extensions need build dependencies
            needed_prereqs.add('build-deps')
        if 'go install' in install_cmd:
            needed_prereqs.add('golang')
        if 'gem install' in install_cmd:
            needed_prereqs.add('ruby')
        if 'snap install' in install_cmd:
            needed_prereqs.add('snap')
        if 'git clone' in install_cmd:
            needed_prereqs.add('git')

    missing_prereqs = []
    for prereq in needed_prereqs:
        info = PREREQUISITES[prereq]
        # Check if always_install flag is set, or if tool check fails
        if info.get('always_install') or (info.get('check') and not check_tool(info['check'])):
            missing_prereqs.append((prereq, info))

    if not missing_prereqs:
        return

    # Sort to ensure build-deps comes first (needed before pipx installs)
    prereq_order = {'build-deps': 0, 'pipx': 1, 'golang': 2, 'ruby': 3, 'snap': 4, 'git': 5}
    missing_prereqs.sort(key=lambda x: prereq_order.get(x[0], 99))

    console.print()
    console.print("[bold cyan]  Installing prerequisites...[/bold cyan]")

    for prereq, info in missing_prereqs:
        console.print(f"    Installing {prereq}...", end=" ")
        success, _, stderr = _run_command(info['install'], console, capture=True)
        if success:
            console.print("[green]✓[/green]")

            # Run pipx ensurepath if we just installed pipx
            if prereq == 'pipx':
                subprocess.run("pipx ensurepath", shell=True, capture_output=True)
        else:
            console.print(f"[red]✗[/red] {stderr[:50]}")

    # Update PATH after installing prerequisites
    _ensure_path_configured()


def _install_apt_tools(console, tools: List[Dict]) -> bool:
    """Install tools available via apt."""
    packages = []
    for tool in tools:
        cmd = tool['install']
        if 'apt install' in cmd:
            # Extract package name
            parts = cmd.split()
            for i, part in enumerate(parts):
                if part == 'install' and i + 1 < len(parts):
                    pkg = parts[i + 1]
                    if not pkg.startswith('-'):
                        packages.append(pkg)
                    break

    if not packages:
        return True

    console.print(f"    Packages: {', '.join(packages)}")
    cmd = f"sudo apt install -y {' '.join(packages)}"
    success, _, _ = _run_command(cmd, console)

    if success:
        console.print("    [green]✓ apt packages installed[/green]")
    else:
        console.print("    [yellow]⚠ Some packages may have failed[/yellow]")

    return success


def _install_pipx_tool(console, tool: Dict) -> bool:
    """Install a tool using pipx."""
    name = tool['name']
    cmd = tool['install']

    console.print(f"    {name}...", end=" ")

    # pipx install commands
    success, _, stderr = _run_command(cmd, console, capture=True)

    if success:
        console.print("[green]✓[/green]")
        return True
    else:
        # Check if already installed
        if "already exists" in stderr or "already installed" in stderr.lower():
            console.print("[green]✓ (already installed)[/green]")
            return True
        console.print(f"[red]✗[/red]")
        if stderr:
            console.print(f"      [dim]{stderr[:80]}[/dim]")
        return False


def _install_go_tool(console, tool: Dict) -> bool:
    """Install a tool using go install."""
    name = tool['name']
    cmd = tool['install']

    console.print(f"    {name}...", end=" ")

    # Ensure GOPATH is set
    gopath = os.path.expanduser("~/go")
    env = os.environ.copy()
    env["GOPATH"] = gopath
    env["PATH"] = f"{gopath}/bin:" + env.get("PATH", "")

    try:
        result = subprocess.run(
            cmd,
            shell=True,
            capture_output=True,
            text=True,
            env=env,
            timeout=300  # 5 minute timeout for go installs
        )
        if result.returncode == 0:
            console.print("[green]✓[/green]")
            return True
        else:
            console.print("[red]✗[/red]")
            if result.stderr:
                console.print(f"      [dim]{result.stderr[:80]}[/dim]")
            return False
    except Exception as e:
        console.print(f"[red]✗[/red] {str(e)[:50]}")
        return False


def _install_gem_tool(console, tool: Dict) -> bool:
    """Install a tool using gem."""
    name = tool['name']
    cmd = tool['install']

    console.print(f"    {name}...", end=" ")

    success, _, stderr = _run_command(cmd, console, capture=True)

    if success:
        console.print("[green]✓[/green]")
        return True
    else:
        console.print(f"[red]✗[/red]")
        if stderr:
            console.print(f"      [dim]{stderr[:80]}[/dim]")
        return False


def _install_snap_tool(console, tool: Dict) -> bool:
    """Install a tool using snap."""
    name = tool['name']
    cmd = tool['install']

    console.print(f"    {name}...", end=" ")

    success, _, stderr = _run_command(cmd, console, capture=True)

    if success:
        console.print("[green]✓[/green]")
        return True
    else:
        console.print(f"[red]✗[/red]")
        if stderr:
            console.print(f"      [dim]{stderr[:80]}[/dim]")
        return False


def _install_git_tool(console, tool: Dict) -> bool:
    """Install a tool from git (requires sudo for /opt)."""
    name = tool['name']
    cmd = tool['install']

    console.print(f"    {name}...", end=" ")

    # These typically clone to /opt and need sudo
    success, _, stderr = _run_command(cmd, console, capture=True)

    if success:
        console.print("[green]✓[/green]")
        return True
    else:
        # Check if directory already exists
        if "already exists" in stderr:
            console.print("[green]✓ (already installed)[/green]")
            return True
        console.print(f"[red]✗[/red]")
        if stderr:
            console.print(f"      [dim]{stderr[:80]}[/dim]")
        return False


def _install_other_tool(console, tool: Dict) -> bool:
    """Install tools with custom installation methods (like metasploit)."""
    name = tool['name']
    cmd = tool['install']

    console.print(f"    {name}...", end=" ")

    # Special handling for metasploit installer (takes a while)
    if 'msfinstall' in cmd:
        console.print()
        console.print("      [yellow]Installing Metasploit (this may take several minutes)...[/yellow]")
        console.print("      [dim]Installing postgresql and downloading Metasploit installer...[/dim]")

        # Run the full install command (includes postgresql setup)
        try:
            result = subprocess.run(
                cmd,
                shell=True,
                timeout=1800  # 30 minute timeout for metasploit
            )
            if result.returncode == 0:
                console.print("      [green]✓ Metasploit installed[/green]")

                # Initialize MSF database (must run as normal user, not root)
                console.print("      [dim]Initializing MSF database...[/dim]")
                try:
                    # Use full path since /opt/metasploit-framework/bin isn't in PATH yet
                    msfdb_path = '/opt/metasploit-framework/bin/msfdb'
                    init_result = subprocess.run(
                        [msfdb_path, 'init'],
                        capture_output=True,
                        timeout=300  # 5 minute timeout
                    )
                    if init_result.returncode == 0:
                        console.print("      [green]✓ MSF database initialized[/green]")

                        # Copy database.yml to /root/.msf4/ so sudo msfconsole can connect
                        # msfdb init creates config in ~/.msf4/, but sudo looks in /root/.msf4/
                        import os
                        from pathlib import Path
                        user_msf4 = Path.home() / '.msf4'
                        user_db_yml = user_msf4 / 'database.yml'
                        if user_db_yml.exists():
                            console.print("      [dim]Configuring sudo access to MSF database...[/dim]")
                            try:
                                # Create /root/.msf4/ directory and copy database.yml
                                copy_result = subprocess.run(
                                    ['sudo', 'mkdir', '-p', '/root/.msf4'],
                                    capture_output=True
                                )
                                copy_result = subprocess.run(
                                    ['sudo', 'cp', str(user_db_yml), '/root/.msf4/database.yml'],
                                    capture_output=True
                                )
                                if copy_result.returncode == 0:
                                    console.print("      [green]✓ Sudo MSF access configured[/green]")
                                else:
                                    console.print("      [yellow]⚠ Could not configure sudo access (run manually: sudo cp ~/.msf4/database.yml /root/.msf4/)[/yellow]")
                            except Exception as e:
                                console.print(f"      [yellow]⚠ Could not configure sudo access: {str(e)[:30]}[/yellow]")
                    else:
                        console.print("      [yellow]⚠ MSF database init returned non-zero (may already be initialized)[/yellow]")
                except Exception as e:
                    console.print(f"      [yellow]⚠ Could not init MSF database: {str(e)[:40]}[/yellow]")

                return True
            else:
                console.print("      [red]✗ Metasploit installation failed[/red]")
                return False
        except subprocess.TimeoutExpired:
            console.print("      [red]✗ Installation timed out[/red]")
            return False
        except Exception as e:
            console.print(f"      [red]✗ Error: {str(e)[:50]}[/red]")
            return False

    success, _, stderr = _run_command(cmd, console, capture=True)

    if success:
        console.print("[green]✓[/green]")
        return True
    else:
        console.print(f"[red]✗[/red]")
        if stderr:
            console.print(f"      [dim]{stderr[:80]}[/dim]")
        return False
