#!/usr/bin/env python3
"""
轮询指定 agent 的 sessions.json，对「超阈值且非 running」的会话执行 /compact。
通过 openclaw gateway call chat.send 向目标 session 发送 /compact 控制命令。

Heartbeat 须在独立会话中运行（heartbeat.isolatedSession=true），否则执行脚本的
会话在 store 里为 running，主会话无法被选中。脚本会跳过键名以 :heartbeat 结尾的
会话，以及环境变量标明的当前执行会话。

局部 compact 失败仅写入 JSON 汇总，进程仍以 exit 0 结束，便于 Heartbeat 稳定回复 HEARTBEAT_OK。

默认状态目录: /home/node/.openclaw
--agent-id 必填（由调用方传入当前 openclaw Agent ID）
可选环境变量 OPENCLAW_STATE_DIR 覆盖状态目录

用法:
  python3 compact_sessions_over_threshold.py --agent-id parent-toddler --json
  python3 compact_sessions_over_threshold.py --agent-id parent-toddler --dry-run
  python3 compact_sessions_over_threshold.py --agent-id parent-toddler --threshold 100000
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import uuid
from pathlib import Path
from typing import Any, Iterable, Optional

DEFAULT_STATE_DIR = Path("/home/node/.openclaw")
DEFAULT_TOKEN_THRESHOLD = 80_000
COMPACT_MESSAGE = "/compact"
DEFAULT_AGENT_TIMEOUT_SEC = 600


def _as_int_or_none(v: Any) -> Optional[int]:
    if v is None or isinstance(v, bool):
        return None
    if isinstance(v, int):
        return v
    if isinstance(v, float) and v.is_integer():
        return int(v)
    return None


def iter_session_rows(store: dict[str, Any]) -> Iterable[dict[str, Any]]:
    for key, entry in store.items():
        if not isinstance(entry, dict):
            continue
        st = entry.get("status")
        yield {
            "key": key,
            "sessionId": entry.get("sessionId"),
            "totalTokens": _as_int_or_none(entry.get("totalTokens")),
            "statusRunning": st == "running",
        }


def list_sessions_from_store(path: Path) -> list[dict[str, Any]]:
    raw = path.read_text(encoding="utf-8")
    data = json.loads(raw)
    if not isinstance(data, dict):
        raise ValueError("根节点必须是 JSON 对象（键为 session key）")
    return list(iter_session_rows(data))


def filter_sessions_by_agent(rows: list[dict[str, Any]], agent_id: str) -> list[dict[str, Any]]:
    aid = agent_id.strip()
    if not aid:
        raise ValueError("agent_id 不能为空")
    prefix = f"agent:{aid}:"
    return [row for row in rows if str(row.get("key", "")).startswith(prefix)]


def filter_sessions_over_threshold(
    rows: list[dict[str, Any]],
    threshold: int,
) -> list[dict[str, Any]]:
    if threshold < 0:
        raise ValueError("threshold 不能为负数")
    out: list[dict[str, Any]] = []
    for row in rows:
        total = row.get("totalTokens")
        if total is None:
            continue
        if total > threshold:
            out.append(row)
    return out


def filter_sessions_idle(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    return [row for row in rows if row.get("statusRunning") is not True]


def _resolve_executor_session_key() -> Optional[str]:
    for name in ("OPENCLAW_TASK_SESSION_KEY", "OPENCLAW_SESSION_KEY"):
        value = os.environ.get(name, "").strip()
        if value:
            return value
    return None


def _resolve_executor_session_id() -> Optional[str]:
    value = os.environ.get("OPENCLAW_TASK_SESSION_ID", "").strip()
    return value or None


def filter_sessions_excluded_from_compact(
    rows: list[dict[str, Any]],
    *,
    executor_session_key: Optional[str] = None,
    executor_session_id: Optional[str] = None,
) -> list[dict[str, Any]]:
    """跳过 Heartbeat 隔离会话及当前正在执行本脚本的会话。"""
    key = (executor_session_key or "").strip()
    sid = (executor_session_id or "").strip()
    out: list[dict[str, Any]] = []
    for row in rows:
        row_key = str(row.get("key", ""))
        if row_key.endswith(":heartbeat"):
            continue
        if key and row_key == key:
            continue
        row_sid = row.get("sessionId")
        if sid and row_sid is not None and str(row_sid) == sid:
            continue
        out.append(row)
    return out


def filter_sessions_compactable(
    rows: list[dict[str, Any]],
    *,
    executor_session_key: Optional[str] = None,
    executor_session_id: Optional[str] = None,
) -> list[dict[str, Any]]:
    rows = filter_sessions_excluded_from_compact(
        rows,
        executor_session_key=executor_session_key,
        executor_session_id=executor_session_id,
    )
    return filter_sessions_idle(rows)


def resolve_state_dir() -> Path:
    override = os.environ.get("OPENCLAW_STATE_DIR", "").strip()
    if override:
        return Path(override).expanduser()
    return DEFAULT_STATE_DIR


def resolve_sessions_store_path(agent_id: str, explicit: Optional[Path]) -> Path:
    if explicit is not None:
        return explicit.expanduser().resolve()
    return resolve_state_dir() / "agents" / agent_id / "sessions" / "sessions.json"


def find_compact_targets(
    store_path: Path,
    agent_id: str,
    threshold: int,
    *,
    executor_session_key: Optional[str] = None,
    executor_session_id: Optional[str] = None,
) -> list[dict[str, Any]]:
    rows = list_sessions_from_store(store_path)
    rows = filter_sessions_by_agent(rows, agent_id)
    rows = filter_sessions_over_threshold(rows, threshold)
    return filter_sessions_compactable(
        rows,
        executor_session_key=executor_session_key or _resolve_executor_session_key(),
        executor_session_id=executor_session_id or _resolve_executor_session_id(),
    )


def run_compact_for_session(
    session_key: str,
    *,
    timeout_sec: int,
    dry_run: bool,
) -> tuple[bool, str]:
    key = session_key.strip()
    if not key:
        return False, "missing sessionKey"
    if dry_run:
        return True, "dry-run"

    params = {
        "sessionKey": key,
        "message": COMPACT_MESSAGE,
        "idempotencyKey": str(uuid.uuid4()),
        "timeoutMs": timeout_sec * 1000,
    }
    # CLI --timeout 为毫秒，略大于 params.timeoutMs 以等待 gateway 收尾。
    cli_timeout_ms = (timeout_sec + 15) * 1000
    cmd = [
        "openclaw",
        "gateway",
        "call",
        "chat.send",
        "--params",
        json.dumps(params, ensure_ascii=False),
        "--expect-final",
        "--timeout",
        str(cli_timeout_ms),
    ]
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=(timeout_sec + 30),
            check=False,
        )
    except FileNotFoundError:
        return False, "openclaw 命令未找到"
    except subprocess.TimeoutExpired:
        return False, f"超时（>{timeout_sec}s）"

    if result.returncode != 0:
        detail = (result.stderr or result.stdout or "").strip()
        return False, detail or f"exit {result.returncode}"
    return True, "ok"


def poll_and_compact(
    *,
    store_path: Path,
    agent_id: str,
    threshold: int,
    dry_run: bool,
    timeout_sec: int,
) -> dict[str, Any]:
    targets = find_compact_targets(store_path, agent_id, threshold)
    results: list[dict[str, Any]] = []
    ok_count = 0
    fail_count = 0

    for row in targets:
        key = row.get("key")
        session_id = row.get("sessionId")
        total = row.get("totalTokens")

        success, detail = run_compact_for_session(
            str(key) if key else "",
            timeout_sec=timeout_sec,
            dry_run=dry_run,
        )
        if success:
            ok_count += 1
            status = "dry-run" if dry_run else "compacted"
        else:
            fail_count += 1
            status = "failed"
        results.append(
            {
                "key": key,
                "sessionId": session_id,
                "totalTokens": total,
                "statusRunning": False,
                "status": status,
                "detail": detail,
            }
        )

    return {
        "agentId": agent_id,
        "storePath": str(store_path),
        "threshold": threshold,
        "dryRun": dry_run,
        "targetCount": len(targets),
        "okCount": ok_count,
        "failCount": fail_count,
        "results": results,
    }


def main() -> None:
    p = argparse.ArgumentParser(
        description=(
            "对 totalTokens 超过阈值、非 running、且非 Heartbeat/执行器会话的条目执行 /compact"
        )
    )
    p.add_argument(
        "--store",
        type=Path,
        default=None,
        help="sessions.json 路径（默认 /home/node/.openclaw/agents/<agent>/sessions/sessions.json）",
    )
    p.add_argument(
        "--agent-id",
        type=str,
        required=True,
        help="当前 openclaw Agent ID（必填，由 Heartbeat 等调用方传入）",
    )
    p.add_argument(
        "--threshold",
        type=int,
        default=DEFAULT_TOKEN_THRESHOLD,
        help=f"totalTokens 阈值，严格大于该值才处理（默认 {DEFAULT_TOKEN_THRESHOLD}）",
    )
    p.add_argument(
        "--dry-run",
        action="store_true",
        help="只列出将压缩的会话（已排除 running），不调用 gateway chat.send",
    )
    p.add_argument(
        "--timeout",
        type=int,
        default=DEFAULT_AGENT_TIMEOUT_SEC,
        help=f"单次 chat.send /compact 超时秒数（默认 {DEFAULT_AGENT_TIMEOUT_SEC}）",
    )
    p.add_argument(
        "--json",
        action="store_true",
        help="以 JSON 输出汇总（便于日志/监控）",
    )
    args = p.parse_args()

    agent_id = args.agent_id.strip()
    if not agent_id:
        print("错误：--agent-id 不能为空", file=sys.stderr)
        sys.exit(1)
    store_path = resolve_sessions_store_path(agent_id, args.store)

    if not store_path.is_file():
        print(f"错误：找不到 session store: {store_path}", file=sys.stderr)
        sys.exit(1)

    try:
        summary = poll_and_compact(
            store_path=store_path,
            agent_id=agent_id,
            threshold=args.threshold,
            dry_run=args.dry_run,
            timeout_sec=max(30, args.timeout),
        )
    except json.JSONDecodeError as e:
        print(f"错误：JSON 解析失败: {e}", file=sys.stderr)
        sys.exit(1)
    except ValueError as e:
        print(f"错误：{e}", file=sys.stderr)
        sys.exit(1)

    if args.json:
        print(json.dumps(summary, ensure_ascii=False, indent=2))
    else:
        print(
            f"agent={summary['agentId']} store={summary['storePath']} "
            f"threshold>{summary['threshold']} targets={summary['targetCount']} "
            f"ok={summary['okCount']} failed={summary['failCount']}"
        )
        for item in summary["results"]:
            print(
                f"  [{item['status']}] {item.get('key')} "
                f"tokens={item.get('totalTokens')} "
                f"sessionId={item.get('sessionId')} "
                f"{item.get('detail', '')}"
            )

    sys.exit(0)


if __name__ == "__main__":
    main()
