import argparse
import os
import sys

import apt
import apt_pkg


def get_manual_installed_packages(cache):
    manual_packages = []
    for pkg in cache:
        if pkg.is_installed and not pkg.is_auto_installed:
            manual_packages.append(pkg.name)
    return manual_packages


def check_root_privileges():
    if os.geteuid() != 0:
        print(
            "This script needs root privileges to apply changes. Please run with sudo."
        )
        sys.exit(1)


def identify_markable_packages(cache, dep_cache, manual_packages, raw=False):
    successfully_markable = []
    reverted_changes = []

    for pkg_name in manual_packages:
        pkg = cache.get(pkg_name)
        if not pkg or not pkg.is_installed:
            continue

        dep_cache.mark_auto(pkg._pkg, True)

        if dep_cache.is_garbage(pkg._pkg):
            dep_cache.mark_auto(pkg._pkg, False)
            reverted_changes.append(pkg_name)
        else:
            successfully_markable.append(pkg_name)

    return successfully_markable, reverted_changes


def main():
    parser = argparse.ArgumentParser(
        description="Suggests and optionally applies marking manually installed packages as automatic."
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="Apply the suggested changes (mark packages as automatic). Requires root privileges.",
    )
    parser.add_argument(
        "--raw",
        "-r",
        action="store_true",
        help="Output only the list of packages -- suitable for scripting.",
    )
    parser.add_argument(
        "-o",
        "--apt-option",
        action="append",
        help="Set an APT configuration option (e.g., -o APT::Get::Assume-Yes=true). Can be used multiple times.",
    )
    args = parser.parse_args()

    if args.apply:
        check_root_privileges()

    apt_pkg.init_system()
    if args.apt_option:
        for option_str in args.apt_option:
            try:
                key, value = option_str.split("=", 1)
                apt_pkg.Configuration().set(key, value)
            except ValueError:
                print(
                    f"Warning: Invalid APT option format: {option_str}. Expected 'Key=Value'."
                )
                sys.exit(1)

    try:
        cache = apt.Cache()
    except Exception as e:
        if not args.raw:
            print(f"Error initializing apt cache: {e}")
        sys.exit(1)

    manual_packages = sorted(get_manual_installed_packages(cache))
    if not manual_packages:
        if not args.raw:
            print("No manually installed packages found.")
        sys.exit(0)

    if not args.raw:
        print(f"Found {len(manual_packages)} manually installed packages.")
        print("Simulating changes to identify safely markable packages...")

    dep_cache = apt_pkg.DepCache(cache._cache)
    try:
        dep_cache.init()
    except Exception as e:
        if not args.raw:
            print(f"Error initializing dependency cache: {e}")
        sys.exit(1)

    successfully_markable, reverted_changes = identify_markable_packages(
        cache, dep_cache, manual_packages
    )

    if args.raw:
        for pkg_name in successfully_markable:
            print(pkg_name)
        sys.exit(0)
    else:
        if successfully_markable:
            print(
                f"Successfully identified {len(successfully_markable)} packages that can be safely marked as automatic:"
            )
            for pkg_name in successfully_markable:
                print(f"  - {pkg_name}")
        else:
            print(
                "No packages could be safely marked as automatic without triggering autoremoval."
            )

        if reverted_changes:
            print(
                f"Need to keep {len(reverted_changes)} packages as manual to prevent their autoremoval:"
            )
            for pkg_name in reverted_changes:
                print(f"  - {pkg_name}")

        print(
            "\n*** This was a simulation. No actual changes were made to your system. ***"
        )

        if args.apply and successfully_markable:
            print("\nApplying changes...")
            try:
                for pkg_name in successfully_markable:
                    pkg = cache.get(pkg_name)
                    if pkg and pkg.is_installed:
                        pkg.mark_auto()
                        print(
                            f"  Marked '{pkg_name}' as automatically installed."
                        )
                cache.commit()
                print("Changes applied successfully.")
                sys.exit(0)
            except Exception as e:
                print(f"Error applying changes: {e}")
                print(
                    "Some packages might not have been marked as automatic."
                )
                sys.exit(1)
        elif args.apply and not successfully_markable:
            print("\nNo packages to apply changes for based on simulation.")
            sys.exit(0)

    sys.exit(0)