Video-Tree-TRM5/app/harness/workspace.py

"""Workspace 生命周期管理 + manifest 读写 + Protocol 实现。

Workspace 是一次实验的独立工作区，通过 manifest.json 引用 Store 中的
特定版本资源并记录实验过程。Skills/Prompts 权重拷入 workspace 本地，
训练产物只进 workspace 不污染 Store。

VersionedSkillStore / VersionedPromptStore 实现 core/evolution/protocols.py
中定义的只读端口，供 core/ 层以 Protocol 方式读取技能和提示词。
"""

from __future__ import annotations

import json
import os
import shutil
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import TYPE_CHECKING

from loguru import logger

from app.harness.store import read_seed

if TYPE_CHECKING:
    from pathlib import Path


@dataclass(frozen=True)
class ResolvedPaths:
    """manifest 解析后的绝对路径集合。

    属性:
        store_dir: Store 根目录绝对路径。
        videos_dir: 视频数据目录。
        questions_dir: 当前引用的题目目录。
        skills_dir: 当前引用的 Skill 版本目录（workspace 内）。
        prompts_dir: 当前引用的 Prompt 版本目录（workspace 内）。
        workspace_dir: Workspace 根目录。
        db_path: harness.db 路径。
        analyses_dir: 分析报告目录。
        runs_dir: 运行临时状态目录。
    """

    store_dir: Path
    videos_dir: Path
    questions_dir: Path
    skills_dir: Path
    prompts_dir: Path
    workspace_dir: Path
    db_path: Path
    analyses_dir: Path
    runs_dir: Path


# ---------------------------------------------------------------------------
# 内部工具
# ---------------------------------------------------------------------------

_MANIFEST_CURRENT_KEYS = {"videos", "questions", "skills", "prompts"}


def _now_iso() -> str:
    """返回当前 UTC 时间的 ISO 格式字符串。"""
    return datetime.now(UTC).isoformat()


# ---------------------------------------------------------------------------
# Workspace 核心函数
# ---------------------------------------------------------------------------


def _scaffold_workspace(
    workspace_dir: Path,
    store_dir: Path,
    questions: str,
    skills_version: str,
    prompts_version: str,
) -> None:
    """写 manifest + 建 analyses/runs 目录（不拷权重；权重由调用方按来源拷入）。

    参数:
        workspace_dir: 目标 workspace（由调用方保证不存在）。
        store_dir: Store 根目录。
        questions: 题目相对路径，如 ``'benchmarks/Video-MME'``。
        skills_version: manifest.current.skills 初始版本号。
        prompts_version: manifest.current.prompts 初始版本号。

    关键实现:
        不依赖任何外部资源源（store 中的 skills/prompts 是否存在不在此校验），
        因此可被 init_workspace 与种子初始化复用；store 引用以相对路径写入 manifest。
    """
    workspace_dir.mkdir(parents=True)
    (workspace_dir / "analyses").mkdir()
    (workspace_dir / "runs").mkdir()

    store_abs = store_dir.resolve()
    store_rel = os.path.relpath(store_abs, workspace_dir.resolve())

    manifest = {
        "name": workspace_dir.name,
        "created_at": _now_iso(),
        "store": store_rel,
        "current": {
            "videos": "videos",
            "questions": f"questions/{questions}",
            "skills": f"skills/{skills_version}",
            "prompts": f"prompts/{prompts_version}",
        },
        "history": [],
    }
    (workspace_dir / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2))


def init_workspace(
    workspace_dir: Path,
    store_dir: Path,
    questions: str,
    skills_version: str,
    prompts_version: str,
) -> None:
    """创建 Workspace 目录并写入初始 manifest.json，拷贝种子权重。

    参数:
        workspace_dir: Workspace 目标路径（不得已存在）。
        store_dir: Store 根目录。
        questions: 题目在 questions/ 下的相对路径，如 ``"benchmarks/Video-MME"``。
        skills_version: Skills 版本号，如 ``"v1"``。
        prompts_version: Prompts 版本号，如 ``"v1"``。

    异常:
        FileExistsError: Workspace 目录已存在。
        FileNotFoundError: 引用的资源在 Store 中不存在。
    """
    if workspace_dir.exists():
        raise FileExistsError(f"Workspace 已存在: {workspace_dir}")

    store_abs = store_dir.resolve()
    refs = {
        "skills": f"skills/{skills_version}",
        "prompts": f"prompts/{prompts_version}",
        "questions": f"questions/{questions}",
    }
    for label, rel in refs.items():
        full = store_abs / rel
        if not full.is_dir():
            raise FileNotFoundError(f"Store 中不存在 {label}: {full}")

    _scaffold_workspace(workspace_dir, store_dir, questions, skills_version, prompts_version)

    # 拷种子权重进 workspace：v2+ 训练产物只进 workspace，不污染 store
    shutil.copytree(store_abs / refs["skills"], workspace_dir / refs["skills"])
    shutil.copytree(store_abs / refs["prompts"], workspace_dir / refs["prompts"])
    logger.info("Workspace 初始化完成: {}", workspace_dir)


def init_workspace_from_seed(
    workspace_dir: Path,
    store_dir: Path,
    seed_name: str,
    questions: str,
) -> str:
    """从种子全新建 workspace：拷权重 -> v1、baseline.db -> harness.db、读 baseline_run_id。

    参数:
        workspace_dir: 目标 workspace（不得已存在）。
        store_dir: Store 根目录。
        seed_name: 种子名（store/seeds 下）。
        questions: 题目相对路径，如 ``'benchmarks/Video-MME'``。

    返回:
        baseline_run_id（供 build_pools 使用）。

    异常:
        FileExistsError: workspace 已存在。
        FileNotFoundError: 种子不存在（由 read_seed 抛出），或 questions ref 目录不存在。

    关键实现:
        破坏性/创建操作前先校验 questions ref 存在：fresh 路径在 runner 侧已先
        归档旧 ws，若到 build_pools 才发现 questions 缺失则旧 ws 已被毁；
        故在此尽早报错（fail-fast），让新 ws 在创建前失败。
    """
    if workspace_dir.exists():
        raise FileExistsError(f"Workspace 已存在: {workspace_dir}")

    # 校验种子存在 + 取 baseline_run_id
    meta = read_seed(store_dir, seed_name)

    # fail-fast：校验 questions ref 存在
    questions_ref = store_dir / "questions" / questions
    if not questions_ref.is_dir():
        raise FileNotFoundError(f"questions ref 目录不存在: {questions_ref}")

    seed_dir = store_dir / "seeds" / seed_name
    _scaffold_workspace(workspace_dir, store_dir, questions, "v1", "v1")
    shutil.copytree(seed_dir / "skills", workspace_dir / "skills" / "v1")
    shutil.copytree(seed_dir / "prompts", workspace_dir / "prompts" / "v1")
    shutil.copy2(seed_dir / "baseline.db", workspace_dir / "harness.db")

    logger.info("Workspace 从种子 '{}' 初始化完成: {}", seed_name, workspace_dir)
    return meta["baseline_run_id"]


def load_manifest(workspace_dir: Path) -> dict:
    """读取并返回 workspace 的 manifest.json。

    参数:
        workspace_dir: Workspace 根目录。

    返回:
        manifest 字典。

    异常:
        FileNotFoundError: manifest.json 不存在。
    """
    manifest_path = workspace_dir / "manifest.json"
    if not manifest_path.exists():
        raise FileNotFoundError(f"manifest.json 不存在: {manifest_path}")
    return json.loads(manifest_path.read_text())


def resolve_paths(workspace_dir: Path) -> ResolvedPaths:
    """读取 manifest，解析 current 中所有资源的绝对路径。

    skills_dir/prompts_dir 解析到 workspace（非 store），
    videos_dir/questions_dir 解析到 store。

    参数:
        workspace_dir: Workspace 根目录。

    返回:
        ResolvedPaths 实例，包含所有资源的绝对路径。
    """
    manifest = load_manifest(workspace_dir)
    ws_abs = workspace_dir.resolve()
    store_abs = (ws_abs / manifest["store"]).resolve()
    current = manifest["current"]
    return ResolvedPaths(
        store_dir=store_abs,
        videos_dir=store_abs / current["videos"],
        questions_dir=store_abs / current["questions"],
        skills_dir=ws_abs / current["skills"],
        prompts_dir=ws_abs / current["prompts"],
        workspace_dir=ws_abs,
        db_path=ws_abs / "harness.db",
        analyses_dir=ws_abs / "analyses",
        runs_dir=ws_abs / "runs",
    )


def list_video_ids(workspace_dir: Path) -> list[str]:
    """列出 workspace 引用的所有视频 ID（含 tree.json 的子目录名）。

    参数:
        workspace_dir: Workspace 根目录。

    返回:
        排序后的视频 ID 列表。
    """
    paths = resolve_paths(workspace_dir)
    video_ids = []
    for entry in paths.videos_dir.iterdir():
        if entry.is_dir() and (entry / "tree.json").exists():
            video_ids.append(entry.name)
    return sorted(video_ids)


def update_manifest(workspace_dir: Path, **version_updates: str) -> None:
    """更新 manifest 的 current 字段。

    参数:
        workspace_dir: Workspace 根目录。
        **version_updates: 要更新的字段及其新值，如 ``skills="skills/v2"``。

    异常:
        KeyError: 更新的字段不在 current 允许的 key 白名单中。
    """
    invalid = set(version_updates) - _MANIFEST_CURRENT_KEYS
    if invalid:
        raise KeyError(f"无效的 manifest current 字段: {invalid}")
    manifest = load_manifest(workspace_dir)
    manifest["current"].update(version_updates)
    (workspace_dir / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2))


def record_run(workspace_dir: Path, run_id: str) -> Path:
    """将 current 版本快照追加到 manifest history，创建 run 目录和 per-video wiki 目录。

    幂等：同 run_id 不重复追加 history（长跑中断后重启 / held-out 复用 run_id 时）。

    参数:
        workspace_dir: Workspace 根目录。
        run_id: 本次运行的唯一标识，如 ``"run_001"``。

    返回:
        创建的 run 目录路径。
    """
    manifest = load_manifest(workspace_dir)
    current = manifest["current"]

    # 幂等：同 run_id 不重复追加 history
    if not any(h["run_id"] == run_id for h in manifest["history"]):
        manifest["history"].append(
            {
                "run_id": run_id,
                "started_at": _now_iso(),
                "skills": current["skills"],
                "prompts": current["prompts"],
                "questions": current["questions"],
            }
        )
        (workspace_dir / "manifest.json").write_text(
            json.dumps(manifest, ensure_ascii=False, indent=2)
        )

    run_dir = workspace_dir / "runs" / run_id
    # exist_ok：同 run_id 重跑时 run 目录已存在不应崩溃
    run_dir.mkdir(parents=True, exist_ok=True)
    for video_id in list_video_ids(workspace_dir):
        (run_dir / video_id / "wiki").mkdir(parents=True, exist_ok=True)

    logger.debug("Run 已记录: {}", run_id)
    return run_dir


def read_best(workspace_dir: Path) -> dict | None:
    """读取 manifest 的 best 指针，未设置时返回 None。

    参数:
        workspace_dir: Workspace 根目录。

    返回:
        best 字典（skills/prompts/val_acc/run_id/epoch），未设置时 None。
    """
    return load_manifest(workspace_dir).get("best")


def update_best(
    workspace_dir: Path,
    skills: str,
    prompts: str,
    val_acc: float,
    run_id: str,
    epoch: int,
) -> None:
    """写入 manifest 的 best 指针（历史最优版本快照，与 current 平级）。

    best 独立于 current——更新 best 不影响 current。

    参数:
        workspace_dir: Workspace 根目录。
        skills: 最优 skills 版本完整 ref，如 ``'skills/v2'``。
        prompts: 最优 prompts 版本完整 ref，如 ``'prompts/v2'``。
        val_acc: 该版本验证集准确率。
        run_id: 该版本验证 run_id。
        epoch: 达成该最优的轮次。
    """
    manifest = load_manifest(workspace_dir)
    manifest["best"] = {
        "skills": skills,
        "prompts": prompts,
        "val_acc": val_acc,
        "run_id": run_id,
        "epoch": epoch,
    }
    (workspace_dir / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2))
    logger.info("Best 已更新: val_acc={}, run={}, epoch={}", val_acc, run_id, epoch)


def archive_workspace(workspace_dir: Path) -> Path:
    """把 workspace 整体移动到同级 .archive/<name>-<ts>，返回归档路径。

    参数:
        workspace_dir: 要归档的 Workspace 根目录。

    返回:
        归档后的目标路径。

    异常:
        FileNotFoundError: workspace 不存在。
    """
    if not workspace_dir.exists():
        raise FileNotFoundError(f"workspace 不存在: {workspace_dir}")

    archive_root = workspace_dir.parent / ".archive"
    archive_root.mkdir(exist_ok=True)
    ts = datetime.now(UTC).strftime("%Y%m%d-%H%M%S")
    target = archive_root / f"{workspace_dir.name}-{ts}"
    shutil.move(str(workspace_dir), str(target))

    logger.info("Workspace 已归档: {} -> {}", workspace_dir, target)
    return target


# ---------------------------------------------------------------------------
# Protocol 实现：VersionedSkillStore / VersionedPromptStore
# ---------------------------------------------------------------------------


class VersionedSkillStore:
    """版本化技能读取端口实现。

    满足 ``core/evolution/protocols.py::SkillStore`` Protocol。
    从指定的 skills 版本目录读取 ``.md`` 文件。

    参数:
        skills_dir: skills 版本目录绝对路径（如 ``workspace/skills/v1``）。
    """

    def __init__(self, skills_dir: Path) -> None:
        if not skills_dir.is_dir():
            raise FileNotFoundError(f"Skills 目录不存在: {skills_dir}")
        self._dir = skills_dir

    def read_skill(self, filename: str) -> str:
        """读取指定 skill 文件的全文内容。

        参数:
            filename: skill 文件名，如 ``'temporal-reasoning.md'``。

        返回:
            文件全文内容。

        异常:
            FileNotFoundError: 文件不存在。
        """
        path = self._dir / filename
        if not path.exists():
            raise FileNotFoundError(f"Skill 文件不存在: {path}")
        return path.read_text()

    def list_skill_files(self) -> list[str]:
        """列出当前版本所有 skill 文件名。

        返回:
            文件名列表（排序）。
        """
        return sorted(entry.name for entry in self._dir.iterdir() if entry.is_file())


class VersionedPromptStore:
    """版本化提示词读取端口实现。

    满足 ``core/evolution/protocols.py::PromptStore`` Protocol。
    从指定的 prompts 版本目录读取 ``.md`` 文件。

    参数:
        prompts_dir: prompts 版本目录绝对路径（如 ``workspace/prompts/v1``）。
    """

    def __init__(self, prompts_dir: Path) -> None:
        if not prompts_dir.is_dir():
            raise FileNotFoundError(f"Prompts 目录不存在: {prompts_dir}")
        self._dir = prompts_dir

    def read_prompt(self, filename: str) -> str:
        """读取指定 prompt 文件的全文内容。

        参数:
            filename: prompt 文件名，如 ``'system.md'``。

        返回:
            文件全文内容。

        异常:
            FileNotFoundError: 文件不存在。
        """
        path = self._dir / filename
        if not path.exists():
            raise FileNotFoundError(f"Prompt 文件不存在: {path}")
        return path.read_text()

    def list_prompt_files(self) -> list[str]:
        """列出当前版本所有 prompt 文件名。

        返回:
            文件名列表（排序）。
        """
        return sorted(entry.name for entry in self._dir.iterdir() if entry.is_file())