"""远程工作环境自动初始化：技能同步 + 依赖软件自动安装（2026-08-28）。

设计：
- 技能同步：本地技能树（公共全局 / 产线 / 项目通用 / 技能集 / 本机构组织技能）
  rsync 到远程 ~/<remote_dir>/skills/。沙箱把该目录挂成 /home，技能在沙箱内
  路径为 /home/skills/，agent 可直接读取。
- 依赖安装：解析同步技能 SKILL.md frontmatter 的 dependencies: 字段，
  机器可解析的 pip 规格（含版本约束）装进远程 ~/<remote_dir>/.venv-skills
  （沙箱内 /home/.venv-skills，技能脚本用该 venv 的 python）。
  非机器可解析依赖（自然语言）不猜，跳过并汇报。
- 重量级依赖（torch/训练框架等，几 GB 且常需 GPU）默认跳过，避免远程机
  被撑爆；跳过名单可通过 appbase 参数 remote_env_heavy_skip 扩展。
- pip 源：appbase 参数 remote_env_pip_index，默认阿里源。
- 全部失败不抛异常：以结构化报告返回，由调用方决定是否阻断/提示。
"""
import base64
import json
import os
import re
import shlex
import asyncio
import logging

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

from .work_env import (
    _ssh_common_args, _ssh_target, _rel_remote_dir,
    _validate_remote_dir, _remote_mkdir,
)

# 重量级 pip 包默认跳过（远程部署机通常无 GPU，torch 系单包 2GB+）。
# 名单可通过 appbase 参数 remote_env_heavy_skip 覆盖/扩展（逗号分隔）。
_DEFAULT_HEAVY = {
    "torch", "torchvision", "torchaudio", "tensorflow", "jax",
    "vllm", "unsloth", "deepspeed", "flash-attn", "bitsandbytes",
    "llama-cpp-python", "llama_cpp_python", "transformers", "accelerate",
    "peft", "trl", "diffusers", "axolotl", "lm-eval", "lm_eval",
    "stable-diffusion", "comfyui", "openai-whisper", "whisper",
    "segment-anything", "modal",
}

# pip 规格：名字（可带 extras）+ 可选版本约束，拒绝路径/URL（注入面）
_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
_SPEC_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*(\[[^\]]+\])?"
                      r"(\s*(==|>=|<=|~=|!=|>|<)\s*[A-Za-z0-9.*+_-]+)*$")

# Python 标准库模块：技能 frontmatter 里常把它们列进 dependencies，
# 但它们不是 pip 包，装会失败（且 venv 里自带）。
_STDLIB_SKIP = {
    "asyncio", "typing", "json", "os", "sys", "re", "math", "time",
    "datetime", "subprocess", "threading", "multiprocessing", "sqlite3",
    "urllib", "http", "socket", "ssl", "hashlib", "hmac", "base64",
    "uuid", "random", "collections", "itertools", "functools", "pathlib",
    "shutil", "tempfile", "logging", "unittest", "argparse", "csv",
    "struct", "ctypes", "dataclasses", "abc", "io", "string", "traceback",
    "inspect", "pickle", "gzip", "zipfile", "tarfile", "email", "xml",
    "html", "http.client", "socketserver", "secrets", "contextlib",
    "operator", "copy", "weakref", "enum", "numbers", "decimal", "fractions",
}


# ── frontmatter 依赖解析（不引 yaml，手写轻量解析）──────────────

def _parse_skill_deps(content: str):
    """解析 SKILL.md frontmatter 的 dependencies: 列表。

    支持两种写法：块式（- item 多行）与内联（[a, b]）。
    只在 frontmatter 区域内解析，避免正文的散文列表混入。
    返回依赖字符串列表。
    """
    if not content.startswith("---"):
        return []
    end = content.find("\n---", 3)
    if end == -1:
        return []
    fm = content[3:end]
    lines = fm.split("\n")
    deps, in_deps = [], False
    for line in lines:
        stripped = line.strip()
        if stripped.startswith("dependencies:"):
            rest = stripped.split(":", 1)[1].strip()
            if rest.startswith("["):
                inner = rest.strip("[]")
                return [x.strip().strip("'\"") for x in inner.split(",") if x.strip()]
            in_deps = True
            continue
        if in_deps:
            if stripped.startswith("- "):
                deps.append(stripped[2:].strip().strip("'\""))
            elif stripped == "":
                continue
            else:
                in_deps = False
    return deps


def _pip_installable(dep: str) -> bool:
    """判断依赖是否为机器可解析的 pip 规格（名字合法 + 规格形态合法 + 非标准库）。"""
    dep = (dep or "").strip()
    if not dep or "/" in dep or ":" in dep:
        return False
    name = re.split(r"[\[<>=!~\s]", dep, 1)[0]
    if name.lower() in _STDLIB_SKIP:
        return False
    return bool(_NAME_RE.match(name)) and bool(_SPEC_RE.match(dep))


def _specs_from_scope(scope_dir: str):
    """收集一个技能目录（含子目录）下所有技能的 pip 依赖规格。"""
    out = {}
    if not os.path.isdir(scope_dir):
        return out
    for root, _dirs, files in os.walk(scope_dir):
        if "SKILL.md" not in files:
            continue
        try:
            with open(os.path.join(root, "SKILL.md"), encoding="utf-8") as f:
                content = f.read()
        except Exception:
            continue
        for dep in _parse_skill_deps(content):
            if _pip_installable(dep):
                name = re.split(r"[\[<>=!~\s]", dep, 1)[0].lower()
                # 带版本约束的规格优先于裸名字
                if name not in out or any(c in dep for c in "=<>~"):
                    out[name] = dep
    return out


# ── 配置（appbase params，禁硬编码）──────────────────────────

def _cfg_defaults() -> dict:
    return {"pip_index": "https://mirrors.aliyun.com/pypi/simple/",
            "heavy_skip": ",".join(sorted(_DEFAULT_HEAVY)),
            "install_timeout": 1200}


async def _load_cfg(sor=None) -> dict:
    cfg = dict(_cfg_defaults())
    if sor is None:
        return cfg
    try:
        from .workspace import get_param
        idx = await get_param(sor, "remote_env_pip_index", "")
        if idx:
            cfg["pip_index"] = idx.strip()
        heavy = await get_param(sor, "remote_env_heavy_skip", "")
        if heavy:
            cfg["heavy_skip"] = heavy.strip()
        timeout = await get_param(sor, "remote_env_install_timeout", "")
        if timeout:
            cfg["install_timeout"] = int(float(timeout))
    except Exception as e:
        logger.warning("remote_env_init: 读取 appbase 参数失败，用默认值: %s", e)
    return cfg


# ── 技能同步 ──────────────────────────────────────────────────

def _skill_scopes(skills_base: str, org_id: str):
    """要同步的技能范围（相对技能根的本地子目录）列表。

    动态枚举技能根顶层目录（除 orgs/ 外全部同步，覆盖公共/产线/项目/技能集
    及未来新增的范围），再加本机构组织技能（含其下项目私有技能）。
    orgs/ 其他机构不同步（隔离）。
    """
    scopes = []
    if os.path.isdir(skills_base):
        for name in sorted(os.listdir(skills_base)):
            if name == "orgs" or name.startswith("."):
                continue
            if os.path.isdir(os.path.join(skills_base, name)):
                scopes.append(name)
    if org_id:
        scopes.append(os.path.join("orgs", str(org_id)))
    return scopes


async def sync_skills_to_remote(env: dict, org_id: str = "", skills_base: str = "") -> dict:
    """把公共/产线/项目/技能集/本机构组织技能同步到远程 ~/<remote_dir>/skills/。

    幂等：--delete 覆盖。返回 {ok, synced:[...], skipped:[...], error}
    """
    err = _validate_remote_dir(env.get("remote_dir", ""))
    if err:
        return {"ok": False, "error": err}
    if not skills_base:
        from pipeline_core.skill_pack import get_skills_base
        skills_base = get_skills_base()
    if not os.path.isdir(skills_base):
        return {"ok": False, "error": f"本地技能树不存在: {skills_base}"}

    rel = _rel_remote_dir(env.get("remote_dir", ""))
    remote_skills = rel + "/skills"
    if not await _remote_mkdir(env, remote_skills):
        return {"ok": False, "error": "远程创建技能目录失败"}

    target_base = _ssh_target(env) + ":~/" + remote_skills
    ssh_opts = " ".join(_ssh_common_args(env))
    synced, skipped = [], []
    for sub in _skill_scopes(skills_base, org_id):
        src = os.path.join(skills_base, sub)
        if not os.path.isdir(src):
            skipped.append(sub)
            continue
        # rsync -R + "./" 锚点：保留相对子路径（如 orgs/<id>），远程结构与本地一致
        cmd = ["rsync", "-az", "--delete", "-R",
               "-e", "ssh " + ssh_opts,
               os.path.join(skills_base, "./" + sub + "/"),
               target_base + "/"]
        try:
            proc = await asyncio.create_subprocess_exec(
                *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
            _out, stderr = await asyncio.wait_for(proc.communicate(), timeout=300)
            if proc.returncode != 0:
                return {"ok": False, "error": f"技能同步失败({sub}): "
                        + stderr.decode('utf-8', 'replace')[-500:]}
            synced.append(sub)
        except asyncio.TimeoutError:
            return {"ok": False, "error": f"技能同步超时({sub})"}
    return {"ok": True, "synced": synced, "skipped": skipped,
            "remote_path": "~/" + remote_skills}


# ── 依赖收集与安装 ────────────────────────────────────────────

def collect_skill_deps(skills_base: str, org_id: str, heavy_skip: str) -> dict:
    """收集各范围技能的 pip 依赖，剔除重量级。返回 {specs, heavy_skipped, unparseable}。"""
    skip = {s.strip().lower() for s in (heavy_skip or "").split(",") if s.strip()}
    specs, heavy_skipped = {}, []
    for sub in _skill_scopes(skills_base, org_id):
        for name, spec in _specs_from_scope(os.path.join(skills_base, sub)).items():
            if name in skip:
                if spec not in heavy_skipped:
                    heavy_skipped.append(spec)
            else:
                specs[name] = spec
    return {"specs": sorted(specs.values()), "heavy_skipped": sorted(heavy_skipped)}


async def launch_remote_install(env: dict, specs: list, pip_index: str, rel: str) -> bool:
    """在远程以 nohup 分离进程启动逐包安装（脱离 SSH 会话，不怕断连/重启）。

    进度写 ~/<rel>/.venv-skills-install.progress（OK:/FAILED: 行），
    完成写 ~/<rel>/.venv-skills-install.done（汇总）。
    本调用秒级返回，是否成功启动由返回值给出。
    """
    payload = base64.b64encode(json.dumps(specs).encode()).decode()
    index_arg = ('-i ' + shlex.quote(pip_index) + ' ') if pip_index else ''
    # 分离脚本：建 venv → 逐包安装 → 写进度/完成标记。
    # 用 set -f 防 pip 规格里 * / [...] 被远程 shell 展开。
    inner = (
        'RD="$HOME/' + rel + '"; VENV="$RD/.venv-skills"; '
        'PROG="$RD/.venv-skills-install.progress"; DONE="$RD/.venv-skills-install.done"; '
        'if [ ! -x "$VENV/bin/python" ]; then '
        'python3 -m venv "$VENV" || { echo "VENV_FAIL" > "$DONE"; exit 4; }; fi; '
        'set -f; : > "$PROG"; '
        'echo "' + payload + '" | base64 -d | python3 -c '
        '"import json,sys;print(chr(10).join(json.load(sys.stdin)))" | '
        'while IFS= read -r pkg; do '
        '[ -n "$pkg" ] || continue; '
        'if "$VENV/bin/pip" install --no-input --disable-pip-version-check '
        + index_arg + '"$pkg" >/dev/null 2>&1; then echo "OK:$pkg" >> "$PROG"; '
        'else echo "FAILED:$pkg" >> "$PROG"; fi; done; '
        'OKC=$(grep -c "^OK:" "$PROG" 2>/dev/null || echo 0); '
        'FC=$(grep -c "^FAILED:" "$PROG" 2>/dev/null || echo 0); '
        'echo "OK=$OKC FAILED=$FC" > "$DONE"'
    )
    script = (
        'RD="$HOME/' + rel + '"; '
        'command -v python3 >/dev/null || { echo "NO_PYTHON3"; exit 3; }; '
        'mkdir -p "$RD"; '
        'nohup bash -c ' + shlex.quote(inner) +
        ' </dev/null >"$RD/.venv-skills-install.log" 2>&1 & echo LAUNCHED'
    )
    ssh_cmd = ["ssh"] + _ssh_common_args(env) + [_ssh_target(env), 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)
    except Exception as e:
        logger.warning("launch_remote_install 启动失败: %s", e)
        return False
    out = stdout.decode("utf-8", "replace")
    if "NO_PYTHON3" in out:
        return False
    return "LAUNCHED" in out


async def poll_remote_install(env: dict, rel: str) -> dict:
    """轮询远程安装进度。短 SSH 调用。

    返回 {done:bool, ok:bool, installed:[], failed:[], venv, ...}
    """
    script = (
        'RD="$HOME/' + rel + '"; DONE="$RD/.venv-skills-install.done"; '
        'PROG="$RD/.venv-skills-install.progress"; '
        'if [ -f "$DONE" ]; then echo "STATE=done"; cat "$DONE"; echo "===PROG==="; '
        'cat "$PROG" 2>/dev/null; '
        'else echo "STATE=running"; cat "$PROG" 2>/dev/null; fi'
    )
    ssh_cmd = ["ssh"] + _ssh_common_args(env) + [_ssh_target(env), 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=30)
    except Exception as e:
        return {"done": False, "error": "轮询失败: " + str(e)[:200]}

    out = stdout.decode("utf-8", "replace")
    done = "STATE=done" in out
    body = out.split("STATE=", 1)[1] if "STATE=" in out else out
    if "===PROG===" in body:
        summary, prog = body.split("===PROG===", 1)
    else:
        summary, prog = body, ""

    installed, failed = [], []
    for line in prog.splitlines():
        line = line.strip()
        if line.startswith("OK:"):
            installed.append(line[3:])
        elif line.startswith("FAILED:"):
            failed.append(line[7:])

    result = {
        "done": done,
        "ok": done and len(failed) == 0 and len(installed) > 0,
        "installed": installed,
        "failed": failed,
        "venv": "~/" + rel + "/.venv-skills",
        "mode": "per-package-detached",
    }
    if done:
        result["summary"] = summary.strip()
    return result


# ── 总编排 ────────────────────────────────────────────────────

async def init_remote_env(env: dict, org_id: str = "", sor=None,
                           skills_base: str = "", skip_bwrap: bool = False) -> dict:
    """远程环境一键初始化：bwrap → 目录 → 技能同步 → 依赖安装。

    bwrap 缺失即中止（沙箱执行的前提）；技能/依赖失败以报告返回不中止后续步骤。
    """
    from .work_env import ensure_remote_bwrap
    report = {"ok": True, "steps": {}}

    # 1. bwrap（前提，失败中止）
    if not skip_bwrap:
        bw = await ensure_remote_bwrap(env)
        report["steps"]["bwrap"] = bw
        if not bw.get("ok"):
            report["ok"] = False
            report["error"] = "远程 bwrap 部署失败: " + str(bw.get("error", ""))
            return report

    # 2. 远程目录
    rel = _rel_remote_dir(env.get("remote_dir", ""))
    if rel:
        ok = await _remote_mkdir(env, rel)
        report["steps"]["mkdir"] = {"ok": ok}
        if not ok:
            report["ok"] = False
            report["error"] = "远程工作目录创建失败"
            return report

    # 3. 技能同步
    sk = await sync_skills_to_remote(env, org_id, skills_base)
    report["steps"]["skills"] = sk
    if not sk.get("ok"):
        report["ok"] = False
        report["error"] = "技能同步失败: " + str(sk.get("error", ""))
        return report

    # 4. 依赖解析（快）+ 启动远程分离式安装（nohup，脱离 SSH 会话，
    #    进度落远程文件，轮询无状态——服务重启/断连都不影响）
    cfg = await _load_cfg(sor)
    deps = collect_skill_deps(skills_base or _default_skills_base(), org_id,
                              cfg["heavy_skip"])
    report["steps"]["deps_found"] = deps
    launched = await launch_remote_install(env, deps["specs"], cfg["pip_index"], rel)
    report["steps"]["deps_launch"] = {"ok": launched, "count": len(deps["specs"])}
    if not launched:
        report["warning"] = "依赖安装进程启动失败（远程可能缺 python3），请重新运行初始化"
    return report


# ── 依赖安装状态查询（无状态，轮询远程进度文件）────────────────
# 安装进程由远程 nohup 托管，进度写 ~/<remote_dir>/.venv-skills-install.*。
# 查询只需 SSH 读文件，不依赖本进程内存，服务重启后仍可查。

async def get_deps_install_status(sor, owner_type: str, owner_id: str) -> dict:
    """查询远程依赖安装进度。从 sd_work_envs 取远程配置后轮询进度文件。"""
    from .work_env import get_work_env
    env = await get_work_env(
        sor, owner_id if owner_type == "user" else "",
        owner_id if owner_type == "org" else "")
    if env.get("mode") != "remote" or not env.get("remote_dir"):
        return {"ok": False, "error": "当前非远程模式，无安装任务"}
    rel = _rel_remote_dir(env.get("remote_dir", ""))
    return await poll_remote_install(env, rel)


def _default_skills_base() -> str:
    from pipeline_core.skill_pack import get_skills_base
    return get_skills_base()