#!/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 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...")
    status, output = run_command(["ansible-galaxy", "collection", "list"], shell=False)
    
    if not status:
        printfe("red", "Failed to list Ansible collections")
        return False
    
    # Check each required collection and install if missing
    for collection in required_collections:
        if collection not in output:
            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")
    
    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 both
    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
    
    # 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"))
        
        # Check if ansible is installed
        status, _ = run_command(["which", "ansible-playbook"], shell=False)
        if not status:
            printfe("yellow", "Ansible is not installed, installing it with pipx...")
            status, output = run_command(["pipx", "install", "--include-deps", "ansible", "ansible-lint"], shell=False)
            if not status:
                printfe("red", f"Failed to install Ansible: {output}")
                return 1
        
        # 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 = [
            "ansible-playbook",
            "-i", f"{dotfiles_path}/config/ansible/inventory.ini",
            f"{dotfiles_path}/config/ansible/main.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())