"""
pipeline_service/work_env.py - 工作环境管理（沙箱的本地/远程切换 + 目录迁移）

工作环境以「用户 / 机构」为单位组织，缺省本地模式（bwrap 沙箱）。
远程模式 = 免密 SSH 到其他主机，远程主机同样部署 bwrap 走沙箱。

目录层级约定：项目目录位于「工作目录」（逻辑账号 deploy_dir）之下，
所以迁移只需同步工作目录顶层，即可覆盖所有项目内容。

- get_work_env: 查询用户/机构工作环境（user 优先 → org 兜底 → 缺省 local）
- set_work_env: 设置工作环境（切换 mode 时触发目录迁移）
- migrate_work_dir: 目录迁移（rsync 本地 deploy_dir ↔ 远程 remote_dir）
- run_remote_sandbox: 远程免密 SSH + bwrap 沙箱执行
"""

import os
import re
import shlex
import shutil
import asyncio
import logging
import subprocess

logger = logging.getLogger("pipeline.work_env")

from .deploy_account import (
    BWRAP, deploy_dir_for, sanitize_username, account_name_for,
)

# rsync 可执行文件（迁移必需）
RSYNC = shutil.which("rsync")


def _row_to_env(row) -> dict:
    """DB 行 → 工作环境 dict。"""
    return {
        "mode": getattr(row, "mode", "local") or "local",
        "remote_host": getattr(row, "remote_host", "") or "",
        "remote_port": getattr(row, "remote_port", 22) or 22,
        "remote_user": getattr(row, "remote_user", "") or "",
        "remote_key_path": getattr(row, "remote_key_path", "") or "",
        "remote_dir": getattr(row, "remote_dir", "") or "",
    }


async def get_work_env(sor, user_id: str, org_id: str = "") -> dict:
    """获取工作环境。优先级：user 级 → org 级 → 缺省 local。"""
    if user_id:
        recs = await sor.sqlExe(
            "SELECT * FROM sd_work_envs WHERE owner_type='user' AND owner_id=${o}$ LIMIT 1",
            {"o": user_id})
        if recs:
            return _row_to_env(recs[0])
    if org_id:
        recs = await sor.sqlExe(
            "SELECT * FROM sd_work_envs WHERE owner_type='org' AND owner_id=${o}$ LIMIT 1",
            {"o": org_id})
        if recs:
            return _row_to_env(recs[0])
    return {"mode": "local", "remote_host": "", "remote_port": 22,
            "remote_user": "", "remote_key_path": "", "remote_dir": ""}


def _ssh_target(env: dict) -> str:
    """构造 scp/rsync 的 user@host 目标。"""
    user = env.get("remote_user", "")
    host = env.get("remote_host", "")
    if user:
        return f"{user}@{host}"
    return host


def _ssh_common_args(env: dict) -> list:
    """ssh/scp/rsync 共用的 SSH 参数。

    StrictHostKeyChecking=accept-new：首次连接记录 host key，之后严格校验，
    防 MITM（host key 变化会拒绝），替代原来的 =no（完全不校验）。
    """
    args = ["-o", "StrictHostKeyChecking=accept-new", "-o", "ConnectTimeout=10"]
    key = env.get("remote_key_path", "")
    if key:
        args += ["-i", key]
    port = env.get("remote_port", 22) or 22
    args += ["-p", str(port)]
    return args


# ── 机构独立 SSH key（每机构一把，防单点泄露） ────────────────

SSH_KEYS_DIR = os.path.expanduser("~/.ssh/org_keys")


def _org_key_dir(org_id: str) -> str:
    """机构 key 目录（org_id sanitize 防路径穿越）。"""
    safe = re.sub(r"[^a-zA-Z0-9_-]", "_", org_id or "")
    if not safe:
        safe = "_"
    return os.path.join(SSH_KEYS_DIR, safe)


def ensure_org_key(org_id: str) -> dict:
    """确保机构有独立 SSH key pair（ed25519），没有则生成。

    Returns: {"key_path", "pubkey_path", "pubkey"}
    """
    d = _org_key_dir(org_id)
    os.makedirs(d, mode=0o700, exist_ok=True)
    priv = os.path.join(d, "id_ed25519")
    pub = priv + ".pub"
    if not (os.path.exists(priv) and os.path.exists(pub)):
        subprocess.run(
            ["ssh-keygen", "-t", "ed25519", "-N", "", "-f", priv,
             "-C", "pipeline-org-" + (re.sub(r"[^a-zA-Z0-9_-]", "_", org_id) or "_")],
            check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        os.chmod(priv, 0o600)
    pubkey = ""
    if os.path.exists(pub):
        with open(pub, "r", encoding="utf-8") as f:
            pubkey = f.read().strip()
    return {"key_path": priv, "pubkey_path": pub, "pubkey": pubkey}


def get_org_pubkey(org_id: str) -> dict:
    """获取机构 public key（供 admin 复制到远程主机 authorized_keys）。"""
    info = ensure_org_key(org_id)
    return {"ok": True, "org_id": org_id, "pubkey": info["pubkey"],
            "key_path": info["key_path"]}


# ── 远程配置安全校验 ───────────────────────────────────────

# 复制粘贴易带入的不可见字符：零宽空格/连字符/BOM 等，str.strip() 去不掉（非 isspace）
_INVISIBLE_RE = re.compile(r"[\s\u200b\u200c\u200d\u2060\ufeff]+")


def _clean_field(s) -> str:
    """清理字段：去所有空白（含 Unicode）与零宽/BOM 不可见字符。"""
    return _INVISIBLE_RE.sub("", s or "")


_HOST_RE = re.compile(r"^[a-zA-Z0-9]([a-zA-Z0-9\-\.]*[a-zA-Z0-9])?$")
_USER_RE = re.compile(r"^[a-zA-Z0-9._-]+$")

# remote_dir 现在是相对路径（从远程登录用户的 home 目录开始）。
# 相对路径天然限制在 home 内，无需系统目录黑名单，只需防路径穿越（..）与 shell 注入。

_REL_DIR_RE = re.compile(r"^[a-zA-Z0-9._/-]+$")


def _rel_remote_dir(remote_dir: str) -> str:
    """规范化相对路径（strip 前后 /），返回纯相对路径（不含 ~，避免 shlex.quote 后 ~ 不展开）。"""
    return (remote_dir or "").strip().strip("/")


def _validate_remote_dir(remote_dir: str) -> str:
    """校验 remote_dir（相对路径），返回错误消息（None 表示通过）。"""
    d = _clean_field(remote_dir)
    if not d:
        return "缺少 remote_dir"
    if d.startswith("/"):
        return "remote_dir 必须是相对路径（从远程用户 home 目录开始）"
    parts = [p for p in d.split("/") if p]
    if ".." in parts:
        return "remote_dir 不能包含 .."
    if not parts:
        return "remote_dir 不能为空"
    if not _REL_DIR_RE.match(d):
        return "remote_dir 只能含字母数字和 . _ - / 字符"
    return None


def _validate_remote(remote_config: dict) -> str:
    """校验远程配置，返回错误消息（None 表示通过）。

    防注入（host/user/port）与防破坏（remote_dir 不能是系统目录、
    remote_key_path 不能读任意私钥文件）。
    """
    host = _clean_field(remote_config.get("remote_host"))
    if not host:
        return "缺少 remote_host"
    # 允许主机名/IPv4；IPv6 用 [] 包裹后取内部
    _host = host[1:-1] if host.startswith("[") and host.endswith("]") else host
    if not _HOST_RE.match(_host):
        return "remote_host 格式非法"

    user = _clean_field(remote_config.get("remote_user"))
    if not user:
        return "缺少 remote_user"
    if not _USER_RE.match(user):
        return "remote_user 格式非法"

    try:
        port = int(remote_config.get("remote_port", 22) or 22)
        if not (1 <= port <= 65535):
            return "remote_port 必须在 1-65535"
    except (ValueError, TypeError):
        return "remote_port 必须是数字"

    err = _validate_remote_dir(remote_config.get("remote_dir", ""))
    if err:
        return err

    key = (remote_config.get("remote_key_path") or "").strip()
    if key:
        ssh_dir = os.path.expanduser("~/.ssh")
        if key.startswith("/"):
            # 绝对路径必须位于 ~/.ssh/ 下，防读任意私钥
            if not key.startswith(ssh_dir.rstrip("/") + "/"):
                return "remote_key_path 必须位于 ~/.ssh/ 目录下"
        elif "/" in key or key.startswith(".."):
            return "remote_key_path 必须是 ~/.ssh/ 下的文件名"

    return None


# ═══════════════════════════════════════════════════════════
# 目录迁移
# ═══════════════════════════════════════════════════════════

async def _remote_mkdir(env: dict, remote_dir: str) -> bool:
    """远程 mkdir -p（相对路径可能多级，rsync 不会递归建父目录）。失败返回 False。

    remote_dir 已校验为安全字符（_REL_DIR_RE），用 $HOME 拼接让远程 shell 展开
    （不用 ~ 前缀，因 shlex.quote 会把 ~ 包进单引号导致不展开）。
    """
    try:
        target = _ssh_target(env)
        ssh_cmd = ["ssh"] + _ssh_common_args(env) + [target, 'mkdir -p "$HOME/' + remote_dir + '"']
        proc = await asyncio.create_subprocess_exec(
            *ssh_cmd, stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.DEVNULL)
        await asyncio.wait_for(proc.communicate(), timeout=30)
        return proc.returncode == 0
    except Exception:
        return False


async def migrate_work_dir(account_name: str, from_mode: str, to_mode: str, env: dict) -> dict:
    """把工作目录内容从旧环境迁移到新环境。

    from_mode/to_mode ∈ {'local', 'remote'}
    - local → remote: rsync 本地 deploy_dir → remote_dir
    - remote → local: rsync remote_dir → 本地 deploy_dir
    """
    local_dir = deploy_dir_for(account_name)
    if to_mode == "local" and not os.path.isdir(local_dir):
        os.makedirs(local_dir, exist_ok=True)

    if not RSYNC:
        return {"ok": False, "error": "rsync 不可用，无法迁移"}

    # remote_dir 安全校验（纵深防御：DB 脏数据/历史记录也拦一道）
    if from_mode == "remote" or to_mode == "remote":
        err = _validate_remote_dir(env.get("remote_dir", ""))
        if err:
            return {"ok": False, "error": err}

    # 构造 rsync 命令
    try:
        if from_mode == "local" and to_mode == "remote":
            # 本地 → 远程
            remote_dir = env.get("remote_dir", "")
            if not remote_dir:
                return {"ok": False, "error": "缺少远程目录 remote_dir"}
            await _remote_mkdir(env, _rel_remote_dir(remote_dir))
            target = _ssh_target(env) + ":~/" + _rel_remote_dir(remote_dir)
            cmd = ["rsync", "-az", "--delete", "-e",
                   "ssh " + " ".join(_ssh_common_args(env)),
                   local_dir.rstrip("/") + "/", target.rstrip("/") + "/"]
        elif from_mode == "remote" and to_mode == "local":
            # 远程 → 本地
            remote_dir = env.get("remote_dir", "")
            if not remote_dir:
                return {"ok": False, "error": "缺少远程目录 remote_dir"}
            src = _ssh_target(env) + ":~/" + _rel_remote_dir(remote_dir)
            cmd = ["rsync", "-az", "--delete", "-e",
                   "ssh " + " ".join(_ssh_common_args(env)),
                   src.rstrip("/") + "/", local_dir.rstrip("/") + "/"]
        else:
            # 同模式或无效，无需迁移
            return {"ok": True, "migrated": False}

        proc = await asyncio.create_subprocess_exec(
            *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=300)
        if proc.returncode != 0:
            return {"ok": False, "error": "迁移失败: " + stderr.decode("utf-8", "replace")[-1000:]}
        return {"ok": True, "migrated": True,
                "detail": stdout.decode("utf-8", "replace")[-500:]}
    except asyncio.TimeoutError:
        return {"ok": False, "error": "迁移超时（300s）"}
    except Exception as e:
        return {"ok": False, "error": str(e)[:500]}


# ═══════════════════════════════════════════════════════════
# 设置工作环境
# ═══════════════════════════════════════════════════════════

async def set_work_env(sor, owner_type: str, owner_id: str, mode: str,
                       remote_config: dict = None, account_name: str = None) -> dict:
    """设置工作环境。切换 mode 时触发目录迁移。

    Returns: {"ok": bool, "mode": str, "migrated": {...}}
    """
    if owner_type not in ("user", "org"):
        return {"ok": False, "error": "owner_type 必须是 user 或 org"}
    if mode not in ("local", "remote"):
        return {"ok": False, "error": "mode 必须是 local 或 remote"}

    remote_config = remote_config or {}
    if owner_type == "org" and mode == "remote":
        # 机构级远程：自动使用机构独立 key（每机构一把，防单点泄露），
        # 不允许 admin 指定任意 key
        _key_info = ensure_org_key(owner_id)
        remote_config["remote_key_path"] = _key_info["key_path"]
    if mode == "remote":
        err = _validate_remote(remote_config)
        if err:
            return {"ok": False, "error": err}

    # 查旧环境（完整记录，迁移时需要用旧的远程配置）
    old_recs = await sor.sqlExe(
        "SELECT * FROM sd_work_envs WHERE owner_type=${t}$ AND owner_id=${o}$ LIMIT 1",
        {"t": owner_type, "o": owner_id})
    old_mode = getattr(old_recs[0], "mode", "local") if old_recs else "local"
    old_env = _row_to_env(old_recs[0]) if old_recs else None

    # upsert
    from appPublic.uniqueID import getID
    data = {
        "owner_type": owner_type,
        "owner_id": owner_id,
        "mode": mode,
        "remote_host": remote_config.get("remote_host", ""),
        "remote_port": int(remote_config.get("remote_port", 22) or 22),
        "remote_user": remote_config.get("remote_user", ""),
        "remote_key_path": remote_config.get("remote_key_path", ""),
        "remote_dir": remote_config.get("remote_dir", ""),
    }
    # 先迁移（切换模式时），成功后再更新 DB（保证状态一致，迁移失败不回写 mode）
    migrated = {"migrated": False}
    if account_name and old_mode != mode:
        # remote → local 用旧的远程配置；local → remote 用新的远程配置
        if old_mode == "remote" and mode == "local":
            migrate_env = old_env or _row_to_env_from_dict(data)
        else:
            migrate_env = _row_to_env_from_dict(data)
        migrated = await migrate_work_dir(account_name, old_mode, mode, migrate_env)
        if not migrated.get("ok"):
            return {"ok": False, "error": "迁移失败，未切换模式: " + migrated.get("error", ""),
                    "migration": migrated}

    # upsert（迁移成功后才写 DB）
    if old_recs:
        await sor.sqlExe(
            "UPDATE sd_work_envs SET mode=${mode}$, remote_host=${h}$, remote_port=${p}$, "
            "remote_user=${u}$, remote_key_path=${k}$, remote_dir=${d}$ "
            "WHERE owner_type=${t}$ AND owner_id=${o}$",
            {"mode": mode, "h": data["remote_host"], "p": data["remote_port"],
             "u": data["remote_user"], "k": data["remote_key_path"],
             "d": data["remote_dir"], "t": owner_type, "o": owner_id})
    else:
        data["id"] = getID()
        await sor.C("sd_work_envs", data)

    # 切换到远程模式（或远程配置变更）后自动初始化远程环境：
    # bwrap → 技能同步（公共/产线/项目/技能集/本机构组织技能）→ 依赖软件安装。
    # 失败不阻断模式切换（报告随返回，可在远程空间页手动重跑）。
    remote_init = None
    if mode == "remote" and (old_mode != "remote" or migrated.get("migrated")):
        try:
            from .remote_env_init import init_remote_env
            remote_init = await init_remote_env(data, org_id=owner_id if owner_type == "org" else "",
                                                 sor=sor)
        except Exception as e:
            remote_init = {"ok": False, "error": "远程环境初始化异常: " + str(e)[:300]}

    return {"ok": True, "mode": mode, "old_mode": old_mode,
            "migration": migrated, "remote_init": remote_init}


def _row_to_env_from_dict(data: dict) -> dict:
    return {
        "mode": data.get("mode", "local"),
        "remote_host": data.get("remote_host", ""),
        "remote_port": data.get("remote_port", 22),
        "remote_user": data.get("remote_user", ""),
        "remote_key_path": data.get("remote_key_path", ""),
        "remote_dir": data.get("remote_dir", ""),
    }


# ═══════════════════════════════════════════════════════════
# 远程沙箱执行
# ═══════════════════════════════════════════════════════════

def _bwrap_static_args() -> list:
    """bwrap 沙箱固定参数（不含 --bind 目录和 command，由调用方拼接）。"""
    return [
        "--unshare-user", "--unshare-pid", "--unshare-ipc", "--unshare-uts",
        "--die-with-parent",
        "--ro-bind", "/usr", "/usr",
        "--ro-bind", "/bin", "/bin",
        "--ro-bind", "/sbin", "/sbin",
        "--ro-bind", "/lib", "/lib",
        "--ro-bind", "/lib64", "/lib64",
        "--ro-bind", "/etc", "/etc",
        "--proc", "/proc",
        "--dev", "/dev",
        "--tmpfs", "/tmp",
    ]


async def run_remote_sandbox(env: dict, deploy_dir: str, command: str,
                             workdir: str = "", timeout: int = 120) -> dict:
    """远程免密 SSH + bwrap 沙箱执行命令。"""
    ssh_args = _ssh_common_args(env)
    target = _ssh_target(env)

    # deploy_dir 是相对路径（已校验安全字符），用 $HOME 拼接让远程 shell 展开
    rel = _rel_remote_dir(deploy_dir)
    inner = ('BW=$(command -v bwrap 2>/dev/null || echo "$HOME/bin/bwrap"); '
             'RD="$HOME/' + rel + '"; '
             'exec "$BW" ' + " ".join(_bwrap_static_args()) + ' '
             '--bind "$RD" /home --chdir /home --setenv HOME /home '
             '-- bash -c ' + shlex.quote(command))

    ssh_cmd = ["ssh"] + ssh_args + [target, inner]

    try:
        proc = await asyncio.create_subprocess_exec(
            *ssh_cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout + 30)
    except asyncio.TimeoutError:
        return {"rc": -1, "stdout": "", "stderr": f"远程命令超时（{timeout}s）", "sandbox": True, "remote": True}
    except Exception as e:
        return {"rc": -1, "stdout": "", "stderr": f"SSH 连接失败: {str(e)[:300]}", "sandbox": True, "remote": True}

    return {
        "rc": proc.returncode or 0,
        "stdout": stdout.decode("utf-8", "replace")[-8000:],
        "stderr": stderr.decode("utf-8", "replace")[-4000:],
        "sandbox": True,
        "remote": True,
    }


async def run_in_work_env(sor, account_name: str, user_id: str, org_id: str,
                          command: str, workdir: str = "", timeout: int = 120,
                          network: str = "shared") -> dict:
    """按工作环境 mode 分发执行：本地 bwrap 或远程 SSH+bwrap。"""
    env = await get_work_env(sor, user_id, org_id)
    if env.get("mode") == "remote":
        remote_dir = env.get("remote_dir", "")
        if not remote_dir:
            return {"rc": -1, "stdout": "", "stderr": "远程模式缺少 remote_dir",
                    "sandbox": True, "remote": True}
        return await run_remote_sandbox(env, remote_dir, command, workdir, timeout)
    from .deploy_account import run_in_sandbox
    return await run_in_sandbox(sor, account_name, command, workdir, timeout, network)


# ═══════════════════════════════════════════════════════════
# 远程 bwrap 部署（远程主机装 bwrap，零 root）
# ═══════════════════════════════════════════════════════════

async def ensure_remote_bwrap(env: dict) -> dict:
    """确保远程主机已安装 bwrap（零 root：apt download + dpkg -x）。"""
    # 校验（host/user/port/key，无需 remote_dir）
    host = _clean_field(env.get("remote_host"))
    _host = host[1:-1] if host.startswith("[") and host.endswith("]") else host
    if not host or not _HOST_RE.match(_host):
        return {"ok": False, "error": "remote_host 格式非法"}
    user = _clean_field(env.get("remote_user"))
    if not user or not _USER_RE.match(user):
        return {"ok": False, "error": "remote_user 格式非法"}
    key = (env.get("remote_key_path") or "").strip()
    if key and key.startswith("/"):
        ssh_dir = os.path.expanduser("~/.ssh")
        if not key.startswith(ssh_dir.rstrip("/") + "/"):
            return {"ok": False, "error": "remote_key_path 必须位于 ~/.ssh/ 目录下"}

    ssh_args = _ssh_common_args(env)
    target = _ssh_target(env)

    script = (
        'mkdir -p ~/bin && '
        'if ! command -v bwrap >/dev/null 2>&1 && [ ! -x ~/bin/bwrap ]; then '
        'tmpd=$(mktemp -d) && (cd "$tmpd" && apt download bubblewrap >/dev/null 2>&1 '
        '&& dpkg -x bubblewrap*.deb . && cp usr/bin/bwrap ~/bin/bwrap) 2>/dev/null; '
        'rm -rf "$tmpd"; fi; '
        'if command -v bwrap >/dev/null 2>&1; then echo PATH_BWRAP; '
        'elif [ -x ~/bin/bwrap ]; then echo HOME_BWRAP; '
        'else echo NO_BWRAP; fi'
    )
    ssh_cmd = ["ssh"] + ssh_args + [target, script]
    try:
        proc = await asyncio.create_subprocess_exec(
            *ssh_cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=60)
        out = stdout.decode("utf-8", "replace").strip()
        if "NO_BWRAP" in out:
            return {"ok": False, "error": "远程主机无法获取 bwrap: " + stderr.decode("utf-8", "replace")[-300:]}
        return {"ok": True, "location": out.strip().splitlines()[-1] if out else "unknown"}
    except Exception as e:
        return {"ok": False, "error": str(e)[:300]}