"""
会话身份：读昵称
================
企业微信全自绘，无障碍树里拿不到任何文本（实测整个主窗口只有 2 个控件、零
Name），所以过去只能拿头像感知哈希 + 名称像素哈希当身份。那套东西同一个人换
个渲染就对不上，同一个联系人被拆成好几份档案，客户从此收不到回复。

这里改成直接把昵称读成文字（「一个小迷糊@微信」「高瑞@微信」），会话身份就是
这串昵称的 md5，一个人只有一个身份。

为什么不用 OCR 自带的文字检测：检测模型找框要 1.6 秒，放不进轮询。而"找出这块
面板里文字在哪"本来就是像素分析的事，自己按墨迹算完框、只让模型做识别，实测
降到 20 毫秒左右，随便调。
"""

import hashlib
import re
import threading
import unicodedata

import numpy as np

# 识别结果低于这个分就当没读到。实测正常昵称在 0.92 以上，
# 裁歪或糊掉时会掉到 0.7 附近
MIN_CONFIDENCE = 0.80
# 与已认识的昵称最多差几个字仍算同一个人。OCR 偶尔会把「打卡」读成「江卡」，
# 差一个字；差两个字以上就不敢认了，宁可当新客户
MAX_CORRECTION_DISTANCE = 1
# 太短的名字差一个字就是另一个人（「小王」vs「小李」），不做纠正
MIN_LENGTH_FOR_CORRECTION = 4
# 画面指纹 → 昵称的缓存上限。同一行的未读态、选中态、悬停态各占一个条目，
# 机器人一跑就是一整天，不封顶迟早把内存吃光
MAX_CACHE_ENTRIES = 512

_WHITESPACE_RE = re.compile(r"\s+")


def normalize_name(value: str) -> str:
    """同一个昵称在标题和列表里要归一成同一串。

    实测标题读出来是「高瑞 @微信」、列表行是「高瑞@微信」，中间多个空格。
    不归一的话同一个人会算出两个 md5，又回到一人多档的老问题上。
    """
    text = unicodedata.normalize("NFKC", str(value or ""))
    text = _WHITESPACE_RE.sub("", text)
    return text.strip()


def session_id_for(name: str) -> str:
    """昵称 → 会话 ID。空名字不给 ID，宁可这一轮不认。"""
    normalized = normalize_name(name)
    if not normalized:
        return ""
    return hashlib.md5(normalized.encode("utf-8")).hexdigest()


def ink_bounds(region: np.ndarray, threshold: float = 45.0):
    """面板里文字实际占的矩形，(x1, y1, x2, y2)；没有文字返回 None。"""
    if region is None or getattr(region, "ndim", 0) != 3 or not region.size:
        return None
    gray = region[:, :, :3].mean(axis=2)
    if not gray.size:
        return None
    background = float(np.median(gray))
    ink = np.abs(gray - background) > threshold
    rows = np.flatnonzero(ink.any(axis=1))
    cols = np.flatnonzero(ink.any(axis=0))
    if not rows.size or not cols.size:
        return None
    return int(cols[0]), int(rows[0]), int(cols[-1]) + 1, int(rows[-1]) + 1


def crop_to_ink(region: np.ndarray, pad: int = 5):
    """把面板裁到只剩文字，四周留一点白边。

    留白不是可有可无：贴着字裁会把首字的一竖切掉，实测「一个小迷糊@微信」会
    读成「个小迷糊@微信」、「高瑞@微信」会读成「瑞@微信」。
    """
    bounds = ink_bounds(region)
    if bounds is None:
        return None
    x1, y1, x2, y2 = bounds
    height, width = region.shape[:2]
    return np.ascontiguousarray(
        region[
            max(0, y1 - pad):min(height, y2 + pad),
            max(0, x1 - pad):min(width, x2 + pad),
            :3,
        ]
    )


def _edit_distance(left: str, right: str, limit: int) -> int:
    """超过 limit 就提前退出，返回 limit + 1。"""
    if abs(len(left) - len(right)) > limit:
        return limit + 1
    previous = list(range(len(right) + 1))
    for i, lch in enumerate(left, 1):
        current = [i]
        for j, rch in enumerate(right, 1):
            current.append(
                min(
                    previous[j] + 1,
                    current[j - 1] + 1,
                    previous[j - 1] + (lch != rch),
                )
            )
        if min(current) > limit:
            return limit + 1
        previous = current
    return previous[-1]


class NameReader:
    """把一块画面读成昵称，并把偶发的一字误读纠回已认识的那个人。"""

    def __init__(self):
        self._engine = None
        self._engine_failed = False
        self._lock = threading.Lock()
        # RapidOCR/ONNX Runtime 的同一个会话不保证可被两个线程同时调用。模型
        # 等待期间的视觉观察线程会和主轮询共享 NameReader，因此识别本身也串行化。
        self._run_lock = threading.Lock()
        self._known = {}          # 归一昵称 -> 见过的次数
        self._cache = {}          # 画面指纹 -> 归一昵称

    # ── OCR 引擎 ─────────────────────────────────────────────────────────────
    def _ensure_engine(self):
        if self._engine is not None or self._engine_failed:
            return self._engine
        with self._lock:
            if self._engine is not None or self._engine_failed:
                return self._engine
            try:
                from rapidocr_onnxruntime import RapidOCR

                try:
                    # onnxruntime 默认按核数开满线程池且空转抢 CPU。识别的
                    # 都是几十像素高的小裁片，两个线程绰绰有余；不限的话
                    # 低配机每轮轮询都会被 OCR 抖一下。
                    self._engine = RapidOCR(
                        intra_op_num_threads=2, inter_op_num_threads=1
                    )
                except Exception:
                    self._engine = RapidOCR()
            except Exception as exc:
                self._engine_failed = True
                print(f"    [身份] OCR 引擎不可用，无法读取会话昵称: {exc}")
        return self._engine

    @property
    def available(self) -> bool:
        return self._ensure_engine() is not None

    # ── 读取 ─────────────────────────────────────────────────────────────────
    def read_raw(self, region: np.ndarray) -> tuple[str, float]:
        """读一块面板，返回（昵称原文, 置信度）。读不到返回 ("", 0.0)。"""
        crop = crop_to_ink(region)
        if crop is None or crop.shape[0] < 6 or crop.shape[1] < 6:
            return "", 0.0
        engine = self._ensure_engine()
        if engine is None:
            return "", 0.0
        try:
            # use_det=False：框已经自己算好了，别再花 1.6 秒让模型找一遍
            with self._run_lock:
                result, _elapse = engine(crop, use_det=False, use_cls=False)
        except Exception as exc:
            print(f"    [身份] 昵称识别失败: {exc}")
            return "", 0.0
        if not result:
            return "", 0.0
        try:
            return str(result[0][0] or ""), float(result[0][1] or 0.0)
        except (IndexError, TypeError, ValueError):
            return "", 0.0

    def read_layout(self, region: np.ndarray, max_results: int = 80) -> list[dict]:
        """低频读取一整块界面的文字框，供异常页面恢复和大模型文字兜底。

        日常昵称识别仍使用 20ms 的无检测路径；只有找不到“消息”入口或视觉模型
        不可用时才调用这里较慢的文字检测，因此不会拖慢正常轮询。
        """
        if region is None or getattr(region, "ndim", 0) != 3 or not region.size:
            return []
        engine = self._ensure_engine()
        if engine is None:
            return []
        try:
            pixels = np.ascontiguousarray(region[:, :, :3])
            with self._run_lock:
                result, _elapse = engine(pixels, use_det=True, use_cls=False)
        except Exception as exc:
            print(f"    [身份] 页面文字布局识别失败: {exc}")
            return []
        parsed = []
        for item in result or []:
            try:
                box, text, score = item[0], str(item[1] or ""), float(item[2] or 0.0)
                points = [[float(point[0]), float(point[1])] for point in box]
            except (IndexError, TypeError, ValueError):
                continue
            if not text.strip() or score < 0.45 or len(points) < 4:
                continue
            xs = [point[0] for point in points]
            ys = [point[1] for point in points]
            parsed.append({
                "text": text.strip(),
                "score": score,
                "x1": int(min(xs)),
                "y1": int(min(ys)),
                "x2": int(max(xs)),
                "y2": int(max(ys)),
            })
            if len(parsed) >= max(1, int(max_results)):
                break
        return parsed

    def read(self, region: np.ndarray, signature: bytes = b"") -> str:
        """读出归一后的昵称；读不出或不够可信返回空串。

        `signature` 是这块画面的像素指纹。同一块画面只识别一次，之后直接命中
        缓存——这让同一个会话在整轮里始终得到同一个身份，不会因为某次识别抖动
        而变成另一个人。
        """
        key = bytes(signature or b"")
        if key and key in self._cache:
            return self._cache[key]
        text, score = self.read_raw(region)
        name = normalize_name(text)
        if not name or score < MIN_CONFIDENCE:
            return ""
        name = self.canonical(name)
        self._known[name] = self._known.get(name, 0) + 1
        if key:
            if len(self._cache) >= MAX_CACHE_ENTRIES:
                # 先进先出即可：留下的都是最近见过的渲染，认错人的风险为零
                # ——缓存没命中只是多花 20 毫秒重识别一次
                for stale in list(self._cache)[: MAX_CACHE_ENTRIES // 4]:
                    self._cache.pop(stale, None)
            self._cache[key] = name
        return name

    # ── 一字误读纠正 ─────────────────────────────────────────────────────────
    def canonical(self, name: str) -> str:
        """把偶发误读纠回已经认识的那个昵称。

        实测同一批识别里「打卡」被读成「江卡」、「行业资讯」被读成「亏业资讯」。
        不纠的话，某一轮读错一个字就会被当成另一个客户、另开一份档案、上下文
        断掉。只纠差一个字的，且名字本身要够长——「小王」和「小李」也只差一个
        字，但那是两个人。
        """
        normalized = normalize_name(name)
        if not normalized or normalized in self._known:
            return normalized
        if len(normalized) < MIN_LENGTH_FOR_CORRECTION:
            return normalized
        best = ""
        best_distance = MAX_CORRECTION_DISTANCE + 1
        best_seen = -1
        for known, seen in self._known.items():
            if len(known) < MIN_LENGTH_FOR_CORRECTION:
                continue
            distance = _edit_distance(normalized, known, MAX_CORRECTION_DISTANCE)
            if distance > MAX_CORRECTION_DISTANCE:
                continue
            # 同样接近时认见得多的那个：偶发误读只会出现一两次
            if distance < best_distance or (
                distance == best_distance and seen > best_seen
            ):
                best, best_distance, best_seen = known, distance, seen
        if not best:
            return normalized
        print(f"    [身份] 昵称疑似误读，已按已认识的联系人纠正：{normalized!r} → {best!r}")
        return best

    def remember(self, name: str) -> None:
        """把一个确定无误的昵称登记为已认识（例如从档案里恢复的）。"""
        normalized = normalize_name(name)
        if normalized:
            self._known.setdefault(normalized, 1)

    def known_names(self) -> list[str]:
        return sorted(self._known)