project/tbsg/ai_helper.py

"""
AI 助手 — 调用 DeepSeek 分析 OCR 文本结果
只负责：接收文本 → 调 AI → 返回结构化数据
不负责 OCR 识别
"""
import json
import os
import re
import time
from pathlib import Path
from typing import Optional

import requests


# ── 配置 ────────────────────────────────────────────────
DEEPSEEK_API_URL = "https://ai-api.dfwy.tech/v1/messages"
DEEPSEEK_MODEL = "deepseek-v4-flash-nothinking"

# 项目根目录（用于读取 config.json）
PROJECT_ROOT = Path(__file__).parent


class AIParser:
    """调用 DeepSeek 分析 OCR 文本，返回结构化店铺列表"""

    # 默认 API Key（写死，也可通过环境变量或 config.json 覆盖）
    _DEFAULT_KEY = ""

    def __init__(self, api_key: Optional[str] = None):
        """
        Args:
            api_key: DeepSeek API Key，不传则从环境变量或 config.json 读取
        """
        self.api_key = api_key or self._load_api_key()
        if not self.api_key:
            raise RuntimeError(
                "未设置 DeepSeek API Key。请:\n"
                "1. 设环境变量 DEEPSEEK_API_KEY\n"
                "2. 或在 config.json 中添加 deepseek_api_key 字段\n"
                "3. 或在初始化时传入 api_key"
            )

    @staticmethod
    def _load_api_key() -> str:
        """从环境变量、config.json 或默认值读取 API Key"""
        key = os.environ.get("DEEPSEEK_API_KEY", "")
        if key:
            return key
        cfg_path = PROJECT_ROOT / "config.json"
        if cfg_path.exists():
            try:
                with open(cfg_path) as f:
                    cfg = json.load(f)
                key = cfg.get("deepseek_api_key", "")
                if key:
                    return key
            except Exception:
                pass
        return AIParser._DEFAULT_KEY

    # ── 核心：分析 OCR 结果 → 店铺列表 ──────────────────

    def parse_shops(
        self,
        ocr_results: list,
        screen_size: tuple = (1220, 2712),
    ) -> list:
        """
        将 OCR 识别结果传给 DeepSeek，返回结构化店铺列表

        Args:
            ocr_results: OCR.recognize(detail="all") 返回值
                [{"text": str, "confidence": float, "bbox": [[x1,y1],...]}, ...]
            screen_size: (宽, 高)，默认 1220x2712

        Returns:
            [[店铺名, 商品标题, 价格, [点击x, 点击y], 配送距离y], ...]
            点击坐标是店铺卡片中心点
            配送距离y = 配送时间/距离文字的y坐标，用于滑动到其上方
        """
        # 1. 格式化 OCR 数据 — 按 y 坐标排序（从上到下）
        sorted_results = sorted(ocr_results, key=lambda r: r["bbox"][0][1])

        ocr_lines = []
        for r in sorted_results:
            box = r["bbox"]
            cy = (box[0][1] + box[2][1]) // 2
            cx = (box[0][0] + box[2][0]) // 2
            ocr_lines.append(
                f"[x={cx:4d}, y={cy:4d}, conf={r['confidence']:.2f}] {r['text']}"
            )
        ocr_text = "\n".join(ocr_lines)

        screen_w, screen_h = screen_size

        # 2. system prompt
        system_prompt = """你是OCR文本结构化助手。从OCR文字列表中找出每个"店铺卡片"的4个字段，必须严格使用OCR原文。

【OCR数据格式】
每行: [x=坐标x, y=坐标y, conf=置信度] OCR识别出的文字
y越大越靠下，文字按从上到下排列。

【判断店铺卡片方法】
- 同卡片的文字y坐标连续相近，卡片间有明显y间隔(>80px)
- 卡片从上到下：商品标题(多行) → 价格(￥/¥开头) → 店铺名(以店/超市/旗舰店等结尾)

【字段规则 — 严格使用OCR原文，不得修改】
1. 店铺名 = 卡片最下方像店名的文字。
   如果文字含"（"，只取"（"之前的部分（括号内是分店名，不重要）。
   店铺名通常以"店/超市/旗舰店"结尾。
2. 商品标题 = 卡片上方描述商品的文字。
   多行时要完整拼接在一起，不要截断。
3. 价格 = 卡片中间带￥或¥符号的文字。
   只取**数值+货币符号**部分（如"￥12.8"、"¥67"），不要含"预估价"等额外文字。
   如果找不到价格，留空字符串""。
4. 点击坐标[x,y] = OCR数据中该卡片区域中间位置的实际坐标，用整数。

【约束】
- 所有字段值必须严格等于或取自OCR原文，不得修改、拼凑、推断
- 店铺名如果找不到，留空字符串""
- 找不到价格的卡片也要列出，价格留空
- 有多少个卡片就输出多少个，不要遗漏
- 输出纯JSON数组，不要任何多余文字

返回格式:
[[店铺名(去括号), 商品标题(完整), 价格(纯数字), [点击x, 点击y], 店铺名y], ...]
店铺名y = 店铺名文字在OCR数据中的y坐标，必须从OCR数据中找"""

        # 3. user prompt
        user_prompt = (
            f"屏幕OCR识别结果（屏幕尺寸 {screen_w}x{screen_h}）：\n\n"
            f"{ocr_text}\n\n请按格式返回店铺列表。"
        )

        # 4. 调 AI
        print(f"[ai] 分析 {len(ocr_results)} 个文本块...")
        response = self._call(system_prompt, user_prompt)

        # 5. 解析返回
        return self._parse_response(response, screen_w, screen_h)

    # ── 页面类型检测 ──────────────────────────────────

    def check_page(self, ocr_results: list) -> dict:
        """
        分析进入商品页后的页面类型

        Returns:
            {"type": "normal"|"risk"|"qrcode"|"unknown",
             "shop": "店铺名(如有)",
             "detail": "额外描述"}
        """
        # 1. 格式化
        sorted_r = sorted(ocr_results, key=lambda r: r["bbox"][0][1])
        lines = []
        for r in sorted_r:
            box = r["bbox"]
            cy = (box[0][1] + box[2][1]) // 2
            cx = (box[0][0] + box[2][0]) // 2
            lines.append(f"[x={cx:4d}, y={cy:4d}] {r['text']}")
        ocr_text = "\n".join(lines)

        system_prompt = """你是手机屏幕分析助手。收到OCR文字列表，判断当前屏幕是什么类型的页面。

【页面类型】
1. normal = 正常商品/店铺详情页（可以看到商品信息、价格、评价、店铺名等）
2. risk = 风控/验证码/拦截页（包含"验证"、"拖动滑块"、"安全验证"、"禁止搜索"、"操作频繁"、"风险"、"请输入验证码"等）
3. qrcode = 二维码/扫码页（包含"二维码"、"扫码"、"请用微信"等）
4. unknown = 其他无法识别的页面

**优先检测 risk 类型**：只要有类似验证码、风控拦截的文字，就返回risk

【输出格式】
只返回一个JSON对象，不要任何多余文字：
{"type": "normal|risk|qrcode|unknown", "shop": "店铺名或空字符串", "detail": "简要说明"}
"""
        user_prompt = f"屏幕OCR文字：\n{ocr_text}\n\n请判断页面类型。"

        print(f"[ai] 检测页面类型...")
        response = self._call(system_prompt, user_prompt)

        import json
        cleaned = response.strip()
        if cleaned.startswith("```"):
            lines = cleaned.split("\n")
            if lines[0].startswith("```"):
                lines = lines[1:]
            if lines and lines[-1].strip() == "```":
                lines = lines[:-1]
            cleaned = "\n".join(lines).strip()

        try:
            data = json.loads(cleaned)
            return {
                "type": data.get("type", "unknown"),
                "shop": data.get("shop", ""),
                "detail": data.get("detail", ""),
            }
        except json.JSONDecodeError:
            if not response.strip():
                return {"type": "normal", "shop": "", "detail": "API调用失败，默认继续"}
            return {"type": "unknown", "shop": "", "detail": response[:100]}

    # ── API 调用 ────────────────────────────────────────

    def _call(self, system_prompt: str, user_content: str, timeout: int = 120, retries: int = 3) -> str:
        """调用 API（Anthropic 格式），失败自动重试"""
        headers = {
            "x-api-key": self.api_key,
            "anthropic-version": "2023-06-01",
            "Content-Type": "application/json",
        }
        payload = {
            "model": DEEPSEEK_MODEL,
            "max_tokens": 8192,
            "system": system_prompt,
            "messages": [
                {"role": "user", "content": user_content},
            ],
        }

        last_err = None
        for attempt in range(retries):
            try:
                resp = requests.post(DEEPSEEK_API_URL, json=payload, headers=headers, timeout=timeout)
                resp.raise_for_status()
                data = resp.json()
                # 从 content 中提取所有 text 类型块
                texts = [c["text"] for c in data.get("content", []) if c.get("type") == "text"]
                return "".join(texts)
            except Exception as e:
                last_err = e
                if attempt < retries - 1:
                    print(f"    [ai] API调用失败(第{attempt+1}次)，重试... ({e})")
                    time.sleep(2)
        print(f"    [ai] API调用{retries}次均失败，跳过: {last_err}")
        return ""

    # ── 解析 ────────────────────────────────────────────

    @staticmethod
    def _parse_response(response: str, screen_w: int, screen_h: int) -> list:
        """解析 AI 返回的 JSON"""
        cleaned = response.strip()

        # 去掉 markdown 代码块包裹
        if cleaned.startswith("```"):
            lines = cleaned.split("\n")
            if lines[0].startswith("```"):
                lines = lines[1:]
            if lines and lines[-1].strip() == "```":
                lines = lines[:-1]
            cleaned = "\n".join(lines).strip()

        # 解析 JSON
        try:
            data = json.loads(cleaned)
        except json.JSONDecodeError:
            print(f"[ai] WARNING: JSON 解析失败，尝试提取数组片段")
            match = re.search(r'\[[\s\S]*\]', cleaned)
            if match:
                try:
                    data = json.loads(match.group(0))
                except json.JSONDecodeError:
                    print(f"[ai] ERROR: JSON 解析失败")
                    return []
            else:
                print(f"[ai] ERROR: 未找到 JSON 数组")
                return []

        if not isinstance(data, list):
            print(f"[ai] ERROR: 返回不是数组")
            return []

        validated = []
        for item in data:
            # 兼容 dict 格式（模型偶发不按 list 模板输出时自动转换）
            if isinstance(item, dict):
                shop = str(item.get("name") or item.get("shop") or item.get("shop_name") or "")
                title = str(item.get("title") or item.get("description") or item.get("product_name") or "")
                price = str(item.get("price") or item.get("min_price") or "")
                coords = item.get("coords") or item.get("click_pos") or item.get("click")
                if not isinstance(coords, list) and item.get("x") is not None and item.get("y") is not None:
                    coords = [item.get("x"), item.get("y")]
                shop_name_y = item.get("shop_name_y") or item.get("y")
                item = [shop, title, price, coords, shop_name_y]
            if not isinstance(item, list) or len(item) < 4:
                continue
            shop = str(item[0]) if item[0] else ""
            title = str(item[1]) if item[1] else ""
            price = str(item[2]) if item[2] else ""
            coords = item[3]
            if isinstance(coords, list) and len(coords) == 2:
                cx, cy = int(coords[0]), int(coords[1])
                if 0 <= cx <= screen_w and 0 <= cy <= screen_h:
                    # 第5个字段=店铺名文字的y坐标
                    shop_name_y = int(item[4]) if len(item) >= 5 and isinstance(item[4], (int, float)) else min(cy + 200, screen_h)
                    validated.append([shop, title, price, [cx, cy], shop_name_y])

        print(f"[ai] 解析到 {len(validated)} 个店铺")
        return validated


# ── 独立测试 ────────────────────────────────────────────

if __name__ == "__main__":
    import sys
    from steps.ocr import OCR

    # 读取已有截图
    test_img = str(PROJECT_ROOT / "screenshots" / "step2_result.png")
    if not os.path.exists(test_img):
        print(f"测试图片不存在: {test_img}")
        print("请先跑一次 main.py 生成截图")
        sys.exit(1)

    # OCR 识别
    ocr = OCR()
    raw = ocr.recognize(test_img, detail="all")
    print(f"OCR 识别到 {len(raw)} 个文本块\n")

    # AI 分析
    parser = AIParser()
    shops = parser.parse_shops(raw)

    print("\n" + "=" * 60)
    print(f"  店铺列表 ({len(shops)} 个)")
    print("=" * 60)
    for s in shops:
        print(f"  店铺: {s[0]}")
        print(f"  商品: {s[1]}")
        print(f"  价格: {s[2]}")
        print(f"  点击: ({s[3][0]}, {s[3][1]})")
        print(f"  配送y: {s[4]}")
        print()