#!/usr/bin/env python3
"""
QQ 图片文字识别 (OCR) — 本地 Tesseract + NapCat 回退

使用策略：
  1. 本地 Tesseract OCR（快速、可靠、无需 GUI）
  2. 如果 Tesseract 不可用，回退到 NapCat 的 /ocr_image 端点

跨主机文件传递（NapCat 回退路径）：
  宿主机 YOUR_SHARED_DIR/  →  NapCat 容器内 /app/napcat/share/

用法:
  python3 qq_ocr_image.py /tmp/screenshot.png           # 本地图片
  python3 qq_ocr_image.py --url https://example.com/a.png  # 远程图片
  python3 qq_ocr_image.py /tmp/image.png --lang eng        # 指定语言

语言参数:
  chi_sim      简体中文（默认）
  chi_tra      繁体中文
  eng          英文
  chi_sim+eng  中英文混合（推荐）

返回:
  JSON: {status, data: {texts: [{text, confidence}], full_text}}
"""

import json
import sys
import os
import argparse
import shutil
import time

# === 本地 Tesseract OCR ===
try:
    from PIL import Image
    import pytesseract
    TESSERACT_AVAILABLE = True
except ImportError:
    TESSERACT_AVAILABLE = False

# === NapCat 回退 ===
import requests

CQHTTP_URL = "http://YOUR_NAPCAT_HOST:25570"
SHARED_DIR = "YOUR_SHARED_DIR"
CONTAINER_FILES = "/app/napcat/share"
OCR_TIMEOUT = 60


def ocr_local(image_path: str, lang: str) -> dict:
    """
    使用本地 Tesseract OCR
    """
    if not TESSERACT_AVAILABLE:
        return None  # 走回退

    if not os.path.exists(image_path):
        return {
            "status": "failed",
            "message": f"文件不存在: {image_path}"
        }

    try:
        img = Image.open(image_path)
        full_text = pytesseract.image_to_string(img, lang=lang)
        # 也拿到详细数据（带置信度）
        data = pytesseract.image_to_data(img, lang=lang, output_type=pytesseract.Output.DICT)

        texts = []
        for i in range(len(data["text"])):
            text = data["text"][i].strip()
            conf = data["conf"][i]
            if text and conf >= 0:  # conf = -1 表示该区域无文本
                texts.append({
                    "text": text,
                    "confidence": int(conf),
                    "bbox": {
                        "x": data["left"][i],
                        "y": data["top"][i],
                        "w": data["width"][i],
                        "h": data["height"][i]
                    }
                })

        if not texts:
            # 没识别到文字，但 full_text 可能有内容
            lines = [l.strip() for l in full_text.strip().split("\n") if l.strip()]
            texts = [{"text": l, "confidence": 0} for l in lines]

        return {
            "status": "ok",
            "source": "tesseract",
            "data": {
                "texts": texts,
                "full_text": full_text.strip(),
                "language": lang
            }
        }

    except Exception as e:
        return {
            "status": "failed",
            "source": "tesseract",
            "message": f"Tesseract OCR 失败: {e}"
        }


def ocr_napcat(image_source: str) -> dict:
    """
    回退到 NapCat /ocr_image
    """
    # 本地文件 → 共享目录桥接
    if not (image_source.startswith("http://") or
            image_source.startswith("https://") or
            image_source.startswith("file://")):
        abs_path = os.path.abspath(image_source)
        if not os.path.exists(abs_path):
            return {
                "status": "failed",
                "message": f"文件不存在: {abs_path}"
            }
        ts = int(time.time())
        basename = os.path.basename(abs_path)
        target_name = f"{ts}_ocr_{basename}"
        target_path = os.path.join(SHARED_DIR, target_name)
        try:
            shutil.copy2(abs_path, target_path)
            image_source = f"file://{CONTAINER_FILES}/{target_name}"
        except Exception as e:
            return {
                "status": "failed",
                "message": f"复制文件到共享目录失败: {e}"
            }

    # 调用 NapCat API
    try:
        resp = requests.post(
            f"{CQHTTP_URL}/ocr_image",
            json={"image": image_source},
            timeout=OCR_TIMEOUT
        )
        resp_data = resp.json()
        if resp_data.get("status") == "ok":
            resp_data["source"] = "napcat"
        return resp_data
    except requests.exceptions.Timeout:
        return {
            "status": "failed",
            "source": "napcat",
            "message": f"NapCat OCR 超时（{OCR_TIMEOUT}秒），在 Docker 无 GUI 环境中不可用"
        }
    except requests.exceptions.ConnectionError as e:
        return {
            "status": "failed",
            "source": "napcat",
            "message": f"无法连接到 NapCat ({CQHTTP_URL}): {e}"
        }


def main():
    parser = argparse.ArgumentParser(description="QQ 图片文字识别 (OCR)")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("image", type=str, nargs="?",
                       help="本地图片路径或 URL")
    group.add_argument("--file", type=str,
                       help="本地图片文件路径")
    group.add_argument("--url", type=str,
                       help="远程图片 URL")

    parser.add_argument("--lang", type=str, default="chi_sim+eng",
                       help="识别语言（默认 chi_sim+eng，支持 eng / chi_sim / chi_tra / chi_sim+eng）")
    parser.add_argument("--force-napcat", action="store_true",
                       help="强制使用 NapCat OCR（跳过本地 Tesseract）")

    args = parser.parse_args()

    if args.file:
        source = args.file
    elif args.url:
        source = args.url
    else:
        source = args.image

    if not source:
        print(json.dumps({
            "status": "failed",
            "message": "请指定图片路径或 URL"
        }))
        sys.exit(1)

    # === 执行 OCR ===
    result = None

    # 优先本地 Tesseract（除非 --force-napcat 或远程 URL）
    if not args.force_napcat and TESSERACT_AVAILABLE:
        if source.startswith("http://") or source.startswith("https://"):
            # 下载远程图片到本地
            try:
                import requests as req
                r = req.get(source, timeout=15)
                ext = source.split(".")[-1].split("?")[0][:4] if "." in source else "png"
                tmp_path = f"/tmp/_ocr_dl_{int(time.time())}.{ext}"
                with open(tmp_path, "wb") as f:
                    f.write(r.content)
                result = ocr_local(tmp_path, args.lang)
                os.remove(tmp_path)
            except Exception as e:
                result = {
                    "status": "failed",
                    "source": "tesseract",
                    "message": f"下载远程图片失败: {e}"
                }
        else:
            result = ocr_local(source, args.lang)

    # 如果本地失败或不可用，走 NapCat 回退
    if result is None or result.get("status") != "ok":
        napcat_result = ocr_napcat(source)
        if napcat_result and napcat_result.get("status") == "ok":
            result = napcat_result
        elif result and result.get("status") != "ok":
            result["napcat_fallback"] = napcat_result.get("message") if napcat_result else None

    if result is None:
        result = {"status": "failed", "message": "所有 OCR 方式均失败"}

    print(json.dumps(result, ensure_ascii=False, indent=2))
    if result.get("status") != "ok":
        sys.exit(1)


if __name__ == "__main__":
    main()