# -*- coding: utf-8 -*-
"""
逐个探测 config 里的 OCR 端点, 看有多少台【真的在服务】。
用来确认"120台只活了20台"这类问题——死端点越多, 活端点被压得越狠 → 超时。

用法:
  python3 tools/查端点存活.py                 # 自动找一张真实帧图做OCR探测
  python3 tools/查端点存活.py --tcp            # 只测TCP端口连通(更快, 不发OCR)
  python3 tools/查端点存活.py --timeout 5
  python3 tools/查端点存活.py --list-dead       # 额外列出死端点IP
"""
import argparse
import base64
import glob
import os
import socket
import sys
import time
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import requests
from config.config_utils import config

OCR_PATH = "/youtu/ocrapi/hpgeneralocr_v3"
DEFAULT_PORT = 18081


def norm(e):
    e = str(e).strip().rstrip("/")
    if not e:
        return None
    if not e.startswith("http"):
        e = "http://" + e
    host = e.split("://", 1)[1]
    if ":" not in host:
        e = e + ":" + str(DEFAULT_PORT)
    return e


def host_port(ep):
    h = ep.split("://", 1)[1]
    host, _, port = h.partition(":")
    return host, int(port or DEFAULT_PORT)


def find_image_b64():
    """找一张真实帧图转 base64; 找不到返回 None(改用TCP探测)。"""
    for pat in ("./split/**/*.jpg", "./composite/**/*.jpg", "./split/*/*.jpg"):
        for p in glob.glob(pat, recursive=True):
            try:
                with open(p, "rb") as f:
                    return base64.b64encode(f.read()).decode()
            except OSError:
                continue
    return None


def probe_ocr(ep, img_b64, timeout):
    url = ep + OCR_PATH
    payload = {"app_id": str(int(time.time() * 1000)), "image": img_b64,
               "ocr_template": "OCR", "options": {"pdf_page_index": 0, "detect_text": True}}
    t0 = time.time()
    try:
        r = requests.post(url, json=payload, timeout=timeout,
                          headers={"Content-Type": "application/json"})
        dt = time.time() - t0
        return (ep, "alive", dt, "") if r.status_code == 200 else (ep, "http_err", dt, f"HTTP {r.status_code}")
    except requests.exceptions.ConnectTimeout:
        return (ep, "dead", time.time() - t0, "连接超时")
    except requests.exceptions.ReadTimeout:
        return (ep, "slow", time.time() - t0, "读超时(连得上不回)")
    except requests.exceptions.ConnectionError:
        return (ep, "dead", time.time() - t0, "拒绝/不可达")
    except Exception as e:
        return (ep, "dead", time.time() - t0, str(e)[:40])


def probe_tcp(ep, timeout):
    host, port = host_port(ep)
    t0 = time.time()
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return (ep, "alive", time.time() - t0, "TCP通")
    except socket.timeout:
        return (ep, "dead", time.time() - t0, "TCP超时")
    except Exception as e:
        return (ep, "dead", time.time() - t0, str(e)[:40])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--timeout", type=float, default=5.0)
    ap.add_argument("--tcp", action="store_true", help="只测TCP端口连通(不发OCR)")
    ap.add_argument("--list-dead", action="store_true")
    args = ap.parse_args()

    eps = config.get_tencent_ocr_config().get("endpoints") or []
    eps = [u for u in (norm(e) for e in eps) if u]

    img = None
    if not args.tcp:
        img = find_image_b64()
        if img is None:
            print("没找到帧图, 改用 TCP 连通探测(--tcp)。")
    mode = "TCP连通" if (args.tcp or img is None) else "真实OCR请求"
    print(f"探测 {len(eps)} 个端点 (方式: {mode}, 超时 {args.timeout}s)...")

    def work(e):
        if args.tcp or img is None:
            return probe_tcp(e, args.timeout)
        return probe_ocr(e, img, args.timeout)

    res = []
    with ThreadPoolExecutor(max_workers=64) as ex:
        for f in as_completed([ex.submit(work, e) for e in eps]):
            res.append(f.result())

    c = Counter(r[1] for r in res)
    alive = [r for r in res if r[1] in ("alive", "http_err")]
    lats = sorted(r[2] for r in alive)
    total = len(eps)
    good = c["alive"] + c["http_err"]
    print("=" * 52)
    print(f"  可用(alive/http_err): {good}")
    print(f"  慢(读超时):           {c['slow']}")
    print(f"  死(连不上/拒绝):      {c['dead']}")
    if lats:
        print(f"  可用端点延迟: 中位 {lats[len(lats)//2]*1000:.0f}ms  最慢 {lats[-1]*1000:.0f}ms")
    print("=" * 52)
    print(f"结论: {total} 个端点里 {good} 个可用 ({good/total*100:.0f}%)")
    if good < total * 0.8:
        print(f"⚠️ {total-good} 个端点不可用! 负载全压在 {good} 个上 → 这极可能就是OCR超时根因。")
        print(f"   把死端点服务拉起 / 或从config删掉死的, OCR容量最高可提升 {total/max(good,1):.1f} 倍。")
    else:
        print("端点基本都活着 → OCR超时不是'端点少', 是'整体容量/CPU抢占'问题。")

    if args.list_dead:
        print("\n死/慢端点:")
        for ep, st, dt, msg in sorted(res, key=lambda x: x[1]):
            if st in ("dead", "slow"):
                print(f"  {st:6s} {ep:32s} {msg}")


if __name__ == "__main__":
    main()