import os, shutil, configparser, pwd
from pathlib import Path

SHARED_ROOT = Path("/home/zhangyi/jupyter-collection")
CTRL_FILE   = SHARED_ROOT / ".hub" / "resource_map.ini"

def prepare_user(spawner):
    username = spawner.user.name                      # 例如 student1
    sys_user = f"jupyter-{username}"                  # 系统账号 jupyter-student1
    user_dir = Path(f"/home/{sys_user}")              # /home/jupyter-student1

    # 确保系统账号存在（TLJH 已建），否则退出
    try:
        pwd.getpwnam(sys_user)
    except KeyError:
        return

    # 提前建家目录 700 + 归系统用户
    user_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
    shutil.chown(user_dir, user=sys_user, group=sys_user)

    # 整目录递归 chown（含 TLJH 留下的 jupyter-xxx 文件）
    for root, dirs, files in os.walk(user_dir):
        for d in dirs:
            shutil.chown(Path(root) / d, user=sys_user, group=sys_user)
        for f in files:
            shutil.chown(Path(root) / f, user=sys_user, group=sys_user)

    # 清空/新建 shared + 拷贝资源（读写副本）
    shared = user_dir / "shared"
    if shared.exists():
        shutil.rmtree(shared)
    shared.mkdir(parents=True, exist_ok=True)
    shutil.chown(shared, user=sys_user, group=sys_user)

    # 拷贝最新资源
    for res in load_map().get(username, []):
        src = SHARED_ROOT / res
        dst = shared / res
        if src.is_dir():
            shutil.copytree(src, dst, dirs_exist_ok=True)

c.Spawner.pre_spawn_hook = prepare_user