#!/usr/bin/env python3
import os
import sys
import subprocess
import shutil
from pathlib import Path

CHROOT_DIRECTORY = "/srv/chroot/ubuntu"
UBUNTU_RELEASE = "noble"
UBUNTU_MIRROR = "http://archive.ubuntu.com/ubuntu"

def execute_command(command, check=True, shell=False):
    if shell:
        return subprocess.run(command, shell=True, check=check)
    return subprocess.run(command, check=check)

def initialize_chroot():
    chroot_path = Path(CHROOT_DIRECTORY)
    if chroot_path.exists() and list(chroot_path.iterdir()):
        print(f"Chroot directory {CHROOT_DIRECTORY} already exists and is not empty.")
        response = input("Reinitialize? This will delete everything. (yes/no): ")
        if response.lower() != "yes":
            print("Aborting.")
            sys.exit(0)
        shutil.rmtree(CHROOT_DIRECTORY)
    
    if not shutil.which("debootstrap"):
        print("Installing debootstrap...")
        execute_command(["apt", "update"])
        execute_command(["apt", "install", "-y", "debootstrap"])
    
    if not shutil.which("systemd-nspawn"):
        print("Installing systemd-container...")
        execute_command(["apt", "update"])
        execute_command(["apt", "install", "-y", "systemd-container"])
    
    print(f"Creating chroot directory: {CHROOT_DIRECTORY}")
    chroot_path.mkdir(parents=True, exist_ok=True)
    
    print(f"Running debootstrap for {UBUNTU_RELEASE}...")
    execute_command([
        "debootstrap",
        "--variant=minbase",
        UBUNTU_RELEASE,
        CHROOT_DIRECTORY,
        UBUNTU_MIRROR
    ])
    
    print("Configuring APT sources...")
    sources_content = f"""deb {UBUNTU_MIRROR} {UBUNTU_RELEASE} main restricted universe multiverse
deb {UBUNTU_MIRROR} {UBUNTU_RELEASE}-updates main restricted universe multiverse
deb http://security.ubuntu.com/ubuntu {UBUNTU_RELEASE}-security main restricted universe multiverse
"""
    (chroot_path / "etc/apt/sources.list").write_text(sources_content)
    
    print("Setting hostname...")
    (chroot_path / "etc/hostname").write_text("ubuntu-container\n")
    
    print("Updating APT inside container...")
    execute_command([
        "systemd-nspawn",
        "-D", CHROOT_DIRECTORY,
        "--pipe",
        "bash", "-c",
        "export DEBIAN_FRONTEND=noninteractive && apt update"
    ])
    
    print(f"\nContainer initialized successfully at {CHROOT_DIRECTORY}")
    print(f"To enter: sudo python3 {sys.argv[0]}")

def enter_chroot():
    chroot_path = Path(CHROOT_DIRECTORY)
    if not chroot_path.exists() or not list(chroot_path.iterdir()):
        print(f"Container not initialized. Run: sudo python3 {sys.argv[0]} init")
        sys.exit(1)
    
    if not shutil.which("systemd-nspawn"):
        print("systemd-nspawn not found. Installing systemd-container...")
        execute_command(["apt", "update"])
        execute_command(["apt", "install", "-y", "systemd-container"])
    
    print(f"Entering container at {CHROOT_DIRECTORY}...")
    os.execvp("/usr/bin/systemd-nspawn", [
        "systemd-nspawn",
        "-D", CHROOT_DIRECTORY,
        "--resolv-conf=bind-host",
        "--private-users=pick",
        "--drop-capability=CAP_SYS_MODULE,CAP_SYS_TIME,CAP_SYS_BOOT",
        "--setenv=DEBIAN_FRONTEND=noninteractive"
    ])

def main():
    if os.geteuid() != 0:
        print("This script must be run as root (use sudo)")
        sys.exit(1)
    
    if len(sys.argv) == 1:
        enter_chroot()
    elif len(sys.argv) > 1:
        match sys.argv[1]:
            case "init":
                initialize_chroot()
            case _:
                print(f"Unknown command: {sys.argv[1]}")
                print(f"Usage: {sys.argv[0]} [init]")
                sys.exit(1)

if __name__ == "__main__":
    main()