import os, grp, pwd, shutil, subprocess, configparser
from pathlib import Path
from jupyterhub.spawner import Spawner

# host端的仓库目录
SHARED_ROOT  = Path("/home/zhangyi/jupyter-collection")
# jyputerhub用户home目录
USER_HOME    = Path("/home/jupyter-{username}")
# 分组信息路径
CTRL_FILE    = SHARED_ROOT / ".hub" / "resource_map.ini"

# ========== 2. 读控制表 ==========
def load_ctrl():
    """返回 组->资源  和  组->用户  两个字典"""
    cfg = configparser.ConfigParser()
    cfg.read(CTRL_FILE)
    res_map = {g: [r.strip() for r in v.split(",")]
               for g, v in cfg.items("map")}
    usr_map = {g: [u.strip() for u in v.split(",")]
               for g, v in cfg.items("users")}
    return res_map, usr_map

# ========== 3. post_auth_hook：自动建用户+加组 ==========
def ensure_user_and_groups(authenticator, handler, authentication):
    username = authentication['name']
    res_map, usr_map = load_ctrl()

    # 该用户应属的组
    wanted_groups = [g for g, users in usr_map.items() if username in users]
    if not wanted_groups:               # 控制文件里没写就什么都不做
        return authentication

    # 确保组存在
    for g in wanted_groups:
        subprocess.run(["groupadd", "-f", g], check=False)

    # 若系统账号不存在则创建（Hub 自己会创 home，这里确保 shell 账号）
    #if subprocess.run(["id", "-u", username], capture_output=True).returncode != 0:
    #    subprocess.run([
    #        "useradd", "-m", "-s", "/bin/bash",
    #        "-d", str(USER_HOME).format(username=username), username
    #    ], check=False)

    # 加组（幂等）
    for g in wanted_groups:
        subprocess.run(["usermod", "-aG", g, username], check=False)

    return authentication

c.Authenticator.post_auth_hook = ensure_user_and_groups

# ========== 4. pre_spawn_hook：重建干净 home + 只读挂 Git 资源 ==========
def prepare_user(spawner: Spawner):
    username = spawner.user.name
    user_dir = Path(str(USER_HOME).format(username=username))

    # 1. 重建空 home（保证无残留）
    if user_dir.exists():
        shutil.rmtree(user_dir)
    user_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
    # 关键：让 systemd 能进去
    shutil.chown(user_dir, user=username, group=username)

    # 2. 读控制表
    res_map, _ = load_ctrl()

    # 3. 取系统组
    try:
        pwnam = pwd.getpwnam(username)
        gid_list = os.getgrouplist(username, pwnam.pw_gid)
    except (KeyError, OSError):
        return
    user_groups = {grp.getgrgid(g).gr_name for g in gid_list}

    # 4. 挂载该用户该看的资源（只读）
    shared = user_dir / "shared"
    shared.mkdir(exist_ok=True)
    bind_list = []                      # 先攒列表
    for grp_name in user_groups:
        for res in res_map.get(grp_name, []):
            res_path = SHARED_ROOT / res
            if res_path.is_dir():
                bind_list.append({
                    'source': str(res_path),
                    'target': str(shared / res),
                    'read-only': True,
                })
    spawner.bind_mounts = bind_list 

c.Spawner.pre_spawn_hook = prepare_user