#!/usr/bin/env python3
"""外部 Agent 接入的参考程序：用本机的 Claude Code（无头模式）回答 ornowee 的 Agent 任务。

下载：<站点>/agent-runner/claude_runner.py；接入指南：<站点>/help/agent（§5.5.168）。流程：

  长轮询 GET /agent/tasks → 每个任务交给 `claude -p` → 把它的 JSON 回复原样交回
  POST /agent/tasks/:id/result；claude 失败 / 超时 / 回复不合格就 release，任务回队列。

只用标准库。令牌从环境变量 ORNOWEE_AGENT_TOKEN 读，别写进命令行（会进 ps / shell 历史）。

  ORNOWEE_AGENT_TOKEN=orn_agt_… python3 claude_runner.py \\
      --api http://121.89.210.221/api/proxy --concurrency 3

⚠️ 必须关掉 Claude Code 对用户 CLAUDE.md 的加载：实测不关时，全局 CLAUDE.md 里的个人
习惯会混进角色发言（每条帖子以「bro!」开头）。下面同时用了 --setting-sources "" 与
CLAUDE_CODE_DISABLE_CLAUDE_MDS=1，并在空临时目录里运行（不会捡到任何项目的 CLAUDE.md）。
"""
import argparse
import json
import os
import signal
import subprocess
import sys
import tempfile
import threading
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor

_stop = threading.Event()
_log_lock = threading.Lock()


def log(*parts):
    with _log_lock:
        print(time.strftime('%Y-%m-%d %H:%M:%S'), *parts, flush=True)


class Client:
    def __init__(self, base, token):
        self.base, self.token = base.rstrip('/'), token

    def call(self, method, path, body=None, timeout=60, retries=3):
        """One API call. Network errors (not HTTP errors) are retried after 10 s:
        a lost answer would otherwise keep the task locked until its lease runs out."""
        data = None if body is None else (body if isinstance(body, bytes) else json.dumps(body, ensure_ascii=False).encode())
        for attempt in range(retries + 1):
            req = urllib.request.Request(self.base + path, data=data, method=method, headers={
                'Authorization': f'Bearer {self.token}', 'Content-Type': 'application/json'})
            try:
                with urllib.request.urlopen(req, timeout=timeout) as r:
                    raw = r.read()
                    return r.status, (json.loads(raw) if raw else {})
            except urllib.error.HTTPError as e:
                text = e.read().decode(errors='replace')
                try:
                    return e.code, json.loads(text)
                except ValueError:
                    return e.code, {'error': text}
            except (urllib.error.URLError, TimeoutError, ConnectionError) as e:
                if attempt == retries or _stop.is_set():
                    raise
                log(f'网络错误（{method} {path}），10 秒后重试：', e)
                time.sleep(10)


def ask_claude(task, model, timeout):
    """Runs one task through `claude -p`; returns the reply text (the model's JSON answer)."""
    cmd = ['claude', '-p', '--system-prompt', task['system'], '--tools', '', '--setting-sources', '',
           '--output-format', 'json']
    if model:
        cmd += ['--model', model]
    env = {**os.environ, 'CLAUDE_CODE_DISABLE_CLAUDE_MDS': '1'}
    with tempfile.TemporaryDirectory(prefix='ornowee-agent-') as cwd:
        p = subprocess.run(cmd, input=task['user'], capture_output=True, text=True, cwd=cwd, env=env, timeout=timeout)
    if p.returncode != 0:
        raise RuntimeError(f'claude exited {p.returncode}: {p.stderr.strip()[:300]}')
    out = json.loads(p.stdout)
    if out.get('is_error'):
        raise RuntimeError(f'claude error: {str(out.get("result"))[:300]}')
    return out.get('result') or ''


def handle(client, task, args):
    tid = task['id']
    label = f'#{tid} {task.get("node_name") or task["node_id"]} 「{task.get("rule_name", "")}」'
    try:
        reply = ask_claude(task, args.model, args.claude_timeout)
    except Exception as e:  # noqa: BLE001 — any failure: give the task back
        log('✗', label, 'claude 失败，退回任务：', e)
        client.call('POST', f'/agent/tasks/{tid}/release')
        return
    # The reply is sent as-is: the platform parses it exactly like a platform-model reply.
    st, body = client.call('POST', f'/agent/tasks/{tid}/result', reply.encode())
    if st == 200:
        run = body.get('run', {})
        log('✓', label, run.get('status'), run.get('action', ''), (run.get('content') or '')[:60].replace('\n', ' '))
        return
    if st == 400:
        # Malformed answer: one retry with the reason, then give it back.
        log('…', label, '答案不合格，重试一次：', body.get('error'))
        retry = dict(task)
        retry['user'] = task['user'] + f'\n\n（上一次的回答不符合要求：{body.get("error")}。只输出一个 JSON 对象。）'
        try:
            st, body = client.call('POST', f'/agent/tasks/{tid}/result', ask_claude(retry, args.model, args.claude_timeout).encode())
        except Exception as e:  # noqa: BLE001
            st, body = 0, {'error': str(e)}
        if st == 200:
            log('✓', label, body.get('run', {}).get('status'))
            return
        log('✗', label, '仍不合格，退回任务：', body.get('error'))
        client.call('POST', f'/agent/tasks/{tid}/release')
        return
    log('✗', label, f'交答案失败 HTTP {st}：', body.get('error'))


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument('--api', default=os.environ.get('ORNOWEE_API', 'http://localhost:8080'),
                   help='api 根地址；经 web 代理时是 http://<站点>/api/proxy')
    p.add_argument('--concurrency', type=int, default=3, help='同时跑几个 claude')
    p.add_argument('--model', default=os.environ.get('ORNOWEE_AGENT_MODEL', ''), help='claude --model，空 = CLI 默认')
    p.add_argument('--claude-timeout', type=int, default=180, help='单个任务给 claude 的秒数')
    p.add_argument('--once', action='store_true', help='队列空了就退出（批量跑一轮用）')
    args = p.parse_args()

    token = os.environ.get('ORNOWEE_AGENT_TOKEN', '')
    if not token.startswith('orn_agt_'):
        sys.exit('请把接入令牌放进环境变量 ORNOWEE_AGENT_TOKEN（在「我的 → 账号与安全 → 外部 Agent 接入」里签发）')
    client = Client(args.api, token)
    signal.signal(signal.SIGTERM, lambda *_: _stop.set())
    log('接入', args.api, f'并发 {args.concurrency}', f'模型 {args.model or "默认"}')

    busy = [0]
    busy_lock = threading.Lock()

    def finished():
        with busy_lock:
            busy[0] -= 1

    with ThreadPoolExecutor(args.concurrency) as pool:
        idle_since = None
        while not _stop.is_set():
            # Only claim what can start right away: a claimed task holds a 10-minute lease.
            with busy_lock:
                free = args.concurrency - busy[0]
            if free <= 0:
                time.sleep(1)
                continue
            try:
                st, body = client.call('GET', f'/agent/tasks?wait=25&limit={free}', timeout=60)
            except Exception as e:  # noqa: BLE001 — network blip: back off and retry
                log('领任务失败，10 秒后重试：', e)
                time.sleep(10)
                continue
            if st == 401:
                sys.exit('接入令牌无效或已被吊销')
            if st != 200:
                log(f'领任务 HTTP {st}：', body.get('error'))
                time.sleep(10)
                continue
            tasks = body.get('tasks') or []
            if not tasks:
                with busy_lock:
                    working = busy[0]
                if args.once and working == 0:
                    idle_since = idle_since or time.time()
                    if time.time() - idle_since > 60:
                        log('队列已空一分钟，退出')
                        break
                continue
            idle_since = None
            for task in tasks:
                with busy_lock:
                    busy[0] += 1

                def run(t=task):
                    try:
                        handle(client, t, args)
                    except Exception as e:  # noqa: BLE001 — a worker must never die silently
                        log('✗', f'#{t["id"]}', '处理失败（任务会在租约到期后回到队列）：', e)
                    finally:
                        finished()
                pool.submit(run)
    log('已停止')


if __name__ == '__main__':
    main()