#!/usr/bin/env python3
"""Exercise existing diagnostic service, aggregate existing per-layer traces."""
import collections
import json
from pathlib import Path
import re
import urllib.request

opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
base = "http://127.0.0.1:19094"
trace = Path("/tmp/ling3-memory-audit-trace.log")
cases = [("digits_128", "1" * 107), ("digits_1024", "1" * 1003),
         ("office_text", "本次会议讨论项目进度、客户需求、交付时间与风险控制。请整理待办事项，明确责任人并安排下一次会议。" * 25)]
for name, text in cases:
    opener.open(urllib.request.Request(base + "/v1/reset", b""), timeout=20).read()
    offset = trace.stat().st_size
    response = opener.open(urllib.request.Request(base + "/v1/chat?max_tokens=2", text.encode()), timeout=300).read()
    events = [json.loads(line) for line in response.splitlines()]
    with trace.open() as f:
        f.seek(offset)
        lines = f.read().splitlines()
    total = collections.Counter()
    current = None
    block = collections.Counter()
    blocks = []
    for line in lines:
        m = re.search(r"batch_attention_ms=([\d.]+) batch_ffn_ms=([\d.]+)", line)
        if m:
            current = (float(m[1]), float(m[2]))
        m = re.search(r"batch_layer=(\d+) ms=([\d.]+)", line)
        if m and current:
            layer = int(m[1]); attn, ffn = current
            typ = "mla_ms" if layer % 4 == 3 else "kda_ms"
            total[typ] += attn; total["ffn_ms"] += ffn
            block[typ] += attn; block["ffn_ms"] += ffn
            if layer == 23:
                blocks.append(dict(block)); block.clear()
            current = None
    print(json.dumps({"case": name, "events": events, "batch_phase_ms": total, "blocks": blocks}, ensure_ascii=False), flush=True)