"""
engine.sandbox — Isolated agent execution environment.

Level 0: In-process (dev, fastest, no isolation)
Level 1: Subprocess (default, process-level isolation)

Run Workspace:
  Each execution creates a persistent run directory:
    {agent_dir}/workspace/run_{YYYYMMDD_HHMMSS}/
  containing input, output, trace, config snapshot, and cost data.
  Old runs are NEVER deleted — all history is preserved.
"""
from __future__ import annotations
import asyncio
import hashlib
import json
import os
import subprocess
import time
import traceback
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


@dataclass
class ResourceLimits:
    cpu_seconds: int = 300
    memory_mb: int = 512
    max_steps: int = 100
    max_tokens: int = 100_000
    network_enabled: bool = True


@dataclass
class ExecutionResult:
    output: str = ""
    status: str = "completed"
    error: str = ""
    duration_ms: int = 0
    steps: int = 0
    input_tokens: int = 0
    output_tokens: int = 0
    workspace_path: str = ""  # run 工作目录路径


# ============================================================
# Run Workspace 管理
# ============================================================

def create_run_workspace(
    agent_dir: str,
    run_id: str,
    input_text: str,
    config: dict,
    bare: bool = False,
    run_dir: str = "",
) -> str:
    """
    为一次 run 创建持久化工作目录。

    目录结构（bare=False，默认）:
        {agent_dir}/workspace/run_{YYYYMMDD_HHMMSS}/
        ├── input.json      ← 本次输入
        ├── config.yml      ← 使用的配置快照
        ├── code/           ← agent 产生的代码
        ├── results/        ← agent 产生的数据
        └── final/          ← agent 最终产物

    当 bare=True 时（run_dir 场景）:
        {agent_dir}/run_{YYYYMMDD_HHMMSS}/   ← 不再追加 workspace/ 子目录

    Args:
        agent_dir: agent 的根目录（bare=False）或运行日志目录（bare=True）
        run_id: run 的唯一标识 (e.g. "run_a1b2c3d4e5f6")
        input_text: 用户输入
        config: agent YAML 配置字典
        bare: 若 True，不追加 workspace/ 子目录，直接在 agent_dir 下建 run_* 目录
        run_dir: 若提供，直接用作 workspace 绝对/相对路径（不做时间戳命名）。
                 供 PipelineRunner 把每个 stage 目录放到 pipeline 目录下用
                 （绕开时间戳发现，见 PRODUCT_GATE_DESIGN §15 P3-F）。
                 优先级高于 bare/agent_dir。

    Returns:
        workspace_path: 创建的 run 目录绝对路径
    """
    import yaml

    timestamp = time.strftime("%Y%m%d_%H%M%S")
    if run_dir:
        workspace = run_dir
    elif bare:
        workspace = os.path.join(agent_dir, f"run_{timestamp}")
    else:
        workspace = os.path.join(agent_dir, "workspace", f"run_{timestamp}")

    # 如果同一秒内有多次 run（罕见），追加 run_id 后缀避免冲突（run_dir 显式路径不冲突处理）
    if not run_dir and os.path.exists(workspace):
        workspace = f"{workspace}_{run_id[-6:]}"

    os.makedirs(workspace, exist_ok=True)
    os.makedirs(os.path.join(workspace, "code"), exist_ok=True)
    os.makedirs(os.path.join(workspace, "results"), exist_ok=True)
    os.makedirs(os.path.join(workspace, "final"), exist_ok=True)

    # 保存输入
    with open(os.path.join(workspace, "input.json"), "w", encoding="utf-8") as f:
        json.dump({
            "run_id": run_id,
            "input": input_text,
            "timestamp": timestamp,
        }, f, ensure_ascii=False, indent=2)

    # 保存配置快照
    with open(os.path.join(workspace, "config.yml"), "w", encoding="utf-8") as f:
        yaml.dump(config, f, allow_unicode=True, default_flow_style=False)

    return os.path.abspath(workspace)


def _compute_cost_usd(usage: dict) -> float:
    """Estimate USD cost from a token-usage dict.

    Dispatches on usage["provider"] when present so non-Claude providers
    use the correct pricing table.  Falls back to Claude Sonnet 4.6 pricing.

    ── Claude Sonnet 4.6 (default, 2026-06) ──────────────────────────────
      Input:          $3.00 / 1M
      Output:        $15.00 / 1M
      Cache read:     $0.30 / 1M
      Cache create:   $3.75 / 1M

    ── Qwen / DashScope (dashscope, qwen) — 2026-06 approximate USD ──────
      qwen-max   Input: ¥0.04/1K → ~$5.56/M   Output: ¥0.12/1K → ~$16.67/M
      qwen-plus  Input: ¥0.0008/1K → ~$0.11/M  Output: ¥0.002/1K → ~$0.28/M
      qwen-turbo Input: ¥0.0003/1K → ~$0.04/M  Output: ¥0.0006/1K → ~$0.08/M
      (Uses qwen-max rates as conservative default; CNY/USD ≈ 7.2)

    ── Ollama (local / self-hosted) ─────────────────────────────────────
      No per-token API cost → always $0.00.
    """
    provider = usage.get("provider", "claude")

    # Local / self-hosted models incur no per-token API cost.
    if provider in ("ollama",):
        return 0.0

    if provider in ("dashscope", "qwen"):
        # Conservative: use qwen-max rates regardless of actual model tier.
        # $5.56/M input, $16.67/M output (¥0.04 and ¥0.12 per 1K ÷ 7.2)
        return (
            usage.get("input_tokens", 0)  * 5.56  / 1_000_000
            + usage.get("output_tokens", 0) * 16.67 / 1_000_000
        )

    # Default: Claude pricing
    return (
        usage.get("input_tokens", 0)                 * 3.00  / 1_000_000
        + usage.get("output_tokens", 0)              * 15.00 / 1_000_000
        + usage.get("cache_read_input_tokens", 0)    * 0.30  / 1_000_000
        + usage.get("cache_creation_input_tokens", 0) * 3.75 / 1_000_000
    )


# ============================================================
# Change Manifest（修改清单）
# ============================================================

# Framework-written files at the run-dir root — NOT agent-produced changes
_MANIFEST_EXCLUDE = {
    "input.json", "config.yml", "output.json",
    "trace.json", "cost.json", "manifest.json",
}
_MANIFEST_HASH_MAX_BYTES = 100 * 1024 * 1024   # don't hash files larger than this
_GIT_DIFF_MAX_CHARS = 200_000                   # cap diff text stored in manifest


def _sha256_file(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def _git_diff_for(path: str) -> dict:
    """Best-effort git capture for a run.

    If *path* lives inside a git work tree, record modifications to *tracked*
    files (this complements the workspace file scan, which already covers newly
    created artifacts).  Honours the original "wire get_diff() into the run
    artifacts" intent for the worktree-isolation case.  Returns
    ``{"available": False}`` whenever git or a repo is absent — never raises.
    """
    # 改进⑧ (AUDIT_2026-06-11): 先用纯文件系统探测向上找 .git，目录树里
    # 没有就直接返回，省掉每次 ~50ms 的 git 子进程启动开销（多 stage
    # pipeline 下每 stage 一次）。找到了再交给 git rev-parse 权威确认
    # （.git 文件型 worktree、裸仓库等边角仍由 git 判定）。
    probe = os.path.realpath(path)
    while True:
        if os.path.exists(os.path.join(probe, ".git")):
            break
        parent = os.path.dirname(probe)
        if parent == probe:  # 到根了，没有仓库
            return {"available": False}
        probe = parent

    try:
        root = subprocess.run(
            ["git", "-C", path, "rev-parse", "--show-toplevel"],
            capture_output=True, text=True, timeout=10,
        )
        if root.returncode != 0 or not root.stdout.strip():
            return {"available": False}
        repo_root = root.stdout.strip()
        status = subprocess.run(
            ["git", "-C", path, "status", "--porcelain"],
            capture_output=True, text=True, timeout=10,
        )
        changed = [ln for ln in status.stdout.splitlines() if ln.strip()]
        diff = subprocess.run(
            ["git", "-C", path, "diff", "HEAD"],
            capture_output=True, text=True, timeout=15,
        )
        diff_text = diff.stdout
        truncated = len(diff_text) > _GIT_DIFF_MAX_CHARS
        if truncated:
            diff_text = diff_text[:_GIT_DIFF_MAX_CHARS]
        return {
            "available": True,
            "repo_root": repo_root,
            "changed_files": changed,
            "changed_count": len(changed),
            "diff": diff_text,
            "diff_truncated": truncated,
        }
    except Exception:
        return {"available": False}


def build_change_manifest(workspace_path: str) -> dict:
    """Scan a run workspace and write ``manifest.json`` — the per-run 修改清单.

    Records every file the agent produced under the run directory (excluding the
    framework's own metadata files in ``_MANIFEST_EXCLUDE``) with relative path,
    size, sha256 and mtime, plus a best-effort git diff of tracked-file edits.

    Because each run gets a fresh workspace, everything under it *is* this run's
    output — making the scan an authoritative evidence trail.  Returns the
    manifest dict (also written to ``{workspace_path}/manifest.json``); returns
    ``{}`` for an invalid path.  Never raises on per-file errors.
    """
    if not workspace_path or not os.path.isdir(workspace_path):
        return {}

    run_id = ""
    input_json = os.path.join(workspace_path, "input.json")
    if os.path.isfile(input_json):
        try:
            with open(input_json, encoding="utf-8") as fh:
                run_id = json.load(fh).get("run_id", "")
        except Exception:
            run_id = ""

    files: list[dict] = []
    total_bytes = 0
    for dirpath, _dirs, filenames in os.walk(workspace_path):
        for name in filenames:
            abspath = os.path.join(dirpath, name)
            rel = os.path.relpath(abspath, workspace_path)
            if rel in _MANIFEST_EXCLUDE:        # framework metadata, not a change
                continue
            try:
                st = os.stat(abspath)
            except OSError:
                continue
            total_bytes += st.st_size
            if st.st_size <= _MANIFEST_HASH_MAX_BYTES:
                try:
                    digest = _sha256_file(abspath)
                except OSError:
                    digest = None
            else:
                digest = None                    # too large to hash
            files.append({
                "path": rel.replace(os.sep, "/"),
                "size_bytes": st.st_size,
                "sha256": digest,
                "modified_at": time.strftime(
                    "%Y-%m-%dT%H:%M:%S", time.localtime(st.st_mtime)
                ),
            })

    files.sort(key=lambda e: e["path"])

    manifest = {
        "run_id": run_id,
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
        "file_count": len(files),
        "total_bytes": total_bytes,
        "files": files,
        "git": _git_diff_for(workspace_path),
    }

    with open(os.path.join(workspace_path, "manifest.json"), "w", encoding="utf-8") as fh:
        json.dump(manifest, fh, ensure_ascii=False, indent=2)

    return manifest


def save_run_artifacts(
    workspace_path: str,
    result: str,
    trace: list,
    duration_ms: int,
    input_tokens: int = 0,
    output_tokens: int = 0,
    steps: int = 0,
    error: str = "",
    status: str = "completed",
    usage: dict = None,
) -> None:
    """
    执行结束后保存 run 产物。

    保存:
        output.json  ← 最终输出 + 状态
        trace.json   ← 完整执行轨迹
        cost.json    ← 成本统计（含实际 token 数 + USD 估算）

    `usage` 字典（可选）：从 ClaudeCodeProvider.get_usage() 读取的精确 token 数，
    包含 input_tokens / output_tokens / cache_creation_input_tokens /
    cache_read_input_tokens。若提供则优先写入 cost.json；否则退回 trace 汇总。
    """
    if not workspace_path or not os.path.isdir(workspace_path):
        return

    # 保存输出
    with open(os.path.join(workspace_path, "output.json"), "w", encoding="utf-8") as f:
        json.dump({
            "status": status,
            "output": result,
            "error": error,
        }, f, ensure_ascii=False, indent=2)

    # 保存 trace
    trace_data = []
    for entry in trace:
        trace_data.append({
            "term_name": getattr(entry, "term_name", ""),
            "term_id": getattr(entry, "term_id", ""),
            "input": str(getattr(entry, "input", ""))[:500],
            "output": str(getattr(entry, "output", ""))[:500],
            "duration_ms": getattr(entry, "duration_ms", 0),
            "model": getattr(entry, "model", ""),
            "tokens_used": getattr(entry, "tokens_used", 0),
        })
    with open(os.path.join(workspace_path, "trace.json"), "w", encoding="utf-8") as f:
        json.dump(trace_data, f, ensure_ascii=False, indent=2)

    # 保存成本统计
    if usage:
        # Accurate counts from ClaudeCode JSONL (preferred)
        real_input   = usage.get("input_tokens", 0)
        real_output  = usage.get("output_tokens", 0)
        cache_read   = usage.get("cache_read_input_tokens", 0)
        cache_create = usage.get("cache_creation_input_tokens", 0)
        cost_usd     = _compute_cost_usd(usage)
        _prov        = usage.get("provider", "claude")
        # Accurate source label: Claude reads JSONL transcripts; OpenAI-compatible
        # providers (ollama/qwen/...) report usage from the API response.
        source       = "claude_jsonl" if _prov in ("claude", "claude-code") else f"{_prov}_api"
    else:
        # Fallback: sum ctx.trace tokens_used (usually 0 for claude-code sessions)
        real_input   = input_tokens
        real_output  = output_tokens
        cache_read   = 0
        cache_create = 0
        cost_usd     = _compute_cost_usd({
            "input_tokens": real_input,
            "output_tokens": real_output,
        })
        source       = "trace_fallback"

    with open(os.path.join(workspace_path, "cost.json"), "w", encoding="utf-8") as f:
        json.dump({
            "status": status,
            "duration_ms": duration_ms,
            "steps": steps,
            "source": source,
            "input_tokens": real_input,
            "output_tokens": real_output,
            "cache_read_tokens": cache_read,
            "cache_creation_tokens": cache_create,
            "total_tokens": real_input + real_output + cache_read + cache_create,
            "cost_usd": round(cost_usd, 6),
        }, f, ensure_ascii=False, indent=2)

    # 保存修改清单（change manifest）— never let manifest failure break a run
    try:
        build_change_manifest(workspace_path)
    except Exception:
        pass


# ============================================================
# Sandbox 执行器
# ============================================================

class Sandbox:
    def __init__(self, level: int = 0, limits: ResourceLimits = None):
        self.level = level
        self.limits = limits or ResourceLimits()

    async def execute(
        self,
        config: dict,
        input_text: str,
        timeout: int = 120,
        agent_dir: str = "",
        run_id: str = "",
    ) -> ExecutionResult:
        if self.level == 0:
            return await self._exec_inprocess(config, input_text, timeout, agent_dir, run_id)
        elif self.level == 1:
            return await self._exec_subprocess(config, input_text, timeout, agent_dir, run_id)
        else:
            return await self._exec_inprocess(config, input_text, timeout, agent_dir, run_id)

    async def _exec_inprocess(
        self, config: dict, input_text: str, timeout: int,
        agent_dir: str = "", run_id: str = "",
    ) -> ExecutionResult:
        t0 = time.time()

        # 创建 run workspace（如果有 agent_dir）
        workspace_path = ""
        if agent_dir:
            workspace_path = create_run_workspace(agent_dir, run_id, input_text, config)

        try:
            import yaml
            # 配置写入 workspace（如果有）或 /tmp
            if workspace_path:
                config_path = os.path.join(workspace_path, "config.yml")
            else:
                import tempfile
                with tempfile.NamedTemporaryFile(mode='w', suffix='.yml', delete=False, encoding='utf-8') as f:
                    yaml.dump(config, f, allow_unicode=True)
                    config_path = f.name

            try:
                from lambdagent.fromconfig import from_config
                from lambdagent.core import Context
                term = from_config(config_path, workspace_path=workspace_path)
                ctx = Context(workspace_path=workspace_path, run_id=run_id)
                # audit #10: offload sync term.apply off the event loop;
                # see matching fix at api/v1/agents.py.
                import asyncio as _asyncio
                _loop = _asyncio.get_event_loop()
                result = await _loop.run_in_executor(None, term.apply, input_text, ctx)
                duration = int((time.time() - t0) * 1000)

                # ── Accurate token counts from ClaudeCode JSONL transcript ──
                # _compile_react attaches _think_ref (ConversationLam) to the
                # Loop term so we can call get_usage() after the run finishes.
                # This reads ~/.claude/projects/*/{session_id}.jsonl and sums
                # all assistant-turn usage entries for the real numbers.
                think_ref = getattr(term, "_think_ref", None)
                usage = None
                if think_ref is not None and callable(getattr(think_ref, "get_usage", None)):
                    try:
                        usage = think_ref.get_usage()
                    except Exception:
                        usage = None

                trace_tokens = sum(getattr(e, 'tokens_used', 0) for e in ctx.trace)
                real_input = (usage or {}).get("input_tokens", trace_tokens)

                exec_result = ExecutionResult(
                    output=str(result), status="completed", duration_ms=duration,
                    steps=len(ctx.trace),
                    input_tokens=real_input,
                    workspace_path=workspace_path,
                )

                # 保存 run 产物（含精确 token 数 + USD 成本估算）
                save_run_artifacts(
                    workspace_path, str(result), ctx.trace,
                    duration, real_input, 0, len(ctx.trace),
                    usage=usage,
                )

                return exec_result
            finally:
                # 只有 /tmp 的临时文件需要删除，workspace 里的保留
                if not workspace_path:
                    try:
                        os.unlink(config_path)
                    except OSError:
                        pass

        except Exception as e:
            duration = int((time.time() - t0) * 1000)
            error_msg = str(e)

            # 即使失败也保存产物
            save_run_artifacts(
                workspace_path, "", [], duration,
                error=error_msg, status="failed",
            )

            return ExecutionResult(
                status="failed", error=error_msg,
                duration_ms=duration,
                workspace_path=workspace_path,
            )

    async def _exec_subprocess(
        self, config: dict, input_text: str, timeout: int,
        agent_dir: str = "", run_id: str = "",
    ) -> ExecutionResult:
        import yaml

        # 创建 run workspace（如果有 agent_dir）
        workspace_path = ""
        if agent_dir:
            workspace_path = create_run_workspace(agent_dir, run_id, input_text, config)

        # 配置写入 workspace 或 /tmp
        if workspace_path:
            config_path = os.path.join(workspace_path, "config.yml")
        else:
            import tempfile
            with tempfile.NamedTemporaryFile(mode='w', suffix='.yml', delete=False, encoding='utf-8') as f:
                yaml.dump(config, f, allow_unicode=True)
                config_path = f.name

        # SEC-09: Pass input via JSON file, NOT string interpolation
        if workspace_path:
            input_path = os.path.join(workspace_path, "input.json")
            # input.json already written by create_run_workspace, but subprocess expects {"input": ...}
        else:
            import tempfile
            with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, encoding='utf-8') as inf:
                json.dump({"input": input_text}, inf, ensure_ascii=False)
                input_path = inf.name

        # Subprocess runner — pass workspace_path as argv[3]
        runner_code = (
            "import sys, json, time, os\n"
            "t0 = time.time()\n"
            "try:\n"
            "    config_path = sys.argv[1]\n"
            "    with open(sys.argv[2]) as _if: input_text = json.load(_if)['input']\n"
            "    workspace_path = sys.argv[3] if len(sys.argv) > 3 else ''\n"
            "    from lambdagent.fromconfig import from_config\n"
            "    from lambdagent.core import Context\n"
            "    term = from_config(config_path, workspace_path=workspace_path)\n"
            "    ctx = Context(workspace_path=workspace_path)\n"
            "    result = term.apply(input_text, ctx)\n"
            "    duration = int((time.time() - t0) * 1000)\n"
            "    # Save trace to workspace\n"
            "    if workspace_path:\n"
            "        trace_data = [{'term_name': e.term_name, 'duration_ms': e.duration_ms, "
            "'tokens_used': e.tokens_used} for e in ctx.trace]\n"
            "        with open(os.path.join(workspace_path, 'trace.json'), 'w') as tf:\n"
            "            json.dump(trace_data, tf, indent=2)\n"
            "        with open(os.path.join(workspace_path, 'output.json'), 'w') as of:\n"
            "            json.dump({'status': 'completed', 'output': str(result)}, of, indent=2)\n"
            "        with open(os.path.join(workspace_path, 'cost.json'), 'w') as cf:\n"
            "            json.dump({'duration_ms': duration, 'steps': len(ctx.trace), "
            "'input_tokens': sum(getattr(e,'tokens_used',0) for e in ctx.trace)}, cf, indent=2)\n"
            "    print(json.dumps({'output': str(result), 'status': 'completed', "
            "'duration_ms': duration, 'steps': len(ctx.trace), "
            "'workspace_path': workspace_path}))\n"
            "except Exception as e:\n"
            "    duration = int((time.time() - t0) * 1000)\n"
            "    workspace_path = sys.argv[3] if len(sys.argv) > 3 else ''\n"
            "    if workspace_path:\n"
            "        with open(os.path.join(workspace_path, 'output.json'), 'w') as of:\n"
            "            json.dump({'status': 'failed', 'error': str(e)}, of)\n"
            "    print(json.dumps({'status': 'failed', 'error': str(e), "
            "'duration_ms': duration, 'workspace_path': workspace_path}))\n"
        )

        import tempfile
        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False, encoding='utf-8') as sf:
            sf.write(runner_code)
            script_path = sf.name

        try:
            proc = await asyncio.create_subprocess_exec(
                'python', script_path, config_path, input_path, workspace_path,
                stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE,
            )
            stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
            output = stdout.decode().strip()
            if output:
                data = json.loads(output)
                return ExecutionResult(**{k: v for k, v in data.items() if hasattr(ExecutionResult, k)})
            return ExecutionResult(status="failed", error=stderr.decode()[:500],
                                   workspace_path=workspace_path)
        except asyncio.TimeoutError:
            proc.kill()
            return ExecutionResult(status="failed",
                                   error="Execution timeout ({}s)".format(timeout),
                                   workspace_path=workspace_path)
        finally:
            # 只清理 runner script 和非 workspace 的临时文件
            try:
                os.unlink(script_path)
            except OSError:
                pass
            # 保存修改清单（subprocess 路径不走 save_run_artifacts）
            if workspace_path:
                try:
                    build_change_manifest(workspace_path)
                except Exception:
                    pass
            if not workspace_path:
                for p in (config_path, input_path):
                    try:
                        os.unlink(p)
                    except OSError:
                        pass
