# -*- coding: utf-8 -*-
"""send_campaign.py —— 真实发送**批量实验**（从仓库里跑，不再用 /tmp 临时脚本）。

和 `send_test.py` 的区别：这个是"打 N 发 + 落图 + 统计学"的实验脚手架，用来判定
`errorCode` 到底跟踪什么（答案对不对？点位准不准？），并把每次点击的**顺序**画在图上。

用法:
    # 基线：可复现靶子上连续跑 6 次图标题
    .venv/bin/python scripts/verify/send_campaign.py --target mliev --kind icon --n 6

    # 容差标定：把答案整体平移一个固定偏移（判定服务端命中半径有多紧）
    .venv/bin/python scripts/verify/send_campaign.py --target mliev --kind icon --n 2 --offset 30,30
    .venv/bin/python scripts/verify/send_campaign.py --target mliev --kind icon --n 2 --offset 0,60

    # 负对照：故意交换前两个点（必错 -> 期望 50）
    .venv/bin/python scripts/verify/send_campaign.py --target mliev --kind icon --n 2 --swap

产物（默认 /tmp/icon_campaign/<时间戳>/）:
    run{i}/bg.jpg          底图
    run{i}/prompt.jpg      提示图（仅图标/文字点选题）
    run{i}/annotated.jpg   底图 + **按点击顺序编号**的红圈（+ 估计的图标质心十字）
    run{i}/points.json     点位 / 每点离图标质心的距离 / errorCode
    summary.json           汇总（含每种 errorCode 计数）
"""

from __future__ import annotations

import argparse
import base64
import json
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional

import numpy as np

ROOT = Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import tcaptcha.protocol.client as client_mod                          # noqa: E402
import tcaptcha.solvers.icon_click as icon_mod                         # noqa: E402
from tcaptcha.protocol.client import (KIND_ICON_CLICK, KIND_SLIDER,    # noqa: E402
                                      KIND_TEXT_CLICK, TencentCaptchaClient)
from tcaptcha.render import annotate_points, icon_centroid_offset       # noqa: E402
from tcaptcha.targets import TARGETS, get_target                       # noqa: E402

KIND_MAP = {
    "text": (KIND_TEXT_CLICK,),
    "icon": (KIND_ICON_CLICK,),
    "slider": (KIND_SLIDER,),
    "any": (KIND_TEXT_CLICK, KIND_ICON_CLICK, KIND_SLIDER),
}

# 点位的"量尺"（点击点 vs 图标暗色笔画质心）与标注图渲染统一在 tcaptcha.render，
# 与 Web 界面共用同一份实现 —— 别在这里再抄一遍。


def main() -> int:
    ap = argparse.ArgumentParser(description="真实发送批量实验（打 N 发 + 落图 + 统计）")
    ap.add_argument("--target", default="mliev", choices=list(TARGETS))
    ap.add_argument("--kind", default="icon", choices=list(KIND_MAP))
    ap.add_argument("--n", type=int, default=6)
    ap.add_argument("--offset", default="0,0",
                    help="给答案整体平移 DX,DY 像素（容差标定用）")
    ap.add_argument("--swap", action="store_true", help="故意交换前两个点（负对照，必错）")
    ap.add_argument("--max-refreshes", type=int, default=40)
    ap.add_argument("--interval", type=float, default=1.0,
                    help="两次求解之间的间隔秒数（1ms 连打会触发频控，建议 3~5）")
    ap.add_argument("--outdir", default=None)
    ap.add_argument("--max-retries", type=int, default=1)
    args = ap.parse_args()

    tgt = get_target(args.target)
    aid, entry, referer, wb = tgt["aid"], tgt["entry_url"], tgt["referer"], tgt["wb"]
    dx, dy = (int(v) for v in args.offset.split(","))
    tag = "%s_%s_n%d_off%s%s" % (args.target, args.kind, args.n, args.offset,
                                 "_swap" if args.swap else "")
    # 默认落在**项目内**（har/ 只跟踪 .md，所以图片不会进 git，但会留在项目树里便于复核）
    outdir = Path(args.outdir or (ROOT / "har" / "1ms_dianxuan" / "icon_campaign" / tag))
    outdir.mkdir(parents=True, exist_ok=True)

    # ---- hook：抓每次求解的图 + 点位，并施加扰动
    state: Dict[str, object] = {}
    real_icon_solve = icon_mod.CaptchaSolver.solve

    def _as_pil(img):
        from PIL import Image
        import io as _io
        if isinstance(img, Image.Image):
            return img.convert("RGB")
        if isinstance(img, bytes):
            return Image.open(_io.BytesIO(img)).convert("RGB")
        return Image.open(img).convert("RGB")

    def _diag(self, bg, prompt) -> Dict[str, object]:
        """记录求解器看到的东西：提示目标类 + 每个候选的 top3 类概率。

        用来判"是不是把某个图标认成了别的类/挑错了同类的另一个实例"。
        """
        try:
            bgr, pim = _as_pil(bg), _as_pil(prompt)
            targets = self.parse_prompt_targets(pim)
            cands = self._extract_candidates(bgr)
            blobs = np.stack([self._preprocess_patch(c[2]) for c in cands],
                             axis=0).astype(np.float32)
            logits = self.cls_sess.run(None, {"input": blobs})[0]
            exp = np.exp(logits - np.max(logits, axis=1, keepdims=True))
            probs = exp / np.sum(exp, axis=1, keepdims=True)
            out = []
            for (cx, cy, _p), pr in zip(cands, probs):
                order = np.argsort(pr[:12])[::-1][:3]
                out.append({"xy": [int(cx), int(cy)],
                            "top": [[int(i), round(float(pr[i]), 3)] for i in order]})
            return {"targets": targets, "candidates": out}
        except Exception as exc:                                   # noqa: BLE001
            return {"diag_err": str(exc)[:120]}

    def spy(self, bg, prompt):
        pts = list(real_icon_solve(self, bg, prompt))
        state["bg_raw"], state["prompt_raw"] = bg, prompt
        state["diag"] = _diag(self, bg, prompt)
        if pts and (dx or dy):
            pts = [(x + dx, y + dy) for x, y in pts]
        if pts and args.swap and len(pts) >= 2:
            pts = [pts[1], pts[0]] + pts[2:]
        state["points"] = pts
        return pts

    icon_mod.CaptchaSolver.solve = spy

    real_gen = client_mod.generate_telemetry
    last_gen: Dict[str, object] = {}

    def hooked(tdc_code, points, **kw):
        out = real_gen(tdc_code, points, **kw)
        info: Dict[str, object] = {"source": out.get("source"),
                                   "collect_len": len(out.get("collect") or "")}
        try:
            from tcaptcha.telemetry.crypto import split_collect_segments, xtea_decrypt_bytes
            plain = xtea_decrypt_bytes(base64.b64decode(out["collect"]),
                                       out["key"]).decode("latin1").rstrip("\x00")
            info["segments"] = [len(s) for s in split_collect_segments(plain)]
        except Exception as exc:                                   # noqa: BLE001
            info["segments_err"] = str(exc)[:160]
            try:
                Path("/tmp/icon_campaign_last_fail.txt").write_bytes(plain.encode("latin1"))
            except Exception:                                      # noqa: BLE001
                pass
        last_gen.clear()
        last_gen.update(info)
        return out

    client_mod.generate_telemetry = hooked

    c = TencentCaptchaClient(aid=aid, entry_url=entry, referer=referer, wb=wb,
                             accept_kinds=KIND_MAP[args.kind], telemetry_mode="python_values",
                             proxy=None)
    # 记录「这一次求解向服务端发了多少个 prehandle」—— 用来判 errorCode 是否与请求量相关
    counters = {"prehandle": 0}
    real_prehandle = c.prehandle

    def counting_prehandle(tkid):
        counters["prehandle"] += 1
        return real_prehandle(tkid)

    c.prehandle = counting_prehandle

    rows: List[Dict[str, object]] = []
    try:
        for n in range(1, args.n + 1):
            state.clear()
            counters["prehandle"] = 0
            t0 = time.time()
            try:
                res = c.solve(max_retries=args.max_retries, max_refreshes=args.max_refreshes)
            except Exception as exc:                               # noqa: BLE001
                rows.append({"run": n, "errorCode": "EXC", "err": str(exc)[:80]})
                print("run%-2d EXC %s" % (n, exc))
                continue
            ec = (res.get("raw_response") or {}).get("errorCode")
            pts = list(state.get("points") or [])
            bg_raw, prompt_raw = state.get("bg_raw"), state.get("prompt_raw")

            d = outdir / ("run%d" % n)
            d.mkdir(parents=True, exist_ok=True)
            offsets: List[Optional[float]] = []
            centroids: List[Optional[List[float]]] = []
            if bg_raw is not None:
                from PIL import Image
                import io as _io
                bg_pil = Image.open(_io.BytesIO(bg_raw)).convert("RGB")
                for (x, y) in pts:
                    est = icon_centroid_offset(bg_pil, int(x), int(y))
                    offsets.append(None if est is None else round(est[2], 1))
                    centroids.append(None if est is None else [round(est[0], 1), round(est[1], 1)])
                (d / "bg.jpg").write_bytes(bg_raw)
                (d / "annotated.jpg").write_bytes(
                    annotate_points(bg_raw, pts, prompt=prompt_raw, offsets=offsets))
            if prompt_raw is not None:
                (d / "prompt.jpg").write_bytes(prompt_raw)
            (d / "points.json").write_text(json.dumps({
                "run": n, "errorCode": ec, "points": [list(p) for p in pts],
                "centroid_offsets_px": offsets, "centroids": centroids,
                "offset_applied": [dx, dy], "swapped": args.swap,
                "segments": last_gen.get("segments") or last_gen.get("segments_err"),
                "solver_diag": state.get("diag"),
                "prehandles": counters["prehandle"],
            }, ensure_ascii=False, indent=1), encoding="utf-8")

            row = {"run": n, "errorCode": ec, "kind": res.get("kind"),
                   "instruction": res.get("instruction"),
                   "points": [list(p) for p in pts],
                   "collect_len": last_gen.get("collect_len"),
                   "off_px": offsets, "secs": round(time.time() - t0, 1),
                   "prehandles": counters["prehandle"]}
            rows.append(row)
            print("run%-2d ec=%-4s kind=%-10s 偏离质心px=%-20s prehandle=%-3d points=%s %.0fs" % (
                n, ec, res.get("kind") or "?", str(offsets), counters["prehandle"], pts,
                time.time() - t0))
            time.sleep(max(0.0, args.interval))
    finally:
        c.close()
        icon_mod.CaptchaSolver.solve = real_icon_solve
        client_mod.generate_telemetry = real_gen

    counts: Dict[str, int] = {}
    by_kind: Dict[str, Dict[str, int]] = {}
    for r in rows:
        ec_s = str(r.get("errorCode"))
        counts[ec_s] = counts.get(ec_s, 0) + 1
        k = str(r.get("kind") or "?")
        by_kind.setdefault(k, {})
        by_kind[k][ec_s] = by_kind[k].get(ec_s, 0) + 1
    all_off = [v for r in rows for v in (r.get("off_px") or []) if isinstance(v, (int, float))]
    summary = {
        "target": args.target, "kind": args.kind, "n": args.n,
        "offset_applied": [dx, dy], "swapped": args.swap,
        "errorCode_counts": counts,
        "by_kind": by_kind,
        "centroid_offset_px": {
            "count": len(all_off),
            "min": min(all_off) if all_off else None,
            "max": max(all_off) if all_off else None,
            "mean": round(float(np.mean(all_off)), 1) if all_off else None,
        },
        "rows": rows,
    }
    (outdir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=1),
                                        encoding="utf-8")
    print("\n" + "=" * 74)
    print("靶子=%s 题型=%s 次数=%d 答案平移=%s 交换=%s" % (args.target, args.kind, args.n,
                                                        args.offset, args.swap))
    print("errorCode 计数: %s" % json.dumps(counts, ensure_ascii=False))
    if len(by_kind) > 1 or "?" not in by_kind:
        for k in sorted(by_kind):
            print("  按题型 %-12s %s" % (k, json.dumps(by_kind[k], ensure_ascii=False)))
    if all_off:
        print("点击点偏离图标笔画质心: 均值 %.1fpx / 最小 %.1f / 最大 %.1f（%d 个点）"
              % (np.mean(all_off), min(all_off), max(all_off), len(all_off)))
    print("产物: %s" % outdir)
    print("=" * 74)
    return 0


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