#!/usr/bin/env python3
"""Gentoo installation orchestrator for Legion S7 15ACH6.

Usage:
    python setup.py [command]
    python setup.py --install  # Full install up to first reboot
    python setup.py --desktop  # Post-reboot desktop setup

Major Phases:
    --install   Full installation to first reboot (auto-detects chroot state)
                Outside chroot: disk -> stage3 -> config -> fstab -> chroot
                Inside chroot:  sync -> world -> firmware -> kernel -> services
                                -> users -> nvidia -> bootloader
    --desktop   Post-reboot desktop setup (Hyprland + optional fingerprint)

Individual Commands:
    disk        Partition, encrypt, and mount disks
    stage3      Download and extract stage3 tarball
    chroot      Prepare and enter chroot environment
    config      Copy portage and system configuration
    fstab       Generate fstab, crypttab, and dmcrypt config
    sync        Initial portage sync, profile, and overlays (inside chroot)
    world       Update @world (inside chroot)
    firmware    Install linux-firmware (inside chroot)
    kernel      Install kernel (inside chroot)
    services    Install and configure system services (inside chroot)
    users       Configure shell, sudo, and create user (inside chroot)
    nvidia      Install NVIDIA drivers and configure (inside chroot)
    bootloader  Install GRUB and verify initramfs (inside chroot)
    fingerprint Set up fingerprint authentication (post-install)
    swap-on     Activate encrypted swap (inside chroot)
    swap-off    Deactivate encrypted swap (inside chroot)
    all         Run pre-chroot steps (disk -> chroot)

Examples:
    python setup.py --install  # Run full install phase (pre-reboot)
    python setup.py --desktop  # Run desktop setup (post-reboot)
    python setup.py disk       # Just partition disks
    python setup.py            # Interactive menu
"""

import argparse
import sys
from pathlib import Path

from install.utils import info, success, error, fatal, check_root, check_uefi, prompt, emerge
from install.disk import DiskConfig, DiskLayout, prepare_disk
from install.stage3 import Stage3Config, fetch_stage3
from install.chroot import prepare_chroot, enter_chroot
from install.fstab import generate_all as generate_fstab_all
from install.portage import copy_portage_config
from install.sync import initial_sync, activate_swap, deactivate_swap
from install.services import setup_services
from install.users import setup_users
from install.nvidia import setup_nvidia
from install.bootloader import setup_bootloader


# Path to config files (same directory as this script)
SCRIPT_DIR = Path(__file__).parent.resolve()
CONFIG_DIR = SCRIPT_DIR  # Config files live alongside setup.py


# Global state (set during disk preparation)
_disk_layout: DiskLayout | None = None
_disk_config: DiskConfig | None = None


def cmd_disk() -> tuple[DiskLayout, DiskConfig]:
    """Partition, encrypt, and mount disks."""
    global _disk_layout, _disk_config

    check_uefi()
    layout, config = prepare_disk()

    _disk_layout = layout
    _disk_config = config

    return layout, config


def cmd_stage3() -> None:
    """Download and extract stage3."""
    config = Stage3Config(
        init_system="openrc",
        mount_root=Path("/mnt/gentoo"),
    )
    fetch_stage3(config)


def cmd_config() -> None:
    """Copy portage configuration."""
    copy_portage_config(
        source_dir=CONFIG_DIR,
        mount_root=Path("/mnt/gentoo"),
    )


def cmd_fstab(layout: DiskLayout | None = None, config: DiskConfig | None = None) -> None:
    """Generate filesystem configuration files."""
    if layout is None or config is None:
        fatal("Disk layout not available. Run 'disk' command first.")

    generate_fstab_all(layout, config)


def cmd_chroot() -> None:
    """Prepare and enter chroot."""
    prepare_chroot()

    print()
    response = prompt("Enter chroot now? (y/n): ")
    if response.lower() == "y":
        enter_chroot()


def cmd_sync() -> None:
    """Initial portage sync and profile setup (run inside chroot)."""
    initial_sync()


def cmd_swap_on() -> None:
    """Activate encrypted swap (inside chroot)."""
    activate_swap()


def cmd_swap_off() -> None:
    """Deactivate encrypted swap (inside chroot)."""
    deactivate_swap()


def cmd_world() -> None:
    """Update @world (inside chroot)."""
    info("=== Updating @world ===")
    emerge("@world", extra_args=["--update", "--deep", "--newuse"])
    success("@world update complete.")

    # Handle circular dependencies - rebuild with full USE flags
    circ_deps_file = Path("/etc/portage/package.use/circular-dependencies")
    if circ_deps_file.exists():
        packages = parse_circular_deps(circ_deps_file)
        if packages:
            print()
            info("=== Rebuilding Circular Dependency Packages ===")
            info(f"Packages: {', '.join(packages)}")

            # Clear the file (keep header)
            clear_circular_deps(circ_deps_file)

            # Rebuild with full USE flags
            emerge(*packages, extra_args=["--oneshot"], ask=False)
            success("Circular dependency packages rebuilt with full USE flags.")


def parse_circular_deps(filepath: Path) -> list[str]:
    """Parse package names from circular-dependencies file."""
    packages = []
    for line in filepath.read_text().splitlines():
        line = line.strip()
        # Skip comments and empty lines
        if not line or line.startswith("#"):
            continue
        # Extract package name (first token before USE flags)
        parts = line.split()
        if parts:
            packages.append(parts[0])
    return packages


def clear_circular_deps(filepath: Path) -> None:
    """Clear circular-dependencies file, keeping header comments."""
    lines = filepath.read_text().splitlines()
    header = []
    for line in lines:
        # Keep comments and blank lines in header
        if line.startswith("#") or line.strip() == "":
            header.append(line)
        else:
            # Stop at first package line
            break
    # Write header, ensure trailing newline
    if header:
        filepath.write_text("\n".join(header) + "\n")
    else:
        filepath.write_text("")


def cmd_firmware() -> None:
    """Install linux-firmware (inside chroot)."""
    info("=== Installing Linux Firmware ===")
    emerge("sys-kernel/linux-firmware")
    success("Linux firmware installed.")


def cmd_kernel() -> None:
    """Install kernel (inside chroot)."""
    info("=== Installing Kernel ===")
    print()
    info("Options:")
    print("  1) gentoo-kernel-bin  - Precompiled, fastest")
    print("  2) gentoo-kernel      - Compiled locally with your USE flags")
    print()

    choice = prompt("Select kernel [1]: ").strip()

    if choice == "2":
        emerge("sys-kernel/gentoo-kernel")
    else:
        emerge("sys-kernel/gentoo-kernel-bin")

    success("Kernel installed.")
    print()
    info("Kernel modules installed. Dracut will generate initramfs.")


def cmd_services() -> None:
    """Install and configure services (inside chroot)."""
    setup_services(source_dir=CONFIG_DIR)


def cmd_users() -> None:
    """Configure shell, sudo, and create user (inside chroot)."""
    setup_users(source_dir=CONFIG_DIR)


def cmd_nvidia() -> None:
    """Install NVIDIA drivers and configure (inside chroot)."""
    setup_nvidia(source_dir=CONFIG_DIR)


def cmd_bootloader() -> None:
    """Install GRUB and verify initramfs (inside chroot)."""
    setup_bootloader(source_dir=CONFIG_DIR)


def cmd_fingerprint() -> None:
    """Set up fingerprint authentication (post-install)."""
    from install.fingerprint import setup_fingerprint
    setup_fingerprint(source_dir=CONFIG_DIR)


def is_in_chroot() -> bool:
    """Detect if we're running inside chroot vs live environment."""
    # If /mnt/gentoo is mounted, we're in the live environment
    # If not, we're either in chroot or on the installed system
    from install.utils import is_mounted
    return not is_mounted(Path("/mnt/gentoo"))


def cmd_install() -> None:
    """Full installation up to first reboot."""
    if is_in_chroot():
        # Inside chroot: run sync through bootloader
        info("=== Install Phase (inside chroot) ===")
        print()
        info("Running: sync -> world -> firmware -> kernel -> services -> users -> nvidia -> bootloader")
        print()

        cmd_sync()
        print()
        cmd_world()
        print()
        cmd_firmware()
        print()
        cmd_kernel()
        print()
        cmd_services()
        print()
        cmd_users()
        print()
        cmd_nvidia()
        print()
        cmd_bootloader()
        print()

        success("=== Install phase complete ===")
        print()
        print("Next steps:")
        print("  1. Exit chroot: exit")
        print("  2. Unmount: umount -R /mnt/gentoo")
        print("  3. Reboot: reboot")
        print("  4. After reboot, run: python setup.py --desktop")
    else:
        # Outside chroot: run disk through chroot, then prompt
        info("=== Install Phase (pre-chroot) ===")
        print()

        layout, config = cmd_disk()
        print()
        cmd_stage3()
        print()
        cmd_config()
        print()
        cmd_fstab(layout, config)
        print()
        cmd_chroot()

        # After chroot exits, remind user
        print()
        success("=== Pre-chroot phase complete ===")
        print()
        print("You exited the chroot. To continue installation:")
        print("  1. Re-enter chroot: chroot /mnt/gentoo /bin/bash")
        print("  2. Source profile: source /etc/profile && export PS1='(chroot) $PS1'")
        print("  3. Run: python /root/gentoo/v4/setup.py --install")


def cmd_desktop() -> None:
    """Post-reboot desktop setup."""
    info("=== Desktop Setup (post-reboot) ===")
    print()

    # Install Hyprland desktop
    info("Installing Hyprland desktop environment...")
    emerge("@hyprland", extra_args=["--update", "--deep", "--newuse"])
    success("Hyprland desktop installed.")
    print()

    # Offer fingerprint setup
    response = prompt("Set up fingerprint authentication? [y/N]: ").strip().lower()
    if response in ("y", "yes"):
        print()
        cmd_fingerprint()
    else:
        info("Skipped fingerprint setup. Run 'python setup.py fingerprint' later if needed.")

    print()
    success("=== Desktop setup complete ===")
    print()
    print("Next steps:")
    print("  1. Reboot to apply display manager")
    print("  2. Log in via SDDM")
    print("  3. Copy Hyprland configs:")
    print("     cp ~/gentoo/v4/hypr/ENVariables.conf ~/.config/hypr/UserConfigs/")
    print("     cp ~/gentoo/v4/hypr/monitors.conf ~/.config/hypr/")


def cmd_all() -> None:
    """Full installation workflow."""
    info("=== Gentoo Full Installation ===")
    print()

    # Step 1: Disk preparation
    info("Step 1: Disk Preparation")
    layout, config = cmd_disk()

    # Step 2: Stage3
    print()
    info("Step 2: Stage3 Download & Extract")
    cmd_stage3()

    # Step 3: Configuration
    print()
    info("Step 3: Copy Configuration")
    cmd_config()

    # Step 4: Generate fstab/crypttab
    print()
    info("Step 4: Generate Filesystem Config")
    cmd_fstab(layout, config)

    # Step 5: Chroot
    print()
    info("Step 5: Chroot Preparation")
    cmd_chroot()


def interactive_menu() -> None:
    """Interactive command selection."""
    print()
    info("=== Gentoo Installer for Legion S7 15ACH6 ===")
    print()
    print("Commands (pre-chroot):")
    print("  1) disk     - Partition, encrypt, and mount disks")
    print("  2) stage3   - Download and extract stage3")
    print("  3) config   - Copy portage configuration")
    print("  4) fstab    - Generate fstab/crypttab")
    print("  5) chroot   - Prepare and enter chroot")
    print("  6) all      - Full pre-chroot installation")
    print()
    print("Commands (inside chroot):")
    print("  7) sync     - Portage sync, profile, overlays (activates swap)")
    print("  8) world    - emerge -avuDN @world")
    print("  9) firmware - Install linux-firmware")
    print(" 10) kernel   - Install kernel")
    print(" 11) services - Install and configure services")
    print(" 12) users    - Configure shell, sudo, create user")
    print(" 13) nvidia   - Install NVIDIA drivers")
    print(" 14) bootloader - Install GRUB, verify initramfs")
    print()
    print("Post-install:")
    print(" 15) fingerprint - Set up fingerprint authentication")
    print()
    print("Utilities:")
    print(" 16) swap-on  - Activate encrypted swap")
    print(" 17) swap-off - Deactivate encrypted swap")
    print()
    print("Full workflows:")
    print("  i) install - Full install to first reboot (auto-detects chroot)")
    print("  d) desktop - Post-reboot desktop setup (Hyprland + fingerprint)")
    print()
    print("  q) quit")
    print()

    choice = prompt("Select command: ").strip().lower()

    if choice in ("1", "disk"):
        cmd_disk()
    elif choice in ("2", "stage3"):
        cmd_stage3()
    elif choice in ("3", "config"):
        cmd_config()
    elif choice in ("4", "fstab"):
        if _disk_layout is None or _disk_config is None:
            error("Disk layout not available.")
            error("Run 'disk' command first in this session, or use 'all' for full installation.")
            return
        cmd_fstab(_disk_layout, _disk_config)
    elif choice in ("5", "chroot"):
        cmd_chroot()
    elif choice in ("6", "all"):
        cmd_all()
    elif choice in ("7", "sync"):
        cmd_sync()
    elif choice in ("8", "world"):
        cmd_world()
    elif choice in ("9", "firmware"):
        cmd_firmware()
    elif choice in ("10", "kernel"):
        cmd_kernel()
    elif choice in ("11", "services"):
        cmd_services()
    elif choice in ("12", "users"):
        cmd_users()
    elif choice in ("13", "nvidia"):
        cmd_nvidia()
    elif choice in ("14", "bootloader"):
        cmd_bootloader()
    elif choice in ("15", "fingerprint"):
        cmd_fingerprint()
    elif choice in ("16", "swap-on"):
        cmd_swap_on()
    elif choice in ("17", "swap-off"):
        cmd_swap_off()
    elif choice in ("i", "install"):
        cmd_install()
    elif choice in ("d", "desktop"):
        cmd_desktop()
    elif choice == "q":
        sys.exit(0)
    else:
        error("Invalid choice")
        interactive_menu()


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Gentoo installation orchestrator",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        "command",
        nargs="?",
        choices=["disk", "stage3", "chroot", "config", "fstab", "sync", "world", "firmware", "kernel", "services", "users", "nvidia", "bootloader", "fingerprint", "swap-on", "swap-off", "all"],
        help="Command to run (interactive menu if not specified)",
    )
    parser.add_argument(
        "--install",
        action="store_true",
        help="Full installation up to first reboot (auto-detects chroot state)",
    )
    parser.add_argument(
        "--desktop",
        action="store_true",
        help="Post-reboot desktop setup (Hyprland + fingerprint)",
    )

    args = parser.parse_args()

    check_root()

    # Handle major phase flags first
    if args.install:
        cmd_install()
    elif args.desktop:
        cmd_desktop()
    elif args.command is None:
        interactive_menu()
    elif args.command == "disk":
        cmd_disk()
    elif args.command == "stage3":
        cmd_stage3()
    elif args.command == "config":
        cmd_config()
    elif args.command == "fstab":
        # For standalone fstab, we need disk info
        error("Run 'disk' first, or use 'all' for full installation")
        sys.exit(1)
    elif args.command == "chroot":
        cmd_chroot()
    elif args.command == "sync":
        cmd_sync()
    elif args.command == "world":
        cmd_world()
    elif args.command == "firmware":
        cmd_firmware()
    elif args.command == "kernel":
        cmd_kernel()
    elif args.command == "services":
        cmd_services()
    elif args.command == "users":
        cmd_users()
    elif args.command == "nvidia":
        cmd_nvidia()
    elif args.command == "bootloader":
        cmd_bootloader()
    elif args.command == "fingerprint":
        cmd_fingerprint()
    elif args.command == "swap-on":
        cmd_swap_on()
    elif args.command == "swap-off":
        cmd_swap_off()
    elif args.command == "all":
        cmd_all()


if __name__ == "__main__":
    main()