"""远程工作环境自动初始化：技能同步 + 依赖软件自动安装（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 install_remote_deps(env: dict, specs: list, pip_index: str,
                              timeout: int = 1200) -> dict:
    """在远程 ~/<remote_dir>/.venv-skills 建 venv 并安装依赖。幂等。

    specs 经 base64 传输（含 >=[] 等特殊字符，避开 shell 展开）。
    返回 {ok, installed:[...], failed:[...], error}
    """
    if not specs:
        return {"ok": True, "installed": [], "failed": [], "skipped_reason": "无可安装依赖"}
    err = _validate_remote_dir(env.get("remote_dir", ""))
    if err:
        return {"ok": False, "error": err}

    rel = _rel_remote_dir(env.get("remote_dir", ""))
    payload = base64.b64encode(json.dumps(specs).encode()).decode()
    # 不用 set -e / 管道：管道会吞掉 pip 的退出码（$? 变成 tail 的状态）。
    # set -f 关 glob：pip 规格里的 * / [...] 不能被远程 shell 展开。
    # 完整输出由 SSH 捕获，报告时截取尾部。
    script = (
        'RD="$HOME/' + rel + '"; VENV="$RD/.venv-skills"; '
        'command -v python3 >/dev/null || { echo "NO_PYTHON3"; exit 3; }; '
        'if [ ! -x "$VENV/bin/python" ]; then python3 -m venv "$VENV" || { echo "VENV_FAIL"; exit 4; }; fi; '
        'set -f; RC=0; "$VENV/bin/pip" install --no-input --disable-pip-version-check '
        + (('-i ' + shlex.quote(pip_index) + ' ') if pip_index else '')
        + '$(echo "' + payload + '" | base64 -d | python3 -c '
        '"import json,sys;print(chr(32).join(json.load(sys.stdin)))") '
        '|| RC=$?; '
        'echo "INSTALL_RC=$RC"'
    )
    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=timeout + 30)
    except asyncio.TimeoutError:
        return {"ok": False, "error": f"依赖安装超时（{timeout}s），venv 可能装了一半，可重跑"}
    except Exception as e:
        return {"ok": False, "error": "SSH 执行失败: " + str(e)[:300]}

    out = stdout.decode("utf-8", "replace")
    if "NO_PYTHON3" in out:
        return {"ok": False, "error": "远程主机缺少 python3"}
    if proc.returncode not in (0, 1):
        return {"ok": False, "error": "venv 创建失败: "
                + (out + stderr.decode('utf-8', 'replace'))[-600:]}

    # 解析 pip 退出码（由脚本末尾 echo INSTALL_RC=$RC 输出）
    rc = -1
    for line in out.splitlines():
        if line.startswith("INSTALL_RC="):
            try:
                rc = int(line.split("=", 1)[1])
            except Exception:
                pass

    # pip 逐包结果解析：新版 pip 解析器失败时要么全装要么全不装，
    # 从 ERROR 行提取包名。常见格式：
    #   ERROR: Could not find a version that satisfies the requirement <name> ...
    #   ERROR: No matching distribution found for <name>
    failed = []
    for line in out.splitlines():
        if "ERROR:" not in line:
            continue
        m = (re.search(r"the requirement ([A-Za-z0-9._\[\]-]+)", line)
             or re.search(r"No matching distribution found for ([A-Za-z0-9._\[\]-]+)", line))
        failed.append(m.group(1) if m else line.strip()[:120])
    if rc != 0 and not failed:
        failed = ["部分包安装失败，详见日志末尾"]

    if rc == 0:
        installed = list(specs)
        failed = []
    else:
        # 批量失败 → 逐包降级：一次 SSH 会话内逐包安装（pip 幂等，已装的秒过），
        # 部分成功优于全败。
        return await _install_per_package(env, specs, pip_index, rel, timeout)
    return {"ok": rc == 0, "installed": installed, "failed": failed,
            "venv": "~/" + rel + "/.venv-skills",
            "tail": out[-800:]}


async def _install_per_package(env: dict, specs: list, pip_index: str,
                               rel: str, timeout: int) -> dict:
    """逐包安装（批量失败后的降级路径）。返回格式与 install_remote_deps 一致。"""
    payload = base64.b64encode(json.dumps(specs).encode()).decode()
    index_arg = ('-i ' + shlex.quote(pip_index) + ' ') if pip_index else ''
    script = (
        'RD="$HOME/' + rel + '"; VENV="$RD/.venv-skills"; '
        '[ -x "$VENV/bin/pip" ] || { echo "VENV_MISSING"; exit 4; }; '
        'set -f; '
        '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"; '
        '  else echo "FAILED:$pkg"; fi; '
        'done'
    )
    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=timeout + 30)
    except asyncio.TimeoutError:
        return {"ok": False, "error": f"逐包安装超时（{timeout}s），可重跑"}
    except Exception as e:
        return {"ok": False, "error": "SSH 执行失败: " + str(e)[:300]}

    out = stdout.decode("utf-8", "replace")
    installed, failed = [], []
    for line in out.splitlines():
        if line.startswith("OK:"):
            installed.append(line[3:])
        elif line.startswith("FAILED:"):
            failed.append(line[7:])
    return {"ok": len(failed) == 0 and len(installed) > 0,
            "installed": installed, "failed": failed,
            "venv": "~/" + rel + "/.venv-skills",
            "mode": "per-package", "tail": out[-500:]}


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

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. 依赖解析（快）+ 安装（慢，后台执行）
    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
    task_key = _start_deps_install(env, deps["specs"], cfg)
    report["deps_task"] = task_key
    report["steps"]["deps_install"] = {"status": "running", "task": task_key}
    return report


# ── 依赖安装后台任务与状态 ────────────────────────────────────
# 依赖安装可能耗时数分钟（超过 nginx 600s 代理超时），不能在 API 请求内同步等。
# 编排只启动后台任务并返回任务号，前端用 init_remote_env_status 轮询结果。
# 状态存进程内存（单进程部署；重启丢状态则重跑，幂等无副作用）。

_INSTALL_TASKS = {}  # {task_key: {"status": ..., "result": ..., "started": ts}}


def _start_deps_install(env: dict, specs: list, cfg: dict) -> str:
    from appPublic.uniqueID import getID
    task_key = "deps_" + getID()
    _INSTALL_TASKS[task_key] = {"status": "running", "result": None,
                                "specs": list(specs)}

    async def _run():
        try:
            r = await install_remote_deps(env, specs, cfg["pip_index"],
                                          cfg["install_timeout"])
            _INSTALL_TASKS[task_key]["result"] = r
            _INSTALL_TASKS[task_key]["status"] = "done"
        except Exception as e:
            _INSTALL_TASKS[task_key]["result"] = {"ok": False, "error": str(e)[:500]}
            _INSTALL_TASKS[task_key]["status"] = "done"

    asyncio.create_task(_run())
    return task_key


def get_deps_install_status(task_key: str) -> dict:
    t = _INSTALL_TASKS.get(task_key or "")
    if not t:
        return {"ok": False, "error": "任务不存在（可能服务重启过，请重新初始化）"}
    if t["status"] == "running":
        return {"ok": True, "status": "running",
                "message": f"依赖安装中（{len(t.get('specs', []))} 个包），请稍候再查询"}
    r = dict(t.get("result") or {})
    r["status"] = "done"
    return r


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