#!/usr/bin/env python3
"""
Codex app-server 客户端（stdio JSON-RPC，零第三方依赖）。

照搬的是官方 codex 二进制自带的 app-server：本文件不实现 agent 循环、
不实现沙箱、不实现工具，只做一件事——按官方协议跟 app-server 对话：
  initialize -> initialized -> thread/start -> turn/start
  -> 流式事件（agent 文本、工具调用、审批请求）-> turn/completed

协议定义见 protocol/codex_app_server_protocol.v2.schemas.json，
由 `codex app-server generate-json-schema` 从官方二进制原样导出。

用法：
  python3 client.py --prompt "修掉 README 里的错别字"
  python3 client.py --prompt "..." --cwd /path/to/repo --model gpt-5.x
  python3 client.py --list-threads
  python3 client.py --thread-id <id> --prompt "继续"
  python3 client.py --prompt "..." --approval-policy on-request --approval deny
  python3 client.py --prompt "..." --sandbox read-only

审批：默认 --approval-policy never，服务端不发审批请求，适合自动化。
若显式用 on-request，--approval auto 会放行、deny 会拒绝（默认 deny）。
注意：子 agent 的审批/沙箱继承父 thread，这里不做额外权限边界，
见 README“子代理”一节。
"""
import argparse
import json
import os
import select
import subprocess
import sys
import time

CODEX = os.environ.get("CODEX_BIN", "codex")


class AppServer:
    def __init__(self, codex_home=None, extra_args=None):
        env = dict(os.environ)
        if codex_home:
            os.makedirs(codex_home, exist_ok=True)
            env["CODEX_HOME"] = codex_home
        cmd = [CODEX, "app-server"]
        if extra_args:
            cmd += extra_args
        self.proc = subprocess.Popen(
            cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL, env=env,
        )
        self.buf = b""
        self.next_id = 1
        self.pending = {}  # id -> message
        self.approval_policy = "deny"

    def send(self, obj):
        line = json.dumps(obj) + "\n"
        self.proc.stdin.write(line.encode())
        self.proc.stdin.flush()
        return obj.get("id")

    def request(self, method, params=None):
        rid = self.next_id
        self.next_id += 1
        self.send({"jsonrpc": "2.0", "id": rid, "method": method,
                   "params": params or {}})
        return rid

    def notify(self, method, params=None):
        self.send({"jsonrpc": "2.0", "method": method, "params": params or {}})

    def _read_msg(self, timeout):
        """读一条服务端消息；是服务端请求时先回包，再继续读。"""
        end = time.time() + timeout
        while time.time() < end:
            while b"\n" in self.buf:
                line, self.buf = self.buf.split(b"\n", 1)
                if not line.strip():
                    continue
                msg = json.loads(line)
                if "method" in msg and "id" in msg:
                    self._answer_server_request(msg)
                    continue
                return msg
            r, _, _ = select.select([self.proc.stdout], [], [], 0.2)
            if r:
                chunk = os.read(self.proc.stdout.fileno(), 65536)
                if not chunk:
                    return None
                self.buf += chunk
            if self.proc.poll() is not None:
                return None
        return None

    def _answer_server_request(self, msg):
        method = msg.get("method", "")
        # 审批类请求：按策略回包。不同审批方法的回包字段在各版本间有差异，
        # 这里用最常见的 decision 字段；若服务端不认，会在事件里看到错误，
        # 此时应改回 approval-policy=never 跑自动化，或在 README 的审批表里查本版字段。
        decision = "approve" if self.approval_policy == "auto" else "deny"
        payload = {"decision": decision}
        if "approval" in method.lower() or "permission" in method.lower():
            print(f"[server-request] {method} -> {decision} "
                  f"(params={json.dumps(msg.get('params', {}))[:400]})",
                  file=sys.stderr)
        self.send({"jsonrpc": "2.0", "id": msg["id"], "result": payload})

    def wait_for(self, rid, timeout=600, on_event=None):
        """等指定 id 的响应；途中的通知/事件交给 on_event 打印。"""
        end = time.time() + timeout
        while time.time() < end:
            msg = self._read_msg(max(1, end - time.time()))
            if msg is None:
                break
            if msg.get("id") == rid:
                return msg
            if on_event:
                on_event(msg)
        return None

    def initialize(self):
        rid = self.request("initialize", {
            "clientInfo": {"name": "codex-app-server-client", "version": "0.1"},
        })
        resp = self.wait_for(rid, timeout=60)
        if not resp or "error" in resp:
            raise RuntimeError(f"initialize 失败: {resp}")
        self.notify("initialized", {})
        return resp["result"]

    def close(self):
        try:
            self.proc.stdin.close()
        except Exception:
            pass
        try:
            self.proc.wait(timeout=5)
        except Exception:
            self.proc.kill()


def handle_event(msg):
    method = msg.get("method", "")
    params = msg.get("params", {}) or {}
    # 文本增量：不同版本的字段名略有差异，尽量都覆盖
    if method in ("thread/item/updated", "item/updated", "agent/message/delta",
                   "thread/item/started", "thread/item/completed"):
        item = params.get("item") or params
        text = item.get("text") or item.get("delta") or ""
        if not text and isinstance(item.get("content"), list):
            text = "".join(c.get("text", "") for c in item["content"]
                           if isinstance(c, dict))
        if text and item.get("type") in ("agentMessage", None, "agent_message"):
            print(text, end="", flush=True)
    elif method == "turn/completed":
        print(f"\n[turn/completed] {json.dumps(params)[:800]}")
    elif "completed" in method or "started" in method:
        print(f"\n[{method}] {json.dumps(params)[:500]}", file=sys.stderr)
    elif "error" in str(params).lower():
        print(f"\n[{method}] {json.dumps(params)[:800]}", file=sys.stderr)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompt", help="本回合发给 agent 的文本")
    ap.add_argument("--cwd", default=os.getcwd())
    ap.add_argument("--model", default=None)
    ap.add_argument("--sandbox", default="workspace-write",
                    choices=["read-only", "workspace-write", "danger-full-access"])
    ap.add_argument("--approval-policy", default="never",
                    choices=["never", "on-request", "untrusted"])
    ap.add_argument("--approval", default="deny", choices=["auto", "deny"])
    ap.add_argument("--effort", default=None,
                    help="本回合推理档位，如 low/medium/high/xhigh/max/ultra")
    ap.add_argument("--thread-id", default=None, help="续跑已有 thread，不传则新建")
    ap.add_argument("--list-threads", action="store_true")
    ap.add_argument("--codex-home", default=None,
                    help="默认用 ~/.codex；测试时可指向空目录隔离")
    ap.add_argument("--timeout", type=int, default=600)
    args = ap.parse_args()

    srv = AppServer(codex_home=args.codex_home)
    srv.approval_policy = args.approval
    try:
        info = srv.initialize()
        print(f"[initialize] codexHome={info.get('codexHome')} "
              f"platform={info.get('platformOs')}", file=sys.stderr)

        if args.list_threads:
            rid = srv.request("thread/list", {})
            print(json.dumps(srv.wait_for(rid, timeout=60), indent=2,
                             ensure_ascii=False))
            return

        if args.thread_id:
            rid = srv.request("thread/resume", {"threadId": args.thread_id})
            resp = srv.wait_for(rid, timeout=60)
            if not resp or "error" in resp:
                print(f"thread/resume 失败: {resp}", file=sys.stderr)
                return
            thread_id = args.thread_id
            print(f"[thread/resumed] {thread_id}", file=sys.stderr)
        else:
            params = {"cwd": os.path.abspath(args.cwd),
                      "approvalPolicy": args.approval_policy,
                      "sandbox": args.sandbox}
            if args.model:
                params["model"] = args.model
            rid = srv.request("thread/start", params)
            resp = srv.wait_for(rid, timeout=60)
            if not resp or "error" in resp:
                print(f"thread/start 失败: {resp}", file=sys.stderr)
                return
            thread = resp["result"]["thread"]
            thread_id = thread["id"]
            print(f"[thread/started] {thread_id} "
                  f"model={resp['result'].get('model')}", file=sys.stderr)

        if not args.prompt:
            print(f"thread_id={thread_id}")
            return
        turn_params = {
            "threadId": thread_id,
            "input": [{"type": "text", "text": args.prompt}],
        }
        if args.model:
            turn_params["model"] = args.model
        if args.effort:
            turn_params["effort"] = args.effort
        rid = srv.request("turn/start", turn_params)
        # turn/start 的响应本身不代表完成；完成以 turn/completed 通知为准
        print("[turn/started] 等待 turn/completed …", file=sys.stderr)
        end = time.time() + args.timeout
        while time.time() < end:
            msg = srv._read_msg(max(1, end - time.time()))
            if msg is None:
                break
            if msg.get("id") == rid and "error" in msg:
                print(f"turn/start 出错: {msg}", file=sys.stderr)
                break
            if msg.get("method") == "turn/completed":
                print(f"\n[done] {json.dumps(msg.get('params', {}))[:1000]}")
                break
            handle_event(msg)
    finally:
        srv.close()


if __name__ == "__main__":
    main()