# -*- coding: utf-8 -*-
"""WorkBuddy 会话日志洪水 —— 一行补丁（就地 · 等长 · 可回滚）

【问题】
    会话诊断日志里 97.7% 的行是 `event-machine:dispatch`，其中 95.8% 是同一条
    `input:"tool_call_update"`；平均每次工具调用产生约 200 帧、约 48 KB。
    病根：抑制"流式帧"的集合只列了 agent_message_chunk / agent_thought_chunk，
    漏掉 tool_call / tool_call_update，于是逐帧噪声全部落盘。

【本补丁怎么修】
    把集合 STREAM_CHUNK_UPDATES 换成"任何 kind 都命中"的对象：
        new Set(["agent_message_chunk", "agent_thought_chunk"])
        -> {has:()=>true}                       （补空格到等长 55 字节）
    于是只在 turn 收口帧返回 true，逐帧噪声不再落盘（体积降约 96%）。

【为什么是"等长"替换】
    app.asar 头部记录每个文件的 size/offset，并带 per-file 完整性块。
    改动长度会让其后所有文件偏移错位，必须重打包。
    本补丁把 55 字节原样替换成 55 字节，头部 size/offset 全部不变；
    只把被改文件自己的完整性校验（1 个总哈希 + N 个分块哈希）原地等长重算。

【用法】
    python logpatch.py --check     只看，不改（含完整性核对）
    python logpatch.py             打补丁（自动备份 .bak-<日期>）
    python logpatch.py --revert    从备份还原（自动找最新的 .bak-*）
    python logpatch.py --verify    只核对完整性
    python logpatch.py --asar <路径>   指定其它 app.asar

【重要】
    · 打补丁后必须重启 WorkBuddy 才生效。
    · 完整性记录必须与被改文件一致，否则有加载失败风险；本脚本会改完立即复核。
"""
from __future__ import annotations

import argparse
import datetime as _dt
import glob
import hashlib
import json
import os
import re
import shutil
import sys

OLD = b'new Set(["agent_message_chunk", "agent_thought_chunk"])'
NEW = b'{has:()=>true}' + b" " * (len(OLD) - len(b'{has:()=>true}'))
assert len(NEW) == len(OLD), (len(OLD), len(NEW))

BLOCK_SIZE = 4194304
HEX64 = re.compile(rb'"[0-9a-fA-F]{64}"')

DEFAULT_ASAR = os.path.join(
    os.environ.get("LOCALAPPDATA", r"C:\Users\Administrator\AppData\Local"),
    "Programs", "WorkBuddy", "resources", "app.asar")


_SYM = {"✅": "[OK]", "⛔": "[!!]", "⚠️": "[!]", "⚠": "[!]", "🔴": "[!]",
        "·": "-", "⇒": "=>", "—": "-", "－": "-"}


def log(msg=""):
    """控制台走 PEP528（Unicode 正确）；重定向走 UTF-8；符号降级为 ASCII 以免控制台显示方框。"""
    s = str(msg)
    for k, v in _SYM.items():
        s = s.replace(k, v)
    try:
        if sys.stdout.isatty():
            sys.stdout.write(s + "\n")
        else:
            sys.stdout.buffer.write((s + "\n").encode("utf-8", "replace"))
        sys.stdout.flush()
    except Exception:
        try:
            sys.stdout.write(s.encode("ascii", "replace").decode() + "\n")
            sys.stdout.flush()
        except Exception:
            pass


# ---------------------------------------------------------------- 头部
def find_header(raw: bytes):
    """返回 (header_json_start, header_len, data_base)"""
    start = raw.find(b'{"files"')
    if start < 0:
        raise RuntimeError('认不出 asar 头部（未找到 {"files"）')
    _, consumed = json.JSONDecoder().raw_decode(raw[start:start + 40_000_000].decode("utf-8"))
    base = start + consumed
    while base % 4:
        base += 1
    return start, consumed, base


def build_index(meta):
    files = []

    def walk(node, prefix):
        for name, ent in (node.get("files") or {}).items():
            p = prefix + "/" + name if prefix else name
            if "files" in ent:
                walk(ent, p)
            else:
                try:
                    off = int(ent.get("offset", "0"))
                except Exception:
                    off = -1
                files.append({"path": p, "offset": off,
                              "size": int(ent.get("size", 0) or 0),
                              "unpacked": bool(ent.get("unpacked")),
                              "integrity": ent.get("integrity")})
    walk(meta, "")
    return files


def owner_packed(files, rel_off):
    """返回**真正**拥有该偏移的打包文件（排除 unpacked：它们 offset=0 会覆盖全区间）。

    补丁要点：旧版本只用 "off <= rel < off+size" 取第一个命中，
    结果被 unpacked 大文件（offset=0）抢走归属 -> 完整性记录漏改。
    """
    cands = [f for f in files
             if (not f["unpacked"]) and f["offset"] >= 0
             and f["offset"] <= rel_off < f["offset"] + f["size"]]
    if not cands:
        return None
    return min(cands, key=lambda f: f["size"])  # 最具体者


# ---------------------------------------------------------------- 完整性
def calc_integrity(blob: bytes):
    return hashlib.sha256(blob).hexdigest(), [
        hashlib.sha256(blob[i:i + BLOCK_SIZE]).hexdigest()
        for i in range(0, len(blob), BLOCK_SIZE)]


def read_content(raw, base, f):
    return raw[base + f["offset"]: base + f["offset"] + f["size"]]


def verify_entry(blob: bytes, integ: dict):
    """返回 (ok|None, 说明)"""
    if not integ:
        return None, "无 integrity 字段"
    if str(integ.get("algorithm", "SHA256")).upper() != "SHA256":
        return None, "算法非 SHA256"
    gh, gb = calc_integrity(blob)
    ok_h = (integ.get("hash") == gh)
    stored = list(integ.get("blocks") or [])
    if len(blob) == 0:
        # 空文件：只认总哈希（分块表示法各家不一，避免假阳性）
        return ok_h, "空文件 hash=%s" % ("OK" if ok_h else "MISMATCH")
    n_ok = sum(1 for a, b in zip(stored, gb) if a == b)
    return (ok_h and n_ok == len(gb) and len(gb) == len(stored)), \
        "hash=%s blocks=%d/%d" % ("OK" if ok_h else "MISMATCH", n_ok, len(gb))


def audit_packed(raw, base, files, only=None):
    """对全部打包文件（或指定集合）做完整性核对；返回 (通过数, 失败列表, 跳过数)"""
    ok_n, skipped, bad = 0, 0, []
    for f in files:
        if only is not None and f["path"] not in only:
            continue
        if not f["integrity"] or f["unpacked"]:
            continue
        if f["offset"] < 0 or base + f["offset"] + f["size"] > len(raw):
            skipped += 1
            continue
        ok, detail = verify_entry(read_content(raw, base, f), f["integrity"])
        if ok is None:
            skipped += 1
        elif ok:
            ok_n += 1
        else:
            bad.append("%s -> %s" % (f["path"], detail))
    return ok_n, bad, skipped


def set_integrity_inplace(buf: bytearray, raw: bytes, start, hlen, base, f):
    """按文件条目（用其唯一 offset 做锚点）就地等长重写 hash / blocks。"""
    """就地重写该文件的 hash / blocks（等长）。成功返回 (True, 说明)。"""
    blob = bytes(buf[base + f["offset"]: base + f["offset"] + f["size"]])
    new_hash, new_blocks = calc_integrity(blob)

    anchor = b'"%d"' % f["offset"]
    if raw.count(anchor) != 1:
        return False, "锚点 %s 不唯一，拒绝改动" % anchor.decode()
    apos = raw.find(anchor)
    ipos = raw.find(b'"integrity":{', apos, start + hlen)
    if ipos < 0 or ipos - apos > 400:
        return False, "定位该文件 integrity 条目失败"

    hp = raw.find(b'"hash":"', ipos, start + hlen)
    if hp < 0:
        return False, "找不到 hash 字段"
    hs = hp + len(b'"hash":"')
    if not all(c in b"0123456789abcdefABCDEF" for c in raw[hs:hs + 64]):
        return False, "hash 字段非 64 位十六进制"
    buf[hs:hs + 64] = new_hash.encode()

    bp = raw.find(b'"blocks":[', ipos, start + hlen)
    if bp < 0:
        return False, "找不到 blocks 数组"
    be = raw.find(b"]", bp, start + hlen)
    if be < 0:
        return False, "blocks 数组未闭合"
    found = list(HEX64.finditer(bytes(raw[bp:be])))
    if len(found) != len(new_blocks):
        return False, "分块数不符（头部 %d / 实算 %d）" % (len(found), len(new_blocks))
    for m, nb in zip(found, new_blocks):
        s = bp + m.start() + 1
        buf[s:s + 64] = nb.encode()
    return True, "hash 已更新 + %d/%d 个分块" % (len(new_blocks), len(new_blocks))


def find_backup(asar):
    d, b = os.path.dirname(asar), os.path.basename(asar)
    cands = glob.glob(os.path.join(d, b + ".bak-*"))
    if not cands:
        return None
    return max(cands, key=os.path.getmtime)


# ---------------------------------------------------------------- main
def main() -> int:
    ap = argparse.ArgumentParser(add_help=True)
    ap.add_argument("--asar", default=DEFAULT_ASAR)
    ap.add_argument("--check", action="store_true", help="只看不改")
    ap.add_argument("--verify", action="store_true", help="只核对完整性")
    ap.add_argument("--revert", action="store_true", help="从备份还原")
    a = ap.parse_args()

    asar = a.asar
    log("=" * 62)
    log("WorkBuddy 会话日志补丁 · %s" % _dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
    log("=" * 62)
    if not os.path.isfile(asar):
        log("⛔ 找不到 app.asar：%s" % asar)
        return 1
    log("目标 = %s" % asar)

    # ---- 还原 ----
    if a.revert:
        bak = find_backup(asar)
        if not bak:
            log("⛔ 没有找到备份（app.asar.bak-*），无法还原")
            return 1
        log("备份 = %s（%d 字节）" % (bak, os.path.getsize(bak)))
        shutil.copy2(bak, asar)
        ok = os.path.getsize(asar) == os.path.getsize(bak)
        log("✅ 已还原（大小一致：%s）" % ok)
        log("")
        log("⇒ 请重启 WorkBuddy 生效。")
        return 0

    raw = open(asar, "rb").read()
    size0 = len(raw)
    start, hlen, base = find_header(raw)
    meta = json.loads(raw[start:start + hlen].decode("utf-8"))
    files = build_index(meta)
    log("大小 = %d 字节；条目 %d 个；数据区起点 %d" % (size0, len(files), base))

    # ---- 只核对 ----
    if a.verify:
        ok_n, bad, skipped = audit_packed(raw, base, files)
        log("")
        log("---- 完整性核对（仅打包文件）----")
        log("✅ 通过 = %d ｜ ⛔ 不通过 = %d ｜ 跳过 = %d" % (ok_n, len(bad), skipped))
        for b in bad:
            log("   ⛔ %s" % b)
        return 1 if bad else 0

    # ---- 找补丁点 ----
    hits = []
    i = base
    while True:
        i = raw.find(OLD, i)
        if i < 0:
            break
        if raw[i + len(OLD): i + len(OLD) + 1] == b";":   # 只认"定义"，跳过 .has(...) 用法
            hits.append(i)
        i += 1
    log("命中「集合定义」 %d 处（等长替换 %d -> %d 字节）" % (len(hits), len(OLD), len(NEW)))

    owners, unresolved = [], []
    for h in hits:
        ow = owner_packed(files, h - base)
        if ow is None:
            unresolved.append(h)
            log("   @%d  归属 = ⛔ 未找到打包文件（将被跳过并告警）" % h)
        else:
            if ow["path"] not in [o["path"] for o in owners]:
                owners.append(ow)
            log("   @%d  归属 = %s (unpacked=False)" % (h, ow["path"]))

    if a.check:
        log("")
        log("---- 当前完整性核对（仅打包文件）----")
        ok_n, bad, skipped = audit_packed(raw, base, files)
        log("✅ 通过 = %d ｜ ⛔ 不通过 = %d ｜ 跳过 = %d" % (ok_n, len(bad), skipped))
        for b in bad[:20]:
            log("   ⛔ %s" % b)
        log("")
        log("⇒ --check 模式，未做任何修改。")
        return 1 if bad else 0
    if not hits:
        log("")
        log("⇒ 未发现可替换的定义：已经是打过补丁的状态（或版本不同）。什么都不做。")
        return 0

    if unresolved:
        log("")
        log("⛔ 有 %d 处命中找不到归属文件，为避免改坏别的文件，**中止，未写盘**。" % len(unresolved))
        return 1

    # ---- 备份 ----
    stamp = _dt.date.today().strftime("%Y%m%d")
    bak = "%s.bak-%s" % (asar, stamp)
    if not os.path.isfile(bak):
        shutil.copy2(asar, bak)
        log("✅ 已备份 → %s" % bak)
    else:
        log("· 备份已存在，复用 → %s" % bak)

    # ---- 改字节 ----
    buf = bytearray(raw)
    for h in hits:
        buf[h:h + len(OLD)] = NEW

    # ---- 改完整性（逐个受影响文件）----
    for ow in owners:
        ok, msg = set_integrity_inplace(buf, raw, start, hlen, base, ow)
        log("   · %s → %s" % (ow["path"], ("✅ " + msg) if ok else ("⛔ " + msg)))
        if not ok:
            log("⛔ 完整性重算失败，**中止，未写盘**。")
            return 1

    if len(buf) != size0:
        log("⛔ 长度变化（%d -> %d），中止，未写盘" % (size0, len(buf)))
        return 1

    # ---- 内存自检 ----
    b2 = bytes(buf)
    s2, h2, base2 = find_header(b2)
    json.loads(b2[s2:s2 + h2].decode("utf-8"))
    ok_n, bad, skipped = audit_packed(b2, base2, build_index(json.loads(
        b2[s2:s2 + h2].decode("utf-8"))), only={o["path"] for o in owners})
    log("")
    log("---- 落盘前自检 ----")
    log("长度未变：True（%d）｜ 头部可解析：True" % len(b2))
    log("被改文件完整性：%s" % ("✅ 全部通过" if not bad else "⛔ " + "; ".join(bad)))
    if bad:
        log("⛔ 自检未通过，**中止，未写盘**。")
        return 1

    with open(asar, "wb") as fp:
        fp.write(b2)

    chk = open(asar, "rb").read()
    log("")
    log("---- 落盘后复核 ----")
    log("大小未变：%s（%d）" % (len(chk) == size0, len(chk)))
    log("旧串残留：%d 处（另一处 .has(...) 用法不在本补丁范围）" % chk.count(OLD))
    log("新串出现：%d 处" % chk.count(NEW))
    try:
        s3, h3, _ = find_header(chk)
        json.loads(chk[s3:s3 + h3].decode("utf-8"))
        log("头部仍可解析：True")
    except Exception as e:
        log("⛔ 头部解析失败：%r" % e)
        return 1
    log("")
    log("✅ 完成。请重启 WorkBuddy 生效。")
    log("   还原：python %s --revert" % os.path.basename(__file__))
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except SystemExit:
        raise
    except Exception as exc:
        log("⛔ 异常：%r" % (exc,))
        raise SystemExit(1)