#!/usr/bin/env python3
"""
Reapply OpenClaw subagent transcript hotfix after package updates.

What it patches:
1) deliverAgentCommandResult(): mirror subagent/CLI output into child transcript
2) runtimeMs math clamp: prevent negative runtime displays

Usage:
  sudo python3 scripts/reapply_subagent_transcript_fix.py --apply --restart
  sudo python3 scripts/reapply_subagent_transcript_fix.py --dry-run
"""

from __future__ import annotations

import argparse
import glob
import pathlib
import shutil
import subprocess
import sys
import time

TARGET_GLOBS = [
    "/usr/lib/node_modules/openclaw/dist/subagent-registry-*.js",
    "/usr/lib/node_modules/openclaw/dist/reply-*.js",
    "/usr/lib/node_modules/openclaw/dist/pi-embedded-*.js",
    "/usr/lib/node_modules/openclaw/dist/plugin-sdk/reply-*.js",
]

MIRROR_MARKER = "[cli-transcript] mirrored output to ${effectiveSessionKey}"
RUNTIME_EXPR = "(entry.endedAt ?? now) - (entry.startedAt ?? entry.createdAt)"
RUNTIME_EXPR_CLAMPED = "Math.max(0, (entry.endedAt ?? now) - (entry.startedAt ?? entry.createdAt))"

MIRROR_BLOCK = """\
\tconst cliProvider = result?.meta?.agentMeta?.provider;
\tconst shouldMirrorTranscript = Boolean(effectiveSessionKey) && (opts.lane === AGENT_LANE_SUBAGENT || cliProvider && isCliProvider(cliProvider, cfg));
\tif (shouldMirrorTranscript && effectiveSessionKey) {
\t\tconst providerLabel = cliProvider ?? "subagent";
\t\tconst modelLabel = result?.meta?.agentMeta?.model ?? "default";
\t\tconst transcriptText = (payloads ?? []).map((payload) => payload?.text ?? "").filter((line) => line.trim().length > 0).join("\\n\\n").trim() || `[${providerLabel}/${modelLabel}] completed with no text output`;
\t\ttry {
\t\t\tawait appendAssistantMessageToSessionTranscript({
\t\t\t\tagentId: resolveSessionAgentId({
\t\t\t\t\tsessionKey: effectiveSessionKey,
\t\t\t\t\tconfig: cfg
\t\t\t\t}),
\t\t\t\tsessionKey: effectiveSessionKey,
\t\t\t\ttext: transcriptText
\t\t\t});
\t\t\truntime.log(`[cli-transcript] mirrored output to ${effectiveSessionKey}`);
\t\t} catch (e) {
\t\t\truntime.error?.(`[cli-transcript] append failed for ${effectiveSessionKey}: ${String(e)}`);
\t\t\tif (!runtime.error) runtime.log(`[cli-transcript] append failed for ${effectiveSessionKey}: ${String(e)}`);
\t\t}
\t}
"""


def find_targets() -> list[pathlib.Path]:
    paths: list[pathlib.Path] = []
    seen = set()
    for pattern in TARGET_GLOBS:
        for p in glob.glob(pattern):
            rp = str(pathlib.Path(p).resolve())
            if rp in seen:
                continue
            seen.add(rp)
            paths.append(pathlib.Path(rp))
    return sorted(paths)


def inject_mirror_block(content: str) -> tuple[str, bool]:
    start = content.find("async function deliverAgentCommandResult(params) {")
    if start == -1:
        return content, False

    end = content.find("\nfunction resolveAgentRunContext(opts)", start)
    if end == -1:
        end = content.find("\n//#endregion", start)
    if end == -1:
        return content, False

    section = content[start:end]
    if MIRROR_MARKER in section:
        return content, False

    insert_at = section.find("if (!payloads || payloads.length === 0) {")
    if insert_at == -1:
        return content, False

    section_new = section[:insert_at] + MIRROR_BLOCK + section[insert_at:]
    return content[:start] + section_new + content[end:], True


def clamp_runtime_exprs(content: str) -> tuple[str, int]:
    out = []
    i = 0
    changed = 0
    while True:
        j = content.find(RUNTIME_EXPR, i)
        if j == -1:
            out.append(content[i:])
            break
        out.append(content[i:j])
        # If already wrapped as Math.max(0, <expr>) keep as-is.
        prefix = content[max(0, j - len("Math.max(0, ")):j]
        if prefix == "Math.max(0, ":
            out.append(RUNTIME_EXPR)
        else:
            out.append(RUNTIME_EXPR_CLAMPED)
            changed += 1
        i = j + len(RUNTIME_EXPR)
    return "".join(out), changed


def patch_content(content: str) -> tuple[str, dict]:
    changes = {"mirror_injected": False, "runtime_clamped": 0}

    content2, injected = inject_mirror_block(content)
    if injected:
        changes["mirror_injected"] = True

    content2, changed = clamp_runtime_exprs(content2)
    changes["runtime_clamped"] = changed

    return content2, changes


def node_check(path: pathlib.Path) -> tuple[bool, str]:
    proc = subprocess.run(["node", "--check", str(path)], text=True, capture_output=True)
    if proc.returncode == 0:
        return True, ""
    return False, (proc.stdout + proc.stderr).strip()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--apply", action="store_true", help="Write changes to files")
    parser.add_argument("--dry-run", action="store_true", help="Show what would change")
    parser.add_argument("--restart", action="store_true", help="Restart gateway service after successful apply")
    args = parser.parse_args()

    if not args.apply and not args.dry_run:
        parser.error("choose one: --apply or --dry-run")

    targets = find_targets()
    if not targets:
        print("No target dist files found.")
        return 1

    backups: list[tuple[pathlib.Path, pathlib.Path]] = []
    changed_files: list[pathlib.Path] = []

    print(f"Found {len(targets)} candidate files")

    for path in targets:
        try:
            original = path.read_text()
        except Exception as e:
            print(f"[skip] {path}: read failed: {e}")
            continue

        patched, changes = patch_content(original)
        changed = patched != original

        if not changed:
            print(f"[ok]   {path.name}: already patched / no applicable change")
            continue

        print(
            f"[patch] {path.name}: mirror_injected={changes['mirror_injected']} runtime_clamped={changes['runtime_clamped']}"
        )

        if args.dry_run:
            continue

        # apply
        ts = int(time.time())
        backup = pathlib.Path(f"{path}.bak.fix-{ts}")
        try:
            shutil.copy2(path, backup)
            backups.append((path, backup))
            path.write_text(patched)
            changed_files.append(path)
        except Exception as e:
            print(f"[error] failed writing {path}: {e}")
            # restore any already changed
            for dst, src in reversed(backups):
                try:
                    shutil.copy2(src, dst)
                except Exception:
                    pass
            return 2

    if args.dry_run:
        print("Dry run complete.")
        return 0

    # syntax check changed files
    for path in changed_files:
        ok, err = node_check(path)
        if not ok:
            print(f"[syntax-fail] {path}\n{err}")
            # rollback all changed files from backups
            for dst, src in reversed(backups):
                try:
                    shutil.copy2(src, dst)
                except Exception:
                    pass
            print("Rollback complete (restored from .bak.fix-* backups).")
            return 3

    print(f"Applied successfully to {len(changed_files)} files.")

    if args.restart:
        cmd = ["systemctl", "--user", "restart", "openclaw-gateway.service"]
        proc = subprocess.run(cmd, text=True, capture_output=True)
        if proc.returncode != 0:
            print("[warn] restart failed:")
            print((proc.stdout + proc.stderr).strip())
            return 4
        print("Gateway restarted (systemctl --user restart openclaw-gateway.service).")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())