#!/usr/bin/python3
"""Add a new AI-enabled user to the host Debian system and optional chroots."""

import click
import glob
import os
import pwd
import grp
import subprocess
import sys

DEFAULT_CHROOT = "/mnt/r/f"


def run_cmd(*args, check=True):
    cmd = " ".join(str(a) for a in args)
    print(f"+ {cmd}")
    subprocess.run(args, check=check)


def read_ids(path, idx):
    """Read numeric IDs from a passwd/group file at the given field index."""
    ids = set()
    try:
        with open(path) as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                parts = line.split(":")
                if len(parts) > idx:
                    try:
                        ids.add(int(parts[idx]))
                    except ValueError:
                        pass
    except FileNotFoundError:
        pass
    return ids


def login_defs(key, default):
    try:
        with open("/etc/login.defs") as f:
            for line in f:
                line = line.strip()
                if line.startswith(key + " "):
                    parts = line.split()
                    if len(parts) >= 2:
                        return int(parts[1])
    except (FileNotFoundError, ValueError):
        pass
    return default


def find_free_uid(system, chroots):
    if system:
        lo = login_defs("SYS_UID_MIN", 100)
        hi = login_defs("SYS_UID_MAX", 999)
    else:
        lo = login_defs("UID_MIN", 1000)
        hi = login_defs("UID_MAX", 60000)
    used = read_ids("/etc/passwd", 2)
    for chroot in chroots:
        used |= read_ids(os.path.join(chroot, "etc", "passwd"), 2)
    ids = range(hi, lo - 1, -1) if system else range(lo, hi + 1)
    for uid in ids:
        if uid not in used:
            return uid
    raise RuntimeError(f"No free UID in range {lo}-{hi}")


def find_free_gid(system, chroots):
    if system:
        lo = login_defs("SYS_GID_MIN", 100)
        hi = login_defs("SYS_GID_MAX", 999)
    else:
        lo = login_defs("GID_MIN", 1000)
        hi = login_defs("GID_MAX", 60000)
    used = read_ids("/etc/group", 2)
    for chroot in chroots:
        used |= read_ids(os.path.join(chroot, "etc", "group"), 2)
    ids = range(hi, lo - 1, -1) if system else range(lo, hi + 1)
    for gid in ids:
        if gid not in used:
            return gid
    raise RuntimeError(f"No free GID in range {lo}-{hi}")


def require_host_group(name):
    try:
        return grp.getgrnam(name).gr_gid
    except KeyError:
        raise RuntimeError(f"Required group '{name}' not found on host")


def ensure_chroot_group(chroot, name, gid):
    path = os.path.join(chroot, "etc", "group")
    try:
        with open(path) as f:
            for line in f:
                parts = line.strip().split(":")
                if len(parts) >= 3 and parts[0] == name:
                    if int(parts[2]) != gid:
                        raise RuntimeError(
                            f"Group {name} exists in chroot {chroot} "
                            f"with GID {parts[2]}, expected {gid} from host"
                        )
                    return
    except FileNotFoundError:
        pass
    run_cmd("chroot", chroot, "groupadd", "-g", str(gid), name)


def user_exists_chroot(chroot, name):
    path = os.path.join(chroot, "etc", "passwd")
    try:
        with open(path) as f:
            for line in f:
                if line.startswith(f"{name}:"):
                    return True
    except FileNotFoundError:
        pass
    return False


def _is_matching_key(line, real_user):
    """
    Return the key string if *line* is an ssh-ed25519 key whose comment
    equals *real_user* or starts with "real_user@".
    """
    line = line.strip()
    if not line or line.startswith("#"):
        return None
    parts = line.split()
    if len(parts) < 2 or parts[0] != "ssh-ed25519":
        return None
    comment = parts[2] if len(parts) > 2 else ""
    if comment == real_user or comment.startswith(f"{real_user}@"):
        return line
    return None


def _ssh_auth_sock_from_user(username):
    """Find SSH_AUTH_SOCK from any running process of *username*."""
    try:
        pw = pwd.getpwnam(username)
        uid = pw.pw_uid
    except KeyError:
        return None
    for pid_dir in glob.glob("/proc/[0-9]*"):
        try:
            if os.stat(pid_dir).st_uid != uid:
                continue
            env_path = os.path.join(pid_dir, "environ")
            with open(env_path, "rb") as f:
                for item in f.read().split(b"\x00"):
                    if item.startswith(b"SSH_AUTH_SOCK="):
                        return item.decode().split("=", 1)[1]
        except (OSError, PermissionError):
            continue
    return None


def get_ssh_key():
    real_user = os.environ.get("SUDO_USER") or os.environ.get("USER") or os.getlogin()

    for ssh_auth_sock in (
        os.environ.get("SSH_AUTH_SOCK"),
        _ssh_auth_sock_from_user(real_user),
    ):
        if not ssh_auth_sock:
            continue
        env = {**os.environ, "SSH_AUTH_SOCK": ssh_auth_sock}
        try:
            result = subprocess.run(
                ["ssh-add", "-L"], capture_output=True, text=True, check=True, env=env
            )
            for line in result.stdout.splitlines():
                key = _is_matching_key(line, real_user)
                if key:
                    return key
        except (subprocess.CalledProcessError, FileNotFoundError):
            pass

    try:
        home = pwd.getpwnam(real_user).pw_dir
    except KeyError:
        home = os.path.expanduser(f"~{real_user}")
    auth_path = os.path.join(home, ".ssh", "authorized_keys")
    try:
        with open(auth_path) as f:
            for line in f:
                key = _is_matching_key(line, real_user)
                if key:
                    return key
    except FileNotFoundError:
        pass

    raise RuntimeError(
        f"No ssh-ed25519 key found for user {real_user!r} "
        f"(checked ssh-add -L and {auth_path})"
    )


@click.command()
@click.option("-s", "system", is_flag=True, help="Create a system user")
@click.option(
    "-r", "--root",
    "chroots",
    multiple=True,
    default=(DEFAULT_CHROOT,),
    show_default=True,
    help=(
        "Chroot to also add the user to. "
        "Pass '-' for local system only. "
        "May be given multiple times."
    ),
)
@click.argument("username")
def main(username, system, chroots):
    if chroots == ("-",):
        chroots = ()
    elif "-" in chroots:
        raise click.UsageError("'-/' must be given alone (no other --root options)")

    if os.geteuid() != 0:
        raise click.UsageError("must be run as root")

    home = f"/home/{username}"

    for chroot in chroots:
        if not os.path.isdir(chroot):
            raise RuntimeError(f"chroot {chroot} not found")

    try:
        pwd.getpwnam(username)
        raise RuntimeError(f"user {username} already exists on host")
    except KeyError:
        pass

    for chroot in chroots:
        if user_exists_chroot(chroot, username):
            raise RuntimeError(f"user {username} already exists in chroot {chroot}")

    uid = find_free_uid(system, chroots)
    gid = find_free_gid(system, chroots)
    click.echo(f"Assigning UID={uid}, GID={gid}")

    ssh_key = get_ssh_key()

    render_gid = require_host_group("render")
    ollama_gid = require_host_group("ollama")

    # Primary group
    run_cmd("groupadd", "-g", str(gid), username)
    for chroot in chroots:
        ensure_chroot_group(chroot, username, gid)

    # Create user on host (with home)
    host_cmd = [
        "/usr/sbin/useradd",
        "-u", str(uid),
        "-g", str(gid),
        "-d", home,
        "-m",
        "-s", "/bin/bash",
    ]
    if system:
        host_cmd.append("-r")
    host_cmd.append(username)
    run_cmd(*host_cmd)

    # Create user in each chroot (home is shared, do not create it there)
    for chroot in chroots:
        chroot_cmd = [
            "chroot", chroot,
            "useradd",
            "-u", str(uid),
            "-g", str(gid),
            "-d", home,
            "-M",
            "-s", "/bin/bash",
        ]
        if system:
            chroot_cmd.append("-r")
        chroot_cmd.append(username)
        run_cmd(*chroot_cmd)

    # Disable password login
    run_cmd("usermod", "-L", username)
    for chroot in chroots:
        run_cmd("chroot", chroot, "usermod", "-L", username)

    # Supplementary groups
    for chroot in chroots:
        ensure_chroot_group(chroot, "render", render_gid)
        ensure_chroot_group(chroot, "ollama", ollama_gid)

    run_cmd("usermod", "-aG", "render,ollama", username)
    for chroot in chroots:
        run_cmd("chroot", chroot, "usermod", "-aG", "render,ollama", username)

    # Authorised SSH key (home is shared, do once on host)
    ssh_dir = os.path.join(home, ".ssh")
    auth_keys = os.path.join(ssh_dir, "authorized_keys")
    os.makedirs(ssh_dir, mode=0o700, exist_ok=True)
    with open(auth_keys, "w") as f:
        f.write(ssh_key + "\n")
    os.chmod(auth_keys, 0o600)
    os.chown(ssh_dir, uid, gid)
    os.chown(auth_keys, uid, gid)

    # Empty .config/systemd directories
    config_dir = os.path.join(home, ".config")
    systemd_dir = os.path.join(config_dir, "systemd")
    os.makedirs(os.path.join(systemd_dir, "user"), mode=0o755, exist_ok=True)
    os.makedirs(os.path.join(systemd_dir, "user.control"), mode=0o755, exist_ok=True)
    run_cmd("chown", "-R", f"{uid}:{gid}", config_dir)

    click.echo(f"User {username} created successfully.")


if __name__ == "__main__":
    try:
        main(standalone_mode=False)
    except click.ClickException as e:
        e.show()
        sys.exit(e.exit_code)
    except RuntimeError as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)
    except subprocess.CalledProcessError as e:
        click.echo(f"Command failed: {e}", err=True)
        sys.exit(1)
