# -*- coding: utf-8 -*-
"""
sandbox.py — Pi Agent 兜底能力 P0 PoC：四层围墙骨架
====================================================

对应 `docs/architecture/fallback.md` 的沙箱围墙设计：

- L1 工具层：由 tool_bridge.py 的工具白名单 + orchestrator 启动参数 `--tools read,grep,find,ls`
  实现（pi 只看得见白名单工具，且本 PoC **不装 bash/edit 给真实路径**，见下）；
- L2 文件层：本模块。每次运行创建 runs/<runId>/{inbox,work,outbox} 三区
  （inbox=可读注入区 / work=可写工作区 / outbox=产物出口），并提供 validate_path()
  路径校验（越界即拒绝）；真实模式下 pi 子进程 cwd 圈禁在 work 区；
- L3 网络层：本 PoC 以「不装 bash 工具（无任意命令执行面）+ 模型出口仅经
  $PI_CODING_AGENT_DIR/models.json 里配置的 provider baseUrl」实现——即网络面收敛为
  唯一可配的 LLM 网关地址（对应方案「模型 API 出口走 gateway 代理（可关）」）。
  pi 自身无沙箱（Agent-A 侦察结论），更严格的端口级限制（仅 127.0.0.1:gateway）
  留待产品化阶段用防火墙规则/网关代理实现，此处以注释与配置约束说明对应关系；
- L4 进程层：由 orchestrator.py 实现（build_child_env 环境清洗 + taskkill /T /F 进程树回收）。

同时本模块负责生成 Pi 扩展守卫（TypeScript，参考 probes/guard-ext.ts 模式）：
tool_call 钩子里 bash/edit 默认 block、文件类工具限制在 run 目录内。
"""

from __future__ import annotations

import time
import uuid
from pathlib import Path
from typing import Optional

# ---------------------------------------------------------------------------
# 集中配置：poc 目录布局（一切产出只在 poc/pi-fallback/ 内，GOAL.md 硬约束）
# ---------------------------------------------------------------------------

POC_ROOT = Path(__file__).resolve().parent                 # poc/pi-fallback/
RUNS_ROOT = POC_ROOT / "runs"                              # 每次运行一个子目录
RUNTIME_ROOT = POC_ROOT / "runtime"                        # Agent-A 装好的 pi 运行时
PI_HOME = RUNTIME_ROOT / "pi-home"                         # PI_CODING_AGENT_DIR 指向这里
PI_CLI = (RUNTIME_ROOT / "node_modules" / "@mariozechner"
          / "pi-coding-agent" / "dist" / "cli.js")         # headless 入口（Agent-A 实测）


class SandboxViolation(Exception):
    """路径越界 / 围墙违规时抛出。调用方必须把它当失败处理，不许静默放行。"""


def new_run_id(prefix: str = "run") -> str:
    """生成 runId：时间戳 + 短随机串，贯穿 propose→execute→verify 全链（方案 §4.5）。"""
    return time.strftime(f"{prefix}-%Y%m%d-%H%M%S-") + uuid.uuid4().hex[:6]


def create_run_dirs(run_id: Optional[str] = None) -> dict:
    """
    为一次运行创建 L2 三区目录。

    返回 {"root", "inbox", "work", "outbox"} 四个 Path：
      - inbox/  可读注入区：用户上传文件/演示数据的副本放这里，语义上只读；
      - work/   可写工作区：pi 子进程 cwd 圈禁于此，中间产物写这里；
      - outbox/ 产物出口：最终报告唯一落点（对应方案 §4.3 report_emit 唯一出口）。
    """
    run_id = run_id or new_run_id()
    root = RUNS_ROOT / run_id
    dirs = {
        "root": root,
        "inbox": root / "inbox",
        "work": root / "work",
        "outbox": root / "outbox",
    }
    for p in dirs.values():
        p.mkdir(parents=True, exist_ok=True)
    return dirs


def validate_path(path, run_root: Path, zone: Optional[str] = None,
                  must_exist: bool = False) -> Path:
    """
    L2 文件层核心校验：确认 path 解析后落在 runs/<runId>/ 内，越界即拒绝。

    参数:
        path:       待校验路径（相对或绝对，可含 ../ 等逃逸企图）。
        run_root:   本次运行的根目录（runs/<runId>）。
        zone:       可选，进一步限定必须落在 "inbox"/"work"/"outbox" 某一区。
        must_exist: 为 True 时目标必须已存在（用于读场景）。

    返回解析后的绝对 Path；违规抛 SandboxViolation。
    实现要点：resolve() 归一化 + is_relative_to 前缀判断，挡住 `..`、
    绝对路径跳转、大小写差异（Windows resolve 会规范化盘符大小写）。
    """
    run_root = Path(run_root).resolve()
    candidate = Path(path)
    if not candidate.is_absolute():
        candidate = run_root / candidate
    candidate = candidate.resolve()

    if must_exist and not candidate.exists():
        raise SandboxViolation(f"路径不存在: {candidate}")
    if not candidate.is_relative_to(run_root):
        raise SandboxViolation(f"路径越界（逃出 run 目录）: {candidate} 不在 {run_root} 内")
    if zone:
        zone_root = (run_root / zone).resolve()
        if not candidate.is_relative_to(zone_root):
            raise SandboxViolation(f"路径越区（须在 {zone}/ 内）: {candidate}")
    return candidate


# ---------------------------------------------------------------------------
# Pi 扩展守卫生成（L1 工具层的可编程部分，参考 probes/guard-ext.ts 模式）
# ---------------------------------------------------------------------------

# TypeScript 模板。占位符：{RUN_ROOT_POSIX} {BASH_WHITELIST_JS} {BLOCKLOG_POSIX}
_GUARD_TS_TEMPLATE = """// AUTO-GENERATED by sandbox.py — PoC 守卫扩展（L1/L2 围墙的 pi 侧执行点）。
// 机制依据：NOTES-pi-runtime.md §4 第二道墙 —— pi.on("tool_call") 返回
// {{ block: true, reason }} 即可在工具执行前拦截（Agent-A 已实测钩子加载）。
import fs from "node:fs";
import path from "node:path";

const RUN_ROOT = path.normalize("{RUN_ROOT_POSIX}");
const BLOCKLOG = path.join(RUN_ROOT, "guard-blocked-calls.jsonl");
// bash 白名单命令前缀（正则字符串数组）；空数组 = bash 全禁（本 PoC 默认）。
const BASH_WHITELIST: RegExp[] = [{BASH_WHITELIST_JS}].map((s) => new RegExp(s));

function inRunRoot(p: string): boolean {{
  const abs = path.resolve(process.cwd(), p);
  const norm = path.normalize(abs);
  return norm === RUN_ROOT || norm.startsWith(RUN_ROOT + path.sep);
}}

function deny(toolName: string, toolCallId: string, reason: string, input: any) {{
  fs.appendFileSync(
    BLOCKLOG,
    JSON.stringify({{ ts: new Date().toISOString(), toolName, toolCallId, reason, input }}) + "\\n",
  );
  return {{ block: true, reason }};
}}

export default function (pi: any) {{
  pi.on("tool_call", async (event: any, _ctx: any) => {{
    const name: string = event.toolName;
    const input: any = event.input || {{}};

    // 1) bash / edit：默认 block；仅放行命中白名单前缀的 bash 命令。
    if (name === "edit") {{
      return deny(name, event.toolCallId, "edit disabled by PoC guard (use write into work/ instead)", input);
    }}
    if (name === "bash") {{
      const cmd: string = String(input.command || "");
      const ok = BASH_WHITELIST.some((re) => re.test(cmd));
      if (!ok) return deny(name, event.toolCallId, "bash command not in whitelist", input);
    }}

    // 2) 文件类工具：路径必须落在 run 目录内（L2 圈禁的 pi 侧执行点）。
    const fileTools = ["read", "write", "edit", "grep", "find", "ls"];
    if (fileTools.includes(name)) {{
      const p: string = String(input.path || input.pattern || ".");
      if (path.isAbsolute(p) || p.includes("..")) {{
        if (!inRunRoot(p)) return deny(name, event.toolCallId, "path escapes run root", input);
      }} else if (!inRunRoot(p)) {{
        return deny(name, event.toolCallId, "path escapes run root", input);
      }}
    }}
    // 放行
  }});
}}
"""

# 本 PoC 默认 bash 全禁（网络限制 L3 的前提：没有任意命令执行面）。
# 需要放行时按命令前缀正则登记，例如：r"^python3? \"?[^\x22']*analyze\.py"
DEFAULT_BASH_WHITELIST: tuple = ()


def write_guard_extension(run_root: Path,
                          bash_whitelist=DEFAULT_BASH_WHITELIST) -> Path:
    """
    为本次运行生成守卫扩展 guard-<runId>.ts（落在 run 目录内，随运行归档）。

    真实模式启动 pi 时以 `-e <该文件>` 加载。返回生成文件路径。
    bash_whitelist: bash 命令前缀正则元组；默认空 = bash 全禁。
    """
    run_root = Path(run_root).resolve()
    bash_js = ", ".join(f'{r!r}'.replace("'", '"') for r in bash_whitelist)
    content = _GUARD_TS_TEMPLATE.format(
        RUN_ROOT_POSIX=run_root.as_posix(),
        BASH_WHITELIST_JS=bash_js,
    )
    out = run_root / f"guard-{run_root.name}.ts"
    out.write_text(content, encoding="utf-8")
    return out