import os
import sys
import shutil

BOOT_DRIVE = "H:\\"
STAGING_DIR = r"C:\Users\Snarf\OneDrive\Desktop\code\SnarfCode\boot_staging"

print(f"Checking for DietPi boot filesystem on {BOOT_DRIVE}...")

if not os.path.exists(BOOT_DRIVE):
    print(f"Error: {BOOT_DRIVE} is not mounted.")
    sys.exit(1)

# Check if config.txt or dietpi.txt exists on H:\
boot_indicators = ["config.txt", "dietpi.txt", "cmdline.txt"]
found = [f for f in boot_indicators if os.path.exists(os.path.join(BOOT_DRIVE, f))]

if not found:
    print(f"Warning: {BOOT_DRIVE} does not look like a freshly flashed DietPi boot partition yet (indicators not found).")
    print(f"Current contents of {BOOT_DRIVE}: {os.listdir(BOOT_DRIVE)}")
    sys.exit(2)

print(f"Found boot filesystem! Detected: {found}")

# 1. Copy dietpi.txt
src_dietpi = os.path.join(STAGING_DIR, "dietpi.txt")
dst_dietpi = os.path.join(BOOT_DRIVE, "dietpi.txt")
shutil.copy2(src_dietpi, dst_dietpi)
print(f"  [OK] Injected pre-configured {dst_dietpi}")

# 2. Copy Automation_Custom_Script.sh
src_script = os.path.join(STAGING_DIR, "Automation_Custom_Script.sh")
dst_script = os.path.join(BOOT_DRIVE, "Automation_Custom_Script.sh")
shutil.copy2(src_script, dst_script)
print(f"  [OK] Injected first-boot {dst_script}")

# 3. Patch cmdline.txt with cgroups
cmdline_path = os.path.join(BOOT_DRIVE, "cmdline.txt")
if os.path.exists(cmdline_path):
    with open(cmdline_path, "r", encoding="utf-8") as f:
        cmdline = f.read().strip()
    
    cgroups = "cgroup_enable=cpuset cgroup_memory=1 cgroup_enable=memory net.ifnames=0"
    if "cgroup_memory=1" not in cmdline:
        cmdline = f"{cmdline} {cgroups}\n"
        with open(cmdline_path, "w", encoding="utf-8", newline="\n") as f:
            f.write(cmdline)
        print("  [OK] Patched cmdline.txt with cgroup parameters")
    else:
        print("  [OK] cmdline.txt already contains cgroup parameters")

# 4. Create empty 'ssh' file if needed
ssh_flag = os.path.join(BOOT_DRIVE, "ssh")
with open(ssh_flag, "w") as f:
    pass
print("  [OK] Created empty ssh flag file")

print("\n=== SD Card Configuration Complete! ===")
print("The card is fully prepared to boot unattended into sninfr2lnxwrk01, join Aerodrome K3s, sync time, and register with Zabbix!")