# -*- encoding: utf-8 -*-
'''
@File    :   tencentcloud_ocr.py
@Time    :   @
@Author  :   eason
@Version :   1.0
@Contact :
@License :   (C)Copyright 2021-2024
@Desc    :   将图片转换为base64编码，并使用内部Youtu OCR服务进行文字识别（app_id自动生成）
'''

import json
import os
import base64
import random
import threading
import time
import requests
from config.config_utils import config   # 如果不再需要配置可以移除，但保留其他配置可能使用


# ---- OCR 端点（多 IP 负载均衡）配置 ----
OCR_PATH = "/youtu/ocrapi/hpgeneralocr_v3"
OCR_DEFAULT_PORT = 18081   # config 里只填 IP 时自动补的端口
DEFAULT_ENDPOINTS = ["http://asr.ocr.general.smart.l20.production.polaris:18081"]
WEIGHTS_FILE = "./config/ocr_weights.json"


def normalize_endpoint(e):
    """把 'IP' / 'IP:port' / 'http://IP:port' 统一成 'http://IP:port'。
    config 里只填裸 IP 时，自动补 http:// 和默认端口。"""
    if not e:
        return None
    e = str(e).strip().rstrip("/")
    if not e:
        return None
    if not (e.startswith("http://") or e.startswith("https://")):
        e = "http://" + e
    host = e.split("://", 1)[1].split("/", 1)[0]   # 取 host[:port] 部分
    if ":" not in host:                            # 没写端口就补默认端口
        e = e + f":{OCR_DEFAULT_PORT}"
    return e


def img2base64(img_path: str) -> str:
    """将图片转换为base64编码"""
    with open(img_path, 'rb') as f:
        image_base64 = base64.b64encode(f.read()).decode()
    return image_base64


def get_ocr_endpoints():
    """从配置读取 OCR 端点列表，自动归一化（只填 IP 也行）；缺省回退到内置默认。"""
    ocr_cfg = config.get_tencent_ocr_config() or {}
    eps = ocr_cfg.get("endpoints")
    if isinstance(eps, str):
        eps = [eps]
    if not eps:
        eps = DEFAULT_ENDPOINTS
    return [u for u in (normalize_endpoint(e) for e in eps) if u]


def load_endpoint_weights():
    """读取压测产出的 QPS 权重文件 {base_url: qps}；不存在则返回 {}。"""
    if not os.path.exists(WEIGHTS_FILE):
        return {}
    try:
        with open(WEIGHTS_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
        return {k.rstrip("/"): float(v) for k, v in data.items() if float(v) > 0}
    except (json.JSONDecodeError, ValueError, TypeError, OSError):
        return {}


class OcrDispatcher:
    """线程安全的 OCR 端点分发器：用 P-of-d-Choices(默认d=3)做【负载/延迟感知】派发——
    随机取 sample_size 个端点, 把请求发给其中【在途更少 × 更快】的那个。慢/积压端点自动少分。
    每个进程独立按后端真实快慢反应 → 无需任何协调/中心LB, 也能在 120 台 OCR 上涌现出全局均衡。
    d 从 2 提到 3(甚至更多): 端点数越多, 随机抽2个都恰好较差的概率越高, 抽样数多一个能明显
    收窄这个方差, 之后收益递减, 不用无限往上加。
    失败端点冷却摘除(冷却时长带随机抖动, 防止多个进程对同一端点同时解冷却又一起打回去,
    造成"整体一起失败-一起恢复-一起打爆"的同步震荡)。多个并发 OCR 线程共享一个实例。"""

    def __init__(self, endpoints, weights=None, cooldown=30.0, sample_size=3):
        self.endpoints = list(endpoints)
        # PID 打乱(历史原因保留): 让各进程初始顺序不同; 现在 P-of-d 是随机取样, 影响不大但无害。
        try:
            random.Random(os.getpid()).shuffle(self.endpoints)
        except Exception:
            pass
        weights = weights or {}
        # 压测权重(QPS): 作为"处理能力", 强端点 score 更低 → 多分。无权重默认 1.0。
        self.weights = {e: max(float(weights.get(e, 1.0)), 0.0001) for e in self.endpoints}
        self.cooldown = cooldown
        self.sample_size = max(2, int(sample_size))    # 每次候选抽样个数(P-of-d 的 d)
        self._lock = threading.Lock()
        self._inflight = {e: 0 for e in self.endpoints}        # 当前在途请求数(本进程视角)
        self._lat = {e: 0.05 for e in self.endpoints}          # 延迟 EWMA(秒), 初始乐观→新端点也会被试
        self._cooldown_until = {e: 0.0 for e in self.endpoints}
        self._rng = random.Random(os.getpid() ^ 0x9E3779B1)
        # —— 派发计数(验证分发是否均匀): 每端点累计派发数, 每 _dump_every 次落盘一次本进程文件 ——
        self._count = {e: 0 for e in self.endpoints}
        self._since_dump = 0
        try:
            self._dump_every = int((config.get_tencent_ocr_config() or {}).get('dispatch_log_every', 0) or 0)
        except (ValueError, TypeError, AttributeError):
            self._dump_every = 0

    def _score(self, e):
        """负载分: (在途+1) × 延迟EWMA / 权重。越低越优先。"""
        return (self._inflight[e] + 1) * self._lat[e] / self.weights[e]

    def pick(self):
        """P-of-d: 随机不重复抽 sample_size(默认3) 个可用端点, 选负载分最低的;
        选中即占一个在途名额(必须配对 done())。全部冷却时退化为全部可用。"""
        now = time.time()
        snapshot = None
        with self._lock:
            available = [e for e in self.endpoints if self._cooldown_until[e] <= now]
            if not available:
                available = list(self.endpoints)        # 全挂了也得发, 给一次机会
            if len(available) == 1:
                chosen = available[0]
            else:
                k = min(self.sample_size, len(available))
                cands = self._rng.sample(available, k)   # 不重复抽样, 比各自 choice() 更有效(不会抽到同一个)
                chosen = min(cands, key=self._score)
            self._inflight[chosen] = self._inflight.get(chosen, 0) + 1
            self._count[chosen] = self._count.get(chosen, 0) + 1
            if self._dump_every > 0:
                self._since_dump += 1
                if self._since_dump >= self._dump_every:
                    self._since_dump = 0
                    snapshot = (dict(self._count),
                                {e: round(self._lat[e], 4) for e in self._lat})
        if snapshot is not None:                        # 落盘在锁外, 不阻塞其它线程派发
            self._write_dump(snapshot[0], snapshot[1])
        return chosen

    def _write_dump(self, counts, lats):
        """把本进程的累计派发计数写到 ./ocr_dispatch/<pid>.json(覆盖式, 单进程单文件无竞争)。"""
        try:
            d = "./ocr_dispatch"
            os.makedirs(d, exist_ok=True)
            path = os.path.join(d, f"{os.getpid()}.json")
            data = {"pid": os.getpid(), "ts": time.time(),
                    "total": sum(counts.values()), "counts": counts, "lat_ewma": lats}
            tmp = path + ".tmp"
            with open(tmp, "w", encoding="utf-8") as f:
                json.dump(data, f, ensure_ascii=False)
            os.replace(tmp, path)
        except Exception:
            pass

    def done(self, endpoint, latency=None, ok=True):
        """请求结束回调: 释放在途名额; 成功则更新延迟EWMA, 失败则冷却摘除(冷却时长带±30%抖动,
        避免同一时刻失败的多个进程在同一时刻解冷却、一起冲回去打同一个刚失败的端点)。"""
        with self._lock:
            if self._inflight.get(endpoint, 0) > 0:
                self._inflight[endpoint] -= 1
            if endpoint not in self._cooldown_until:
                return
            if ok:
                if latency is not None:
                    self._lat[endpoint] = 0.7 * self._lat[endpoint] + 0.3 * max(latency, 0.001)
                self._cooldown_until[endpoint] = 0.0
            else:
                jitter = self._rng.uniform(0.7, 1.3)
                self._cooldown_until[endpoint] = time.time() + self.cooldown * jitter

    # —— 兼容旧接口(内部都走 done) ——
    def report_failure(self, endpoint):
        self.done(endpoint, ok=False)

    def report_success(self, endpoint):
        self.done(endpoint, ok=True)


_dispatcher = None
_dispatcher_lock = threading.Lock()


def get_dispatcher():
    """惰性单例分发器（首次调用时读取端点+权重）。"""
    global _dispatcher
    if _dispatcher is None:
        with _dispatcher_lock:
            if _dispatcher is None:
                eps = get_ocr_endpoints()
                _dispatcher = OcrDispatcher(eps, load_endpoint_weights())
    return _dispatcher


def _build_payload(image_base64: str) -> dict:
    timestamp = int(time.time() * 1000)
    rand_suffix = random.randint(100000, 999999)
    app_id = f"{timestamp}{rand_suffix}"
    return {
        "app_id": app_id,
        "image": image_base64,
        "ocr_template": "OCR",
        "options": {"pdf_page_index": 0, "detect_text": True},
    }


# 跨进程【全局 OCR 并发上限】信号量。进程池下由 main 在父进程创建、经 initializer 注入每个 worker;
# 不管多少进程×多少线程,全局同时在飞的 OCR 请求数被这个信号量硬性卡住,防止打爆端点。
# 为 None 时不限流(线程池模式 / 未注入)。
_ocr_semaphore = None


def set_ocr_semaphore(sem):
    """由 worker 进程初始化时调用,注入全局并发信号量。"""
    global _ocr_semaphore
    _ocr_semaphore = sem


# 每个进程一个持久 Session(连接池/keep-alive 复用连接)。
# 原来用 requests.post() 每次新建+关闭连接 → 高并发下海量 TIME_WAIT → 本地临时端口耗尽
# → "Cannot assign requested address" / 连接超时(表现成"所有端点超时", 其实是本机端口用完)。
# 复用连接后, 到每个端点只保持少量长连接, 端口不再churn。惰性创建 → 进程池下各进程独立 Session。
_session = None
_session_lock = threading.Lock()


def _get_session():
    global _session
    if _session is None:
        with _session_lock:
            if _session is None:
                s = requests.Session()
                adapter = requests.adapters.HTTPAdapter(
                    pool_connections=256,   # 连接池数(≥端点数, 每端点一个池)
                    pool_maxsize=64,        # 每端点最大并发连接(复用)
                    max_retries=0,
                    pool_block=False)
                s.mount("http://", adapter)
                s.mount("https://", adapter)
                _session = s
    return _session


def ocr_request_once(base_url: str, image_base64: str, timeout: int = 30) -> str:
    """向指定端点发一次 OCR 请求，返回原始响应文本；失败抛异常。
    受全局信号量限流;用持久 Session 复用连接(防端口耗尽)。"""
    url = base_url.rstrip("/") + OCR_PATH
    headers = {"Content-Type": "application/json"}
    sem = _ocr_semaphore
    if sem is not None:
        sem.acquire()
    try:
        resp = _get_session().post(url, headers=headers, json=_build_payload(image_base64),
                                   timeout=timeout)
        resp.raise_for_status()
        return resp.text
    finally:
        if sem is not None:
            sem.release()


def _normalize_youtu_to_tencent(raw: str) -> str:
    """将 Youtu OCR 返回格式转为项目其他模块期望的腾讯云 OCR 格式。"""
    try:
        data = json.loads(raw)
    except (json.JSONDecodeError, TypeError):
        return raw

    # 已经是腾讯云格式则直接返回
    if "TextDetections" in data:
        return raw

    if data.get("error_code") != 0:
        return json.dumps({"TextDetections": [], "Language": "zh",
                           "RequestId": data.get("session_id", "")}, ensure_ascii=False)

    text_detections = []
    for recognize_item in data.get("recognize_list", []):
        item_list = recognize_item.get("item_content", {}).get("item_list", [])
        for item in item_list:
            content = item.get("content", "")
            confidence = int(item.get("conf", 0) * 100)  # 0-1 → 0-100

            vertex = item.get("vertex", {})
            if vertex:
                polygon = [
                    {"X": vertex.get("x1", 0), "Y": vertex.get("y1", 0)},
                    {"X": vertex.get("x2", 0), "Y": vertex.get("y2", 0)},
                    {"X": vertex.get("x3", 0), "Y": vertex.get("y3", 0)},
                    {"X": vertex.get("x4", 0), "Y": vertex.get("y4", 0)},
                ]
            else:
                coord = item.get("coord", {})
                x, y = coord.get("x", 0), coord.get("y", 0)
                w, h = coord.get("width", 0), coord.get("height", 0)
                polygon = [
                    {"X": x,     "Y": y},
                    {"X": x + w, "Y": y},
                    {"X": x + w, "Y": y + h},
                    {"X": x,     "Y": y + h},
                ]

            text_detections.append({
                "DetectedText": content,
                "Confidence": confidence,
                "Polygon": polygon,
            })

    return json.dumps({
        "TextDetections": text_detections,
        "Language": "zh",
        "RequestId": data.get("session_id", ""),
    }, ensure_ascii=False)


def _get_ocr_call_params():
    """从 config 读【单次OCR的重试上限/超时】, 带默认。
    封顶重试是为了防"卡1小时": 端点都超时时, 一张图最坏耗时 = max_attempts × timeout,
    而不是原来的 len(endpoints)(119) × 30s ≈ 3570s。"""
    try:
        cfg = config.get_tencent_ocr_config() or {}
        max_attempts = int(cfg.get("max_attempts", 4) or 4)
        timeout = int(cfg.get("timeout", 12) or 12)
    except (ValueError, TypeError, AttributeError):
        max_attempts, timeout = 4, 12
    return max(1, max_attempts), max(1, timeout)


def youtu_ocr_api(image_base64: str) -> str:
    """
    调用内部 Youtu OCR 服务进行文字识别。
    按 QPS 权重自动在多个 OCR IP 间分发；某端点失败时摘除冷却并换下一个端点重试。
    重试次数封顶(config: ocr.max_attempts), 单次超时可配(ocr.timeout), 避免端点全超时时拖死 worker。
    """
    dispatcher = get_dispatcher()
    max_attempts, timeout = _get_ocr_call_params()
    # 最多试 max_attempts 个端点(且不超过端点总数), 而不是把 119 个全试一遍
    attempts = min(max(1, len(dispatcher.endpoints)), max_attempts)
    last_err = None
    for _ in range(attempts):
        url = dispatcher.pick()
        t0 = time.time()
        try:
            raw = ocr_request_once(url, image_base64, timeout=timeout)
            # 把真实延迟反馈给分发器 → 慢端点延迟EWMA升高 → 后续自动少分(负载感知核心)
            dispatcher.done(url, latency=time.time() - t0, ok=True)
            return _normalize_youtu_to_tencent(raw)
        except requests.exceptions.RequestException as e:
            last_err = e
            dispatcher.done(url, latency=time.time() - t0, ok=False)
            continue
    print(f"Youtu OCR 请求失败（已尝试 {attempts} 个端点, 单次超时 {timeout}s）: {last_err}")
    return ""


# 为了保持原函数名兼容，将新接口赋值给原名称
tencentcloud_ocr_api = youtu_ocr_api


def tencent_ocr_video_all_frames(pic_dir_path: str,
                                 jsonl_save_path: str = 'frames_ocr.jsonl') -> list:
    """
    对一个视频抽帧后的全部图片进行 ocr处理并保存起来
    """
    with open(jsonl_save_path, "w", encoding='utf-8') as f:
        for root, _, files in os.walk(pic_dir_path):
            for pic_name in sorted(files):
                if pic_name.endswith(('jpg')):
                    pic_path = os.path.join(root, pic_name)
                    img_base64 = img2base64(img_path=pic_path)
                    json_string = tencentcloud_ocr_api(image_base64=img_base64)
                    f.write("pic_path: {}\n".format(pic_path))
                    f.write(json_string + "\n")