#!/usr/bin/env python3

import os
import sys
import subprocess
import argparse

# Import helper functions
sys.path.append(os.path.join(os.path.expanduser("~/.dotfiles"), "bin"))
from helpers.functions import printfe, run_command


def help_message():
    """Print help message and exit"""
    printfe("green", "Usage: upgrade.py [options]")
    printfe("green", "Options:")
    printfe("green", "  --ha, -H          Upgrade Home Manager packages.")
    printfe("green", "  --ansible, -A     Upgrade Ansible packages.")
    printfe(
        "green",
        "  --ansible-verbose Upgrade Ansible packages with verbose output. (-vvv)",
    )
    printfe(
        "green",
        "  --full-speed, -F  Upgrade packages and use all available cores for compilation. (Default: 8 cores)",
    )
    printfe("green", "  --help, -h        Display this help message.")
    return 0


def check_git_repository():
    """Check for changes in the dotfiles git repository and prompt user to pull if needed"""
    dotfiles_path = os.environ.get("DOTFILES_PATH", os.path.expanduser("~/.dotfiles"))

    printfe("cyan", "Checking for updates in dotfiles repository...")

    # Change to dotfiles directory
    current_dir = os.getcwd()
    os.chdir(dotfiles_path)

    # Check if this is a git repository
    status, _ = run_command(["git", "rev-parse", "--is-inside-work-tree"], shell=False)
    if not status:
        printfe("red", "The dotfiles directory is not a git repository.")
        os.chdir(current_dir)
        return False

    # Get the current branch name
    status, current_branch = run_command(
        ["git", "rev-parse", "--abbrev-ref", "HEAD"], shell=False
    )
    if not status:
        printfe("red", "Failed to determine current branch.")
        os.chdir(current_dir)
        return False

    current_branch = current_branch.strip()

    # Fetch the latest changes
    status, output = run_command(["git", "fetch"], shell=False)
    if not status:
        printfe("yellow", f"Warning: Failed to fetch changes from git repository: {output}")
        printfe("yellow", "Continuing update process without repository check...")
        os.chdir(current_dir)
        return True

    # Check if remote branch exists
    status, output = run_command(
        ["git", "ls-remote", "--heads", "origin", current_branch], shell=False
    )
    if not status or not output.strip():
        printfe(
            "yellow",
            f"Remote branch 'origin/{current_branch}' not found. Using local branch only.",
        )
        os.chdir(current_dir)
        return True

    # Check if we're behind the remote
    status, output = run_command(
        ["git", "rev-list", f"HEAD..origin/{current_branch}", "--count"], shell=False
    )
    if not status:
        printfe("red", f"Failed to check for repository updates: {output}")
        os.chdir(current_dir)
        return False

    behind_count = output.strip()
    if behind_count == "0":
        printfe(
            "green", f"Dotfiles repository is up to date on branch '{current_branch}'."
        )
        os.chdir(current_dir)
        return True

    # Show what changes are available
    status, output = run_command(
        ["git", "log", f"HEAD..origin/{current_branch}", "--oneline"], shell=False
    )
    if status:
        printfe(
            "yellow",
            f"Your dotfiles repository is {behind_count} commit(s) behind on branch '{current_branch}'. Changes:",
        )
        for line in output.strip().splitlines():
            printfe("yellow", f"  • {line}")
    else:
        printfe(
            "yellow",
            f"Your dotfiles repository is {behind_count} commit(s) behind on branch '{current_branch}'.",
        )

    # Ask user if they want to pull changes
    response = input("Do you want to pull these changes? (yes/no): ").strip().lower()

    if response in ["yes", "y"]:
        status, output = run_command(
            ["git", "pull", "origin", current_branch], shell=False
        )
        if not status:
            printfe("red", f"Failed to pull changes: {output}")
            os.chdir(current_dir)
            return False
        printfe("green", "Successfully updated dotfiles repository.")
    else:
        printfe("yellow", "Skipping repository update.")

    os.chdir(current_dir)
    return True


def ensure_ansible_collections():
    """Ensure required Ansible collections are installed"""
    # List of required collections that can be expanded in the future
    required_collections = [
        "community.general",
    ]

    printfe("cyan", "Checking for required Ansible collections...")

    # Get list of installed collections using ansible-galaxy
    status, output = run_command(["ansible-galaxy", "collection", "list"], shell=False)
    if not status:
        printfe("yellow", f"Failed to list Ansible collections: {output}")
        printfe("yellow", "Will try to install all required collections.")
        installed_collections = []
    else:
        # Parse output to get installed collections
        installed_collections = []

        # Split output into lines and process
        lines = output.splitlines()
        collection_section = False

        for line in lines:
            line = line.strip()

            # Skip empty lines
            if not line:
                continue

            # Check if we've reached the collection listing section
            if line.startswith("Collection"):
                collection_section = True
                continue

            # Skip the separator line after the header
            if collection_section and line.startswith("--"):
                continue

            # Process collection entries
            if collection_section and " " in line:
                # Format is typically: "community.general 10.4.0"
                parts = line.split()
                if len(parts) >= 1:
                    collection_name = parts[0]
                    installed_collections.append(collection_name)

    # Check which required collections are missing
    missing_collections = []
    for collection in required_collections:
        if collection not in installed_collections:
            missing_collections.append(collection)

    # Install missing collections
    if missing_collections:
        for collection in missing_collections:
            printfe("yellow", f"Installing {collection} collection...")
            status, install_output = run_command(
                ["ansible-galaxy", "collection", "install", collection], shell=False
            )
            if not status:
                printfe(
                    "yellow",
                    f"Warning: Failed to install {collection} collection: {install_output}",
                )
                printfe(
                    "yellow",
                    f"Continuing anyway, but playbook might fail if it requires {collection}",
                )
            else:
                printfe("green", f"Successfully installed {collection} collection")
    else:
        printfe("green", "All required collections are already installed.")

    return True


def main():
    # Parse arguments
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument(
        "--ha", "-H", action="store_true", help="Upgrade Home Manager packages"
    )
    parser.add_argument(
        "--ansible", "-A", action="store_true", help="Upgrade Ansible packages"
    )
    parser.add_argument(
        "--ansible-verbose",
        action="store_true",
        help="Upgrade Ansible packages with verbose output",
    )
    parser.add_argument(
        "--full-speed", "-F", action="store_true", help="Use all available cores"
    )
    parser.add_argument(
        "--help", "-h", action="store_true", help="Display help message"
    )

    args = parser.parse_args()

    if args.help:
        return help_message()

    # If no specific option provided, run all
    if not args.ha and not args.ansible and not args.ansible_verbose:
        args.ha = True
        args.ansible = True

    # If ansible_verbose is set, also set ansible
    if args.ansible_verbose:
        args.ansible = True

    # Always check git repository first
    if not check_git_repository():
        printfe("red", "Failed to check or update dotfiles repository.")
        return 1

    # Set cores and jobs based on full-speed flag
    if args.full_speed:
        import multiprocessing

        cores = jobs = multiprocessing.cpu_count()
    else:
        cores = 8
        jobs = 1

    printfe("cyan", f"Limiting to {cores} cores with {jobs} jobs.")

    # Home Manager update
    if args.ha:
        dotfiles_path = os.environ.get(
            "DOTFILES_PATH", os.path.expanduser("~/.dotfiles")
        )
        hostname = os.uname().nodename

        printfe("cyan", "Updating Home Manager flake...")
        os.chdir(f"{dotfiles_path}/config/home-manager")
        status, output = run_command(
            [
                "nix",
                "--extra-experimental-features",
                "nix-command",
                "--extra-experimental-features",
                "flakes",
                "flake",
                "update",
            ],
            shell=False,
        )
        if not status:
            printfe("red", f"Failed to update Home Manager flake: {output}")
            return 1

        # Check if home-manager is installed
        status, _ = run_command(["which", "home-manager"], shell=False)
        if status:
            printfe("cyan", "Cleaning old backup files...")
            backup_file = os.path.expanduser("~/.config/mimeapps.list.backup")
            if os.path.exists(backup_file):
                os.remove(backup_file)

            printfe("cyan", "Upgrading Home Manager packages...")
            env = os.environ.copy()
            env["NIXPKGS_ALLOW_UNFREE"] = "1"

            cmd = [
                "home-manager",
                "--extra-experimental-features",
                "nix-command",
                "--extra-experimental-features",
                "flakes",
                "switch",
                "-b",
                "backup",
                f"--flake",
                f".#{hostname}",
                "--impure",
                "--cores",
                str(cores),
                "-j",
                str(jobs),
            ]

            result = subprocess.run(cmd, env=env)
            if result.returncode != 0:
                printfe("red", "Failed to upgrade Home Manager packages.")
                return 1
        else:
            printfe("red", "Home Manager is not installed.")
            return 1

    # Ansible update
    if args.ansible:
        dotfiles_path = os.environ.get(
            "DOTFILES_PATH", os.path.expanduser("~/.dotfiles")
        )
        hostname = os.uname().nodename
        username = os.environ.get("USER", os.environ.get("USERNAME", "user"))

        # Ensure required collections are installed
        if not ensure_ansible_collections():
            printfe(
                "red", "Failed to ensure required Ansible collections are installed"
            )
            return 1

        printfe("cyan", "Running Ansible playbook...")
        ansible_cmd = [
            "/usr/bin/env",
            "ansible-playbook",
            "-i",
            f"{dotfiles_path}/config/ansible/inventory.ini",
            f"{dotfiles_path}/config/ansible/playbook.yml",
            "--extra-vars",
            f"hostname={hostname}",
            "--extra-vars",
            f"ansible_user={username}",
            "--limit",
            hostname,
            "--ask-become-pass",
        ]

        if args.ansible_verbose:
            ansible_cmd.append("-vvv")

        result = subprocess.run(ansible_cmd)
        if result.returncode != 0:
            printfe("red", "Failed to upgrade Ansible packages.")
            return 1

    return 0


if __name__ == "__main__":
    sys.exit(main())