#!/usr/bin/env python3

"""
BredOS FEX-Emu Manager

Manages rootfs creation, package installation, and mode switching for FEX compatibility.
"""

import argparse, subprocess, sys, shutil, json
import tempfile, tarfile, pwd, signal, psutil
from os import makedirs, getenv, getuid, path, chown, listdir
from typing import Any, List, Tuple, Dict, Set, Union, Optional
from datetime import datetime
from pathlib import Path

from bredos import curseapp


"""
---------------------------------------------------------------------------------------------------
Static data

Bred man, if you wanna do globals, do shit that gets reused, not every dict and list you see.
("Also", "use", "tuples,", "they're", "faster.")
"""

DEFAULT_ROOTFS_NAME = "bredos-chroot"
DEFAULT_ROOTFS = "~/.fex-emu/RootFS/" + DEFAULT_ROOTFS_NAME
DEFAULT_PACMANCONF = "/usr/share/fex-emu/bredos/pacman.conf"


"""
---------------------------------------------------------------------------------------------------

airootfs in code

Generate using conv.py from source repo.
Unpack manually using unconv.py from source repo, by copying this into a new file and renaming to
"embedded_files".
"""

airootfs = {
    "etc/group": "root:x:0:root\nsys:x:3:bin,bred\nnetwork:x:90:bred\npower:x:98:bred\nadm:x:999:bred\nwheel:x:998:bred\nuucp:x:987:bred\noptical:x:990:bred\nrfkill:x:983:bred\nvideo:x:986:bred\nstorage:x:988:bred\naudio:x:995:bred\nusers:x:985:bred\nnopasswdlogin:x:966:bred\nautologin:x:967:bred\nbred:x:1000:\n",
    "etc/gshadow": "root:::root\nsys:!!::bred\nnetwork:!!::bred\npower:!!::bred\nadm:!!::bred\nwheel:!!::bred\nuucp:!!::bred\noptical:!!::bred\nrfkill:!!::bred\nvideo:!!::bred\nstorage:!!::bred\naudio:!!::bred\nusers:!!::bred\nnopasswdlogin:!::bred\nautologin:!::bred\nbred:!::\n",
    "etc/hostname": "bredos\n",
    "etc/locale.conf": "LANG=C.UTF-8",
    "etc/locale.gen": "C.UTF-8 UTF-8  \n",
    "etc/makepkg.conf": "#!/hint/bash\n# shellcheck disable=2034\n\nDLAGENTS=('file::/usr/bin/curl -qgC - -o %o %u'\n          'ftp::/usr/bin/curl -qgfC - --ftp-pasv --retry 3 --retry-delay 3 -o %o %u'\n          'http::/usr/bin/curl -qgb \"\" -fLC - --retry 3 --retry-delay 3 -o %o %u'\n          'https::/usr/bin/curl -qgb \"\" -fLC - --retry 3 --retry-delay 3 -o %o %u'\n          'rsync::/usr/bin/rsync --no-motd -z %u %o'\n          'scp::/usr/bin/scp -C %u %o')\n\nVCSCLIENTS=('bzr::breezy'\n            'fossil::fossil'\n            'git::git'\n            'hg::mercurial'\n            'svn::subversion')\n\nCARCH=\"x86_64\"\nCHOST=\"x86_64-pc-linux-gnu\"\n\nCFLAGS=\"-march=x86-64 -mtune=generic -O2 -pipe -fno-plt -fexceptions \\\n        -Wp,-D_FORTIFY_SOURCE=3 -Wformat -Werror=format-security \\\n        -fstack-clash-protection -fcf-protection \\\n        -fno-omit-frame-pointer -mno-omit-leaf-frame-pointer\"\nCXXFLAGS=\"$CFLAGS -Wp,-D_GLIBCXX_ASSERTIONS\"\nLDFLAGS=\"-Wl,-O1 -Wl,--sort-common -Wl,--as-needed -Wl,-z,relro -Wl,-z,now \\\n         -Wl,-z,pack-relative-relocs\"\nLTOFLAGS=\"-flto=auto\"\n\nDEBUG_CFLAGS=\"-g\"\nDEBUG_CXXFLAGS=\"$DEBUG_CFLAGS\"\n\nBUILDENV=(!distcc color !ccache check !sign)\nOPTIONS=(strip docs !libtool !staticlibs emptydirs zipman purge debug lto)\n\nINTEGRITY_CHECK=(sha256)\nSTRIP_BINARIES=\"--strip-all\"\nSTRIP_SHARED=\"--strip-unneeded\"\nSTRIP_STATIC=\"--strip-debug\"\nMAN_DIRS=({usr{,/local}{,/share},opt/*}/{man,info})\nDOC_DIRS=(usr/{,local/}{,share/}{doc,gtk-doc} opt/*/{doc,gtk-doc})\nPURGE_TARGETS=(usr/{,share}/info/dir .packlist *.pod)\nDBGSRCDIR=\"/usr/src/debug\"\nLIB_DIRS=('lib:usr/lib' 'lib32:usr/lib32')\n\nCOMPRESSGZ=(gzip -c -f -n)\nCOMPRESSBZ2=(bzip2 -c -f)\nCOMPRESSXZ=(xz -c -z -)\nCOMPRESSZST=(zstd -c -T0 -)\nCOMPRESSLRZ=(lrzip -q)\nCOMPRESSLZO=(lzop -q)\nCOMPRESSZ=(compress -c -f)\nCOMPRESSLZ4=(lz4 -q)\nCOMPRESSLZ=(lzip -c -f)\n\nPKGEXT='.pkg.tar.zst'\nSRCEXT='.src.tar.gz'\n",
    "etc/pacman.conf": "[options]\nHoldPkg = pacman glibc\nArchitecture = x86_64\nColor\nILoveCandy\nParallelDownloads = 5\nSigLevel = Required DatabaseOptional\nLocalFileSigLevel = Optional\n\n[BredOS]\nInclude = /etc/pacman.d/bredos-mirrorlist\n\n[core]\nInclude = /etc/pacman.d/mirrorlist\n\n[extra]\nInclude = /etc/pacman.d/mirrorlist\n\n[multilib]\nInclude = /etc/pacman.d/mirrorlist\n",
    "etc/passwd": "root:x:0:0:root:/root:/bin/bash\nbred:x:1000:1000::/home/bred:/bin/bash\n",
    "etc/shadow": "root:$6$dXftay4w.KBX1xaf$zhIqWyq4ivvUign/rHA5mHkRSho0DJ3gzmjA0ZVfg.CO7hH3svQcOnkCCIr8ggbmIM8Lxw6lYUme5VER0/M7y0:14871::::::\nbred:$6$dXftay4w.KBX1xaf$zhIqWyq4ivvUign/rHA5mHkRSho0DJ3gzmjA0ZVfg.CO7hH3svQcOnkCCIr8ggbmIM8Lxw6lYUme5VER0/M7y0:14871::::::\n",
    "etc/pacman.d/mirrorlist": "##\n## Arch Linux repository mirrorlist\n## Generated on 2024-07-17\n##\n\n## Worldwide\nServer = https://geo.mirror.pkgbuild.com/$repo/os/$arch\nServer = http://mirror.rackspace.com/archlinux/$repo/os/$arch\nServer = https://mirror.rackspace.com/archlinux/$repo/os/$arch\n\n\n",
    "etc/sudoers.d/g_wheel": "%wheel ALL=(ALL:ALL) NOPASSWD: ALL\n",
    "home/bred/.bashrc": "",
}


"""
---------------------------------------------------------------------------------------------------
Logging
"""


def _log_data_format(obj: Any, indent: int = 0) -> str:
    """
    Format any data type into a readable string representation.

    Args:
        obj: The object to format
        indent: Current indentation level for nested structures

    Returns:
        Formatted string representation of the object
    """

    if obj is None:
        return "None"

    spaces = "  " * indent
    obj_type = type(obj)

    # Handle primitive types
    if obj_type is str:
        return repr(obj) if "\n" in obj or "\t" in obj else obj
    elif obj_type in (int, float, bool):
        return str(obj)

    # Handle collections
    elif obj_type is list:
        if not obj:
            return "[]"
        if len(obj) == 1 and type(obj[0]) not in (list, dict, tuple, set):
            return f"[{_log_data_format(obj[0])}]"

        items = []
        next_indent = indent + 1
        for item in obj:
            items.append(f"{spaces}  {_log_data_format(item, next_indent)}")
        return "[\n" + ",\n".join(items) + f"\n{spaces}]"

    elif obj_type is tuple:
        if not obj:
            return "()"
        if len(obj) == 1:
            return f"({_log_data_format(obj[0])},)"
        if len(obj) <= 3 and all(type(x) not in (list, dict, tuple, set) for x in obj):
            return f'({", ".join(_log_data_format(x) for x in obj)})'

        items = []
        next_indent = indent + 1
        for item in obj:
            items.append(f"{spaces}  {_log_data_format(item, next_indent)}")
        return "(\n" + ",\n".join(items) + f"\n{spaces})"

    elif obj_type is dict:
        if not obj:
            return "{}"
        if len(obj) == 1:
            key, value = next(iter(obj.items()))
            if type(value) not in (list, dict, tuple, set):
                return f"{{{_log_data_format(key)}: {_log_data_format(value)}}}"

        items = []
        next_indent = indent + 1
        for key, value in obj.items():
            items.append(
                f"{spaces}  {_log_data_format(key)}: {_log_data_format(value, next_indent)}"
            )
        return "{\n" + ",\n".join(items) + f"\n{spaces}}}"

    elif obj_type is set:
        if not obj:
            return "set()"
        if len(obj) <= 3 and all(type(x) not in (list, dict, tuple, set) for x in obj):
            return f'{{{", ".join(_log_data_format(x) for x in sorted(obj, key=str))}}}'

        items = []
        next_indent = indent + 1
        for item in sorted(obj, key=str):
            items.append(f"{spaces}  {_log_data_format(item, next_indent)}")
        return "{\n" + ",\n".join(items) + f"\n{spaces}}}"

    # Handle objects with attributes
    elif hasattr(obj, "__dict__"):
        attrs = ", ".join(f"{k}={_log_data_format(v)}" for k, v in obj.__dict__.items())
        return f"{obj_type.__name__}({attrs})"

    return str(obj)


def log(data: Any, level: int = 0) -> None:
    """
    ANSI escape logging function that I like.

    Args:
        data: Any data type to log (will be converted to string)
        level: Log level (0=DEBUG, 1=INFO, 2=WARNING, 3=ERROR, 4=CRITICAL)
    """

    RESET = "\033[0m"
    BOLD = "\033[1m"

    COLORS = (
        "\033[36m",  # 0: DEBUG - Cyan
        "\033[32m",  # 1: INFO - Green
        "\033[33m",  # 2: WARNING - Yellow
        "\033[31m",  # 3: ERROR - Red
        "\033[35m\033[1m",  # 4: CRITICAL - Bold Magenta
    )

    LEVEL_NAMES = ("DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL")

    # Clamp level
    level = max(0, min(4, level))

    # Pre-calculate values
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    color = COLORS[level]
    level_name = LEVEL_NAMES[level]

    # Format the data
    data_str = _log_data_format(data)

    # Create the log prefix
    prefix = f"{color}[{timestamp}] {BOLD}{level_name}{RESET}{color}:{RESET}> "

    # Handle multiline content
    lines = data_str.split("\n")

    # Print first line with full prefix
    print(f"{prefix}{lines[0]}")

    # Print subsequent lines with aligned spacing
    if len(lines) > 1:
        # Pre-calculate indent for alignment
        visible_prefix_len = len(f"[{timestamp}] {level_name}:> ")
        indent = " " * visible_prefix_len

        for line in lines[1:]:
            print(f"{indent}{line}")


def log_warn(data: Any) -> None:
    log(data, 2)


def log_err(data: Any) -> None:
    log(data, 3)


"""
---------------------------------------------------------------------------------------------------
mkimage functions
"""


def pacstrap_packages(pacman_conf: str, packages_file: str, install_dir: str) -> None:
    """
    Install packages using pacstrap

    Args:
        pacman_conf: Path to pacman configuration file
        packages_file: Path to file containing package list
        install_dir: Target installation directory
    """
    with open(packages_file) as f:
        packages = map(lambda package: package.strip(), f.readlines())
        packages = list(
            filter(
                lambda package: not (package.startswith("#") or not len(package)),
                packages,
            )
        )
    log("Install dir is:" + install_dir)
    log("Running pacstrap")
    subprocess.run(
        ["pacstrap", "-c", "-C", pacman_conf, "-M", "-G", install_dir] + packages,
        check=True,
    )
    log("Pacstrap complete")

    # Run post-install commands
    for cmd in [["pacman-key", "--init"], ["pacman-key", "--populate"]]:
        log(f"Running post-install command: {' '.join(cmd)}")
        run_chroot_cmd(install_dir, cmd)


def fixperms(target: str) -> None:
    """
    Fix permissions and ownership for specified paths

    Args:
        target: Base target directory
        perms: Dictionary mapping paths to [owner, group, mode]
    """
    realtarget = path.realpath(target)

    perms = {
        "/etc/": ["0", "0", "755"],
        "/etc/polkit-1/rules.d": ["0", "0", "750"],
        "/etc/sudoers.d": ["0", "0", "750"],
        "/home/bred/": ["1000", "1000", "750"],
        "/home": ["0", "0", "755"],
    }

    for i in perms.keys():
        if path.realpath(realtarget + i) != realtarget + (
            i if not i[-1] == "/" else i[:-1]
        ):
            raise OSError("Out of bounds permission fix!")
        if i[-1] == "/":
            subprocess.run(
                [
                    "chown",
                    "-Rh",
                    "--",
                    perms[i][0] + ":" + perms[i][1],
                    realtarget + i,
                ]
            )
        else:
            subprocess.run(
                [
                    "chown",
                    "-hv",
                    "--",
                    perms[i][0] + ":" + perms[i][1],
                    realtarget + i,
                ]
            )
        subprocess.run(["chmod", "--", perms[i][2], realtarget + i])


def run_chroot_cmd(work_dir: str, cmd: list) -> None:
    """
    Run command in rootfs using arch-chroot

    Args:
        work_dir: Rootfs directory path
        cmd: Command to execute as list
    """
    subprocess.run(["arch-chroot", work_dir] + cmd)


def enter_chroot_as_user(work_dir: str, username: str) -> None:
    """
    Enter rootfs and switch to specified user

    Args:
        work_dir: Rootfs directory path
        username: Username to switch to
    """
    subprocess.run(["arch-chroot", work_dir, "su", "-l", username])


def unpack_airootfs(target_dir: str) -> None:
    """
    Copy airootfs directory to target directory

    Args:
        target_dir: Target directory for copy
    """
    log(f"Unpacking airootfs to {target_dir}")
    if path.exists(target_dir):
        shutil.rmtree(target_dir)

    for p, content in airootfs.items():
        out_path = path.join(target_dir, p)
        makedirs(path.dirname(out_path), exist_ok=True)
        with open(out_path, "w", encoding="utf-8") as f:
            f.write(content)

    log("airootfs copied successfully")


def get_user_home() -> str:
    """
    Get the actual user's home directory, even when running as root.

    Returns:
        User home directory path
    """
    # Try to get the original user if running with sudo
    sudo_user = getenv("SUDO_USER")
    if sudo_user:
        return pwd.getpwnam(sudo_user).pw_dir

    # Try LOGNAME environment variable
    logname = getenv("LOGNAME")
    if logname and logname != "root":
        return pwd.getpwnam(logname).pw_dir

    # Fall back to current user
    return path.expanduser("~")


"""
---------------------------------------------------------------------------------------------------
FEX Stuffs
"""


def get_uid() -> int:
    """
    Get the actual user's UID, even when running as root.

    Returns:
        User UID (integer)
    """
    # Try to get the original user if running with sudo
    sudo_user = getenv("SUDO_USER")
    if sudo_user:
        return pwd.getpwnam(sudo_user).pw_uid

    # Try LOGNAME environment variable
    logname = getenv("LOGNAME")
    if logname and logname != "root":
        return pwd.getpwnam(logname).pw_uid

    # Fall back to current user
    return getuid()


def fixowner() -> None:
    """
    Chown sh*t properly

    Necessari.
    """

    uid = get_uid()
    p = get_user_home() + "/.fex-emu"

    chown(p, uid, uid)
    for e in listdir(p):
        chown(path.join(p, e), uid, uid)


def fexkill() -> None:
    """
    Oh look ma it's a bird, no a plane, no a GNU/LINUX/QEMU/FEX INTERPLANETARY
    GAME RUNNER KILLER!!!
    """

    uid = get_uid()

    for proc in psutil.process_iter(["pid", "name", "uids", "exe"]):
        try:
            if (
                proc.info["uids"]
                and proc.info["uids"].real == uid
                and (
                    (proc.info["name"] and "FEXServer" in proc.info["name"])
                    or (proc.info["exe"] and "FEXServer" in proc.info["exe"])
                )
            ):
                proc.send_signal(signal.SIGTERM)
        except:
            pass


def update_fex_config(user_home: str, custom_rootfsdir: bool) -> None:
    """
    Update FEX config to set RootFS to bredos-rootfs if using default path.

    Args:
        user_home: User's home directory path
        custom_rootfsdir: Whether custom rootfs directory was specified
    """
    if custom_rootfsdir:
        return  # Don't modify config if custom rootfsdir was specified

    config_path = path.join(user_home, ".fex-emu", "Config.json")

    # Load existing config or create new one
    if path.exists(config_path):
        with open(config_path, "r") as f:
            config = json.load(f)
        log("Updating existing FEX config")
    else:
        config = {"Config": {}}
        makedirs(path.dirname(config_path), exist_ok=True)
        log("Creating new FEX config")

    # Set RootFS to bredos-rootfs
    config["Config"]["RootFS"] = DEFAULT_ROOTFS_NAME
    config["Config"]["DISABLE_VIXL_INDIRECT_RUNTIME_CALLS"] = "1"
    config["Config"]["NeedsSeccomp"] = "0"
    config["Config"]["ServerSocketPath"] = ""
    config["Config"]["AOTIRLoad"] = "0"
    config["Config"]["AOTIRGenerate"] = "0"
    config["Config"]["AOTIRCapture"] = "0"
    config["Config"]["StartupSleepProcName"] = ""
    config["Config"]["StartupSleep"] = "0"
    config["Config"]["HideHypervisorBit"] = "0"
    config["Config"]["StallProcess"] = "0"
    config["Config"]["ParanoidTSO"] = "0"
    config["Config"]["ABILocalFlags"] = ""
    config["Config"]["X87ReducedPrecision"] = "0"
    config["Config"]["VolatileMetadata"] = "1"
    config["Config"]["TSOAutoMigration"] = "1"
    config["Config"]["StrictInProcessSplitLocks"] = "0"
    config["Config"]["HalfBarrierTSOEnabled"] = "1"
    config["Config"]["MemcpySetTSOEnabled"] = "0"
    config["Config"]["VectorTSOEnabled"] = "0"
    config["Config"]["SingleStep"] = "0"
    config["Config"]["ThunkConfig"] = ""
    config["Config"]["ThunkGuestLibs"] = "\\/usr\\/share\\/fex-emu\\/GuestThunks"
    config["Config"]["ThunkHostLibs"] = "\\/usr\\/lib\\/fex-emu\\/HostThunks"
    config["Config"]["SmallTSCScale"] = "1"
    config["Config"]["PassManagerDumpIR"] = "0"
    config["Config"]["CacheObjectCodeCompilation"] = "0"
    config["Config"]["TSOEnabled"] = "1"
    config["Config"]["HostFeatures"] = "enableatomics,enablelrcpc,enablecrypto,enablepmull128"
    config["Config"]["DumpIR"] = "no"
    config["Config"]["MaxInst"] = "10000"
    config["Config"]["SMCChecks"] = "1"
    config["Config"]["GdbServer"] = "0"
    config["Config"]["Multiblock"] = "1"
    config["Config"]["ProfileStats"] = "0"
    config["Config"]["DumpGPRs"] = "0"
    config["Config"]["O0"] = "0"
    config["Config"]["GlobalJITNaming"] = "0"
    config["Config"]["LibraryJITNaming"] = "0"
    config["Config"]["BlockJITNaming"] = "0"
    config["Config"]["GDBSymbols"] = "0"
    config["Config"]["InjectLibSegFault"] = "0"
    config["Config"]["Disassemble"] = "0"
    config["Config"]["ForceSVEWidth"] = "0"
    config["Config"]["DisableTelemetry"] = "1"
    config["Config"]["SilentLog"] = "1"
    config["Config"]["OutputLog"] = "server"
    config["Config"]["TelemetryDirectory"] = ""
    config["ThunksDB"]["fex_thunk_test"] = 0
    config["ThunksDB"]["asound"] = 0
    config["ThunksDB"]["drm"] = 0
    config["ThunksDB"]["Vulkan"] = 1
    config["ThunksDB"]["WaylandClient"] = 1
    config["ThunksDB"]["GL"] = 1

    # Write config back
    with open(config_path, "w") as f:
        json.dump(config, f, indent=2)

    log(f"FEX config updated at {config_path}")
    fixowner()
    fexkill()


def get_binfmt_status() -> dict:
    """
    Get current status of FEX and QEMU binfmt entries

    Returns:
        Dictionary mapping binfmt entry paths to enabled status
    """
    status = {}

    # Check FEX binfmt entries
    fex_entries = [
        "/proc/sys/fs/binfmt_misc/FEX-x86_64",
        "/proc/sys/fs/binfmt_misc/FEX-x86_64-priority",
    ]
    for entry in fex_entries:
        if path.exists(entry):
            try:
                with open(entry, "r") as f:
                    content = f.read().strip()
                    status[entry] = content.startswith("enabled")
                    break
            except:
                pass

    # Check QEMU binfmt entry
    qemu_entry = "/proc/sys/fs/binfmt_misc/qemu-x86_64"
    if path.exists(qemu_entry):
        try:
            with open(qemu_entry, "r") as f:
                content = f.read().strip()
                status[qemu_entry] = content.startswith("enabled")
        except:
            status[qemu_entry] = False

    return status


def set_binfmt_status(entry_path: str, enabled: bool) -> None:
    """
    Enable or disable a binfmt entry

    Args:
        entry_path: Path to binfmt entry
        enabled: Whether to enable (True) or disable (False)
    """
    if not path.exists(entry_path):
        return

    try:
        with open(entry_path, "w") as f:
            f.write("1" if enabled else "0")
        log(f"{'Enabled' if enabled else 'Disabled'} {entry_path}")
    except Exception as e:
        log_warn(f"Failed to modify {entry_path}: {e}")


def register_qemu_binfmt() -> None:
    # Register QEMU binfmt if it doesn't exist
    qemu_entry = "/proc/sys/fs/binfmt_misc/qemu-x86_64"
    if path.exists(qemu_entry):
        return

    try:
        with open("/usr/lib/binfmt.d/40-qemu-x86_64-static.conf") as f:
            binfmt_data = f.read().strip()

        with open("/proc/sys/fs/binfmt_misc/register", "w") as f:
            f.write(binfmt_data)

        log("Registered QEMU x86_64 binfmt")
    except Exception as e:
        log_warn(f"Failed to register QEMU binfmt: {e}")


def enable_qemu() -> dict:
    """
    Disable FEX and enable QEMU for rootfs operations

    Returns:
        Original binfmt status dictionary
    """
    original_status = get_binfmt_status()

    # Disable FEX entries
    fex_entries = ["/proc/sys/fs/binfmt_misc/FEX-x86_64"]
    for entry in fex_entries:
        if path.exists(entry):
            set_binfmt_status(entry, False)

    # Enable QEMU
    qemu_entry = "/proc/sys/fs/binfmt_misc/qemu-x86_64"
    if path.exists(qemu_entry):
        set_binfmt_status(qemu_entry, True)

    return original_status


def restore_binfmt_status(original_status: dict) -> None:
    """
    Restore original binfmt status

    Args:
        original_status: Dictionary of original binfmt entry states
    """
    for entry_path, was_enabled in original_status.items():
        set_binfmt_status(entry_path, was_enabled)


def install_packages_sysroot(rootfs_dir: str, packages: list) -> None:
    """
    Install packages using pacman --sysroot

    Args:
        rootfs_dir: Rootfs directory path
        packages: List of package names to install
    """
    log(f"Installing packages to {rootfs_dir}: {' '.join(packages)}")

    # Ensure QEMU is enabled for package installation
    original_status = enable_qemu()
    try:
        subprocess.run(
            ["pacman", "--sysroot", rootfs_dir, "-S", "--noconfirm"] + packages,
            check=True,
        )
    except subprocess.CalledProcessError as e:
        log_err(f"Package installation failed: {e}")
    else:
        log("Package installation complete")
    finally:
        # Restore original binfmt status after installation
        restore_binfmt_status(original_status)


def move_for_fex(rootfs_dir: str) -> list:
    """
    Move files and folders to for-rootfs directory for FEX compatibility

    Args:
        rootfs_dir: Rootfs directory path

    Returns:
        List of moved items
    """
    backup_dir = path.join(rootfs_dir, "for-rootfs")
    makedirs(backup_dir, exist_ok=True)

    moved_items = []

    # Move files
    for file_path in (
        "etc/hosts",
        "etc/resolv.conf",
        "etc/timezone",
        "etc/localtime",
        "etc/passwd",
    ):
        src = path.join(rootfs_dir, file_path)
        if path.exists(src):
            dst_dir = path.join(backup_dir, path.dirname(file_path))
            makedirs(dst_dir, exist_ok=True)
            dst = path.join(backup_dir, file_path)

            try:
                shutil.move(src, dst)
                moved_items.append(file_path)
                log(f"Moved file: {file_path}")
            except Exception as e:
                log_warn(f"Failed to move file {file_path}: {e}")

    # Move folders
    for folder_path in (
        "boot",
        "dev",
        "home",
        "media",
        "mnt",
        "proc",
        "root",
        "srv",
        "tmp",
        "var/cache/apt",
        "var/lib/apt",
        "sys",
        "opt",
    ):
        src = path.join(rootfs_dir, folder_path)
        if path.exists(src):
            dst = path.join(backup_dir, folder_path)

            try:
                shutil.move(src, dst)
                moved_items.append(folder_path)
                log(f"Moved folder: {folder_path}")
            except Exception as e:
                log_warn(f"Failed to move folder {folder_path}: {e}")

    # Save list of moved items
    moved_list_file = path.join(backup_dir, "moved_items.txt")
    with open(moved_list_file, "w") as f:
        for item in moved_items:
            f.write(item + "\n")

    log(f"Moved {len(moved_items)} items to for-rootfs directory")
    return moved_items


def restore_for_rootfs(rootfs_dir: str) -> list:
    """
    Restore files and folders from for-rootfs directory for rootfs operations

    Args:
        rootfs_dir: Rootfs directory path

    Returns:
        List of restored items
    """
    backup_dir = path.join(rootfs_dir, "for-rootfs")
    moved_list_file = path.join(backup_dir, "moved_items.txt")

    if not path.exists(moved_list_file):
        log_warn("No moved items list found, nothing to restore")
        return []

    restored_items = []

    with open(moved_list_file, "r") as f:
        items_to_restore = [line.strip() for line in f.readlines()]

    for item_path in items_to_restore:
        src = path.join(backup_dir, item_path)
        dst = path.join(rootfs_dir, item_path)

        if path.exists(src):
            try:
                # Ensure destination directory path.exists
                dst_dir = path.dirname(dst)
                if dst_dir != rootfs_dir:
                    makedirs(dst_dir, exist_ok=True)

                shutil.move(src, dst)
                restored_items.append(item_path)
                log(f"Restored: {item_path}")
            except Exception as e:
                log_warn(f"Failed to restore {item_path}: {e}")

    log(f"Restored {len(restored_items)} items for rootfs")
    return restored_items


def is_fex_mode(rootfs_dir: str) -> bool:
    """
    Check if rootfs is currently in FEX mode (items moved to for-rootfs)

    Args:
        rootfs_dir: Rootfs directory path

    Returns:
        True if in FEX mode, False if in rootfs mode
    """
    backup_dir = path.join(rootfs_dir, "for-rootfs")
    moved_list_file = path.join(backup_dir, "moved_items.txt")
    return path.exists(moved_list_file)


def remove_rootfs(rootfs_dir: str) -> None:
    """
    Remove rootfs directory

    Args:
        rootfs_dir: Rootfs directory path to remove
    """
    if not path.exists(rootfs_dir):
        log_warn(f"Rootfs directory does not exist: {rootfs_dir}")
        return

    log(f"Removing rootfs at {rootfs_dir}")
    shutil.rmtree(rootfs_dir)
    log("Rootfs removed successfully")


def backup_rootfs(rootfs_dir: str) -> str:
    """
    Create a backup archive of rootfs directory using native tar

    Args:
        rootfs_dir: Rootfs directory path to backup

    Returns:
        Path to created backup archive
    """
    if not path.exists(rootfs_dir):
        raise FileNotFoundError(f"Rootfs directory does not exist: {rootfs_dir}")

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_name = f"bredos-rootfs-backup_{timestamp}.tar.xz"
    backup_path = path.join(path.dirname(rootfs_dir), backup_name)

    log(f"Creating backup of {rootfs_dir} to {backup_path} using native tar")
    # Use -I 'xz -5 -T0' for compression
    try:
        subprocess.run(
            [
                "tar",
                "-c",
                "-I",
                "xz -5 -T0",
                "-f",
                backup_path,
                "-C",
                path.dirname(rootfs_dir),
                basename(rootfs_dir),
            ],
            check=True,
        )
    except subprocess.CalledProcessError as e:
        log_err(f"Backup failed: {e}")
        raise
    log(f"Backup created successfully: {backup_path}")
    return backup_path


def restore_rootfs(archive_path: str, rootfs_dir: str) -> None:
    """
    Restore rootfs from backup archive using native tar

    Args:
        archive_path: Path to backup archive
        rootfs_dir: Target rootfs directory path
    """
    if not path.exists(archive_path):
        raise FileNotFoundError(f"Archive file does not exist: {archive_path}")

    if path.exists(rootfs_dir):
        log(f"Removing existing rootfs at {rootfs_dir}")
        shutil.rmtree(rootfs_dir)

    makedirs(path.dirname(rootfs_dir), exist_ok=True)
    log(f"Restoring rootfs from {archive_path} to {rootfs_dir} using native tar")
    try:
        subprocess.run(
            [
                "tar",
                "-x",
                "-I",
                "xz -T0",
                "-f",
                archive_path,
                "-C",
                path.dirname(rootfs_dir),
            ],
            check=True,
        )
    except subprocess.CalledProcessError as e:
        log_err(f"Restore failed: {e}")
        raise
    log("Rootfs restored successfully")


"""
---------------------------------------------------------------------------------------------------
Generic stuffs
"""


def check_root() -> None:
    """
    Check if running as root, exit if not.

    Yes, this script needs it.
    """
    if getuid() != 0:
        log_err("This script must be run as root (use sudo)")
        sys.exit(1)


"""
---------------------------------------------------------------------------------------------------
Main entrypoint
"""


def main() -> None:
    """
    Main function, handles the whole thing.

    --create          -c : Create new rootfs environment
    --rootfs          -r : Enter rootfs environment
    --install-package -i : Install packages using pacman --sysroot
    --fex-mode        -e : Setup the rootfs for usage with FEX (this fixes  usage with FEX)
    --rootfs-mode     -d : Setup chroot rootfs access          (this breaks usage with FEX)
    --remove          -R : Remove rootfs
    --backup          -b : Backup rootfs
    --slim            -s : Use a slim, minimal package set for installation
    --rootfsdir          : Root filesystem directory
    --pacmanconf         : pacman.conf path override
    --pkglist            : Package list override (Takes precedence over --slim)
    """

    parser = argparse.ArgumentParser(
        description="Setup FEX emulation rootfs environment"
    )
    parser.add_argument(
        "--create", "-c", action="store_true", help="Create new rootfs environment"
    )
    parser.add_argument(
        "--rootfs", "-r", action="store_true", help="Enter rootfs environment"
    )

    parser.add_argument(
        "--install-package",
        "-i",
        nargs="+",
        metavar="PACKAGE",
        help="Install packages using pacman --sysroot",
    )

    parser.add_argument(
        "--fex-mode",
        "-e",
        action="store_true",
        help="Setup the rootfs for usage with FEX (this fixes usage with FEX)",
    )

    parser.add_argument(
        "--rootfs-mode",
        "-d",
        action="store_true",
        help="Setup chroot rootfs access (this breaks usage with FEX)",
    )

    parser.add_argument("--remove", "-R", action="store_true", help="Remove rootfs")
    parser.add_argument("--backup", "-b", action="store_true", help="Backup rootfs")
    parser.add_argument(
        "--restore", metavar="TARFILE", help="Restore a rootfs from archive"
    )

    parser.add_argument(
        "--slim",
        "-s",
        action="store_true",
        help="Use a slim, minimal package set for installation",
    )

    parser.add_argument(
        "--rootfsdir",
        default=DEFAULT_ROOTFS,
        help=f"Root filesystem directory (default: {DEFAULT_ROOTFS})",
    )

    parser.add_argument(
        "--pacmanconf",
        default=DEFAULT_PACMANCONF,
        help=f"pacman.conf path override (default: {DEFAULT_ROOTFS})",
    )

    parser.add_argument(
        "--pkglist",
        help=f"Package list override (Takes precedence over --slim)",
    )

    args = parser.parse_args()

    # Register QEMU binfmt if needed
    register_qemu_binfmt()

    # Check initial binfmt status
    initial_binfmt_status = get_binfmt_status()
    log(f"Initial binfmt status: {initial_binfmt_status}")

    # Check if custom rootfsdir was specified
    custom_rootfsdir = args.rootfsdir != DEFAULT_ROOTFS

    # Handle tilde expansion with proper user home detection
    if args.rootfsdir.startswith("~"):
        user_home = get_user_home()
        rootfs_dir = path.abspath(args.rootfsdir.replace("~", user_home))
    else:
        user_home = get_user_home()
        rootfs_dir = path.abspath(args.rootfsdir)

    # Determine which packages file to use
    packages_file = "/usr/share/fex-emu/bredos/"
    if args.pkglist:
        pacmages_file = args.pkglist
        log("Package list override detected")
    else:
        if args.slim:
            packages_file += "packages.minimal"
        else:
            # Default to full if neither --full nor --minimal specified, or if --full specified
            packages_file += "packages.full"

    if args.create:
        check_root()
        if path.exists(rootfs_dir):
            log(f"Removing existing rootfs at {rootfs_dir}")
            shutil.rmtree(rootfs_dir)

        # Manage binfmt pacstrap
        original_status = enable_qemu()

        # Create rootfs directory
        makedirs(rootfs_dir, exist_ok=True)

        # Copy airootfs
        unpack_airootfs(rootfs_dir)

        # Fix permissions
        fixperms(rootfs_dir)

        try:
            # Run pacstrap
            pacstrap_packages(args.pacmanconf, packages_file, rootfs_dir)

        finally:
            # Restore original binfmt status
            log("Restoring original binfmt status")
            restore_binfmt_status(original_status)

        log(f"Chroot environment ready at {rootfs_dir}")

        # Configure for FEX mode by default after creation
        log("Configuring rootfs for FEX mode")
        move_for_fex(rootfs_dir)

        # Update FEX config
        update_fex_config(user_home, custom_rootfsdir)

    elif args.rootfs:
        check_root()
        if not path.exists(rootfs_dir):
            log_err(f"Rootfs directory does not exist: {rootfs_dir}")
            sys.exit(1)

        # Need rootfs mode for rootfs operations
        was_fex_mode = is_fex_mode(rootfs_dir)
        if was_fex_mode:
            log("Switching to rootfs mode")
            restore_for_rootfs(rootfs_dir)

        log(f"Entering rootfs at {rootfs_dir}")

        # Manage binfmt for rootfs
        original_status = enable_qemu()

        try:
            enter_chroot_as_user(rootfs_dir, "bred")
        finally:
            # Restore original binfmt status
            log("Restoring original binfmt status")
            restore_binfmt_status(original_status)

            # Switch back to FEX mode if it was in FEX mode before
            if was_fex_mode:
                log("Switching back to FEX mode")
                move_for_fex(rootfs_dir)

    elif args.install_package:
        check_root()
        if not path.exists(rootfs_dir):
            log_err(f"Rootfs directory does not exist: {rootfs_dir}")
            sys.exit(1)

        # Need rootfs mode for package installation
        was_fex_mode = is_fex_mode(rootfs_dir)
        if was_fex_mode:
            log("Switching to rootfs mode for package installation")
            restore_for_rootfs(rootfs_dir)

        # Manage binfmt for package installation
        original_status = enable_qemu()

        try:
            install_packages_sysroot(rootfs_dir, args.install_package)
        finally:
            # Restore original binfmt status
            log("Restoring original binfmt status")
            restore_binfmt_status(original_status)

            # Switch back to FEX mode if it was in FEX mode before
            if was_fex_mode:
                log("Switching back to FEX mode")
                move_for_fex(rootfs_dir)

    elif args.fex_mode:
        check_root()
        if not path.exists(rootfs_dir):
            log_err(f"Chroot directory does not exist: {rootfs_dir}")
            sys.exit(1)

        if is_fex_mode(rootfs_dir):
            log("Rootfs is already in FEX mode")
        else:
            log("Configuring rootfs for FEX mode")
            move_for_fex(rootfs_dir)

    elif args.rootfs_mode:
        check_root()
        if not path.exists(rootfs_dir):
            log_err(f"Chroot directory does not exist: {rootfs_dir}")
            sys.exit(1)

        if not is_fex_mode(rootfs_dir):
            log("Rootfs is already in rootfs mode")
        else:
            log("Configuring rootfs for rootfs mode")
            restore_for_rootfs(rootfs_dir)

    elif args.remove:
        check_root()
        remove_rootfs(rootfs_dir)

    elif args.backup:
        check_root()
        try:
            backup_path = backup_rootfs(rootfs_dir)
            log(f"Backup completed: {backup_path}")
        except FileNotFoundError as e:
            log_err(str(e))
            sys.exit(1)

    elif args.restore:
        check_root()
        try:
            restore_rootfs(args.restore, rootfs_dir)
            log("Restore completed successfully")
        except FileNotFoundError as e:
            log_err(str(e))
            sys.exit(1)

    else:
        parser.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()
