#!/usr/bin/env python3
"""
setup-sd.py - Customise a Devuan Pi image and produce a new image.
Usage:
    python3 setup_pi_sd.py input.img output.img [--config pi.toml]
    python3 setup_pi_sd.py input.zip output.img [--config pi.toml]
    python3 setup_pi_sd.py input.img output.img --sudoers
The --sudoers flag prints the exact sudoers rules needed to run this tool.
Install with:
    python3 setup_pi_sd.py input.img output.img --sudoers | sudo tee /etc/sudoers.d/setup-sd
    sudo chmod 0440 /etc/sudoers.d/setup-sd
Dependencies:
    docker (for package installation and user creation)
"""
import os
import sys
import getpass
import subprocess
import argparse
import tomllib
import atexit
import json
import shutil
import zipfile
import struct

DEFAULTS = {
    "hostname":     "tether-pi",
    # ethernet
    "eth_dhcp":     True,
    "eth_ip":       "192.168.0.2",
    "eth_netmask":  "255.255.255.0",
    "eth_gateway":  None,
    # wifi
    "wifi":         False,
    "wifi_dhcp":    True,
    "wlan0_ip":     "192.168.3.200",
    "wifi_netmask": "255.255.255.0",
    "wifi_gateway": "192.168.3.1",
    # packages to install via chroot
    "packages":     [],
    # users to create
    "users":        [],
}

def ask(prompt, default=None):
    suffix = f" [{default}]" if default else ""
    val = input(f"{prompt}{suffix}: ").strip()
    return val or default

def ask_bool(prompt, default=True):
    suffix = " [Y/n]" if default else " [y/N]"
    val = input(f"{prompt}{suffix}: ").strip().lower()
    if not val:
        return default
    return val.startswith('y')

def run(cmd, check=True):
    subprocess.run(cmd, shell=True, check=check)

def check_deps():
    missing = []
    for tool in ["docker"]:
        if not shutil.which(tool):
            missing.append(tool)
    if missing:
        print(f"Error: missing dependencies: {', '.join(missing)}")
        print(f"Install with: sudo apt install docker.io")
        sys.exit(1)

def build_docker_image():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    dockerfile = os.path.join(script_dir, "Dockerfile")
    if not os.path.exists(dockerfile):
        print(f"Error: Dockerfile not found at {dockerfile}")
        sys.exit(1)
    print("Building Docker image pi-setup ...")
    run(f"sudo docker build -t pi-setup {script_dir}")

def get_config(config_file):
    cfg = dict(DEFAULTS)
    if config_file:
        with open(config_file, "rb") as f:
            cfg.update(tomllib.load(f))
        if cfg.get("wifi") and "wifi_pass" not in cfg:
            cfg["wifi_pass"] = getpass.getpass("WiFi password: ")
        if cfg.get("wifi") and "wifi_ssid" not in cfg:
            cfg["wifi_ssid"] = ask("WiFi SSID")
        for user in cfg.get("users", []):
            if "password" not in user:
                user["password"] = getpass.getpass(f"Password for {user['name']}: ")
    else:
        cfg["hostname"] = ask("Hostname", cfg["hostname"])
        cfg["eth_dhcp"] = ask_bool("Use DHCP for ethernet?", cfg["eth_dhcp"])
        if not cfg["eth_dhcp"]:
            cfg["eth_ip"]      = ask("Ethernet IP",      cfg["eth_ip"])
            cfg["eth_netmask"] = ask("Ethernet netmask", cfg["eth_netmask"])
            cfg["eth_gateway"] = ask("Ethernet gateway (blank for none)", cfg["eth_gateway"])
        cfg["wifi"] = ask_bool("Configure WiFi?", cfg["wifi"])
        if cfg["wifi"]:
            cfg["wifi_ssid"]    = ask("WiFi SSID")
            cfg["wifi_pass"]    = getpass.getpass("WiFi password: ")
            cfg["wifi_dhcp"]    = ask_bool("Use DHCP for WiFi?", cfg["wifi_dhcp"])
            if not cfg["wifi_dhcp"]:
                cfg["wlan0_ip"]     = ask("WiFi IP",      cfg["wlan0_ip"])
                cfg["wifi_netmask"] = ask("WiFi netmask", cfg["wifi_netmask"])
                cfg["wifi_gateway"] = ask("WiFi gateway", cfg["wifi_gateway"])
    return cfg

def run_tasks(mountpoint, cfg):
    tasks = {
        "packages": cfg.get("packages", []),
        "users":    cfg.get("users", []),
        "config":   {k: v for k, v in cfg.items() if k not in ("packages", "users")},
    }
    tasks_json = json.dumps(tasks)
    docker_cmd = (
        f"sudo docker run --rm -i"
        f" -v {mountpoint}:/mnt/pi"
        f" -v /dev/null:/mnt/pi/dev/null"
        f" -v /etc/resolv.conf:/etc/resolv.conf"
        f" pi-setup"
    )
    print(f"\nRunning: {docker_cmd}")
    subprocess.run(docker_cmd, input=tasks_json, shell=True, text=True, check=True)

def get_partition2_offset(img_path):
    with open(img_path, 'rb') as f:
        f.seek(446)
        for i in range(4):
            entry = f.read(16)
            if i == 1:
                lba_start = struct.unpack_from('<I', entry, 8)[0]
                lba_size  = struct.unpack_from('<I', entry, 12)[0]
                return lba_start * 512, lba_size * 512
    raise RuntimeError("Could not find partition 2 in image")

def print_sudoers(output_img, offset, mountpoint):
    user       = os.environ.get("USER", "YOUR_USER")
    mount      = shutil.which("mount")  or "/bin/mount"
    umount     = shutil.which("umount") or "/bin/umount"
    docker     = shutil.which("docker") or "/usr/bin/docker"
    script_dir = os.path.dirname(os.path.abspath(output_img))
    print(f"# Generated by setup-sd --sudoers")
    print(f"# Install: setup-sd ... --sudoers | sudo tee /etc/sudoers.d/setup-sd && sudo chmod 0440 /etc/sudoers.d/setup-sd")
    print()
    print(f"{user} ALL=(ALL) NOPASSWD: {mount} -o loop\\,nosuid\\,nodev\\,offset={offset} {output_img} {mountpoint}")
    print(f"{user} ALL=(ALL) NOPASSWD: {umount} {mountpoint}")
    print(f"{user} ALL=(ALL) NOPASSWD: {docker} build -t pi-setup {script_dir}")
    print(f"{user} ALL=(ALL) NOPASSWD: {docker} run --rm -i -v {mountpoint}\\:/mnt/pi -v /dev/null\\:/mnt/pi/dev/null -v /etc/resolv.conf\\:/etc/resolv.conf pi-setup")

def main():
    parser = argparse.ArgumentParser(description="Customise a Devuan Pi image.")
    parser.add_argument("input",  help="Input image or zip file")
    parser.add_argument("output", help="Output image file")
    parser.add_argument("--config",   metavar="FILE", help="TOML config file")
    parser.add_argument("--sudoers",  action="store_true", help="Print sudoers rules and exit")
    parser.add_argument("--shell",    action="store_true", help="Drop into bash shell in Docker container")
    args = parser.parse_args()

    output_img = os.path.abspath(args.output)
    mountpoint = os.path.join(os.path.dirname(output_img), ".mnt")

    if args.sudoers:
        if os.path.exists(output_img):
            offset, _ = get_partition2_offset(output_img)
        elif args.input.endswith(".zip"):
            print("Note: run without --sudoers first to produce the output image, then re-run with --sudoers", file=sys.stderr)
            sys.exit(1)
        else:
            offset, _ = get_partition2_offset(args.input)
        print_sudoers(output_img, offset, mountpoint)
        return

    check_deps()
    print("=== Pi image setup ===\n")
    cfg = get_config(args.config)
    print()

    if args.input.endswith(".zip"):
        print(f"Extracting {args.input} -> {output_img} ...")
        with zipfile.ZipFile(args.input) as z:
            imgs = [n for n in z.namelist() if n.endswith(".img")]
            if not imgs:
                print("Error: no .img file found in zip"); sys.exit(1)
            print(f"  found {imgs[0]}")
            with z.open(imgs[0]) as src, open(output_img, 'wb') as dst:
                shutil.copyfileobj(src, dst)
    else:
        print(f"Copying {args.input} -> {output_img} ...")
        shutil.copy2(args.input, output_img)

    offset, size = get_partition2_offset(output_img)
    print(f"Partition 2: offset={offset}, size={size}")

    os.makedirs(mountpoint, exist_ok=True)

    def cleanup():
        run(f"sudo umount {mountpoint}", check=False)
    atexit.register(cleanup)

    print(f"Mounting {output_img} -> {mountpoint} ...")
    run(f"sudo mount -o loop,nosuid,nodev,offset={offset} {output_img} {mountpoint}")

    build_docker_image()

    if args.shell:
        docker_cmd = (
            f"sudo docker run --rm -it"
            f" --entrypoint bash"
            f" -v {mountpoint}:/mnt/pi"
            f" -v /dev/null:/mnt/pi/dev/null"
            f" -v /etc/resolv.conf:/etc/resolv.conf"
            f" pi-setup"
        )
        print(f"\nDropping into shell: {docker_cmd}")
        subprocess.run(docker_cmd, shell=True)
    else:
        run_tasks(mountpoint, cfg)

    print(f"\nUnmounting...")
    run(f"sudo umount {mountpoint}")
    print("\n=== Done! ===")
    print(f"Flash with: dd if={output_img} of=/dev/sdX bs=4M status=progress && sync")
    print(f"\nTo generate sudoers rules for next run:")
    print(f"  setup-sd {args.input} {args.output} --sudoers | sudo tee /etc/sudoers.d/setup-sd")

if __name__ == "__main__":
    main()