project/tbsg/ai_helper_vision.py

"""
AI 视觉助手 — 方案A：截图直接送视觉大模型识别
=================================================
跳过本地 OCR，把整张截图直接给视觉模型，让它同时完成"识字 + 结构化 + 定位坐标"。
返回格式与 ai_helper1.AIParser.parse_shops 完全一致，可直接替换。

用法（在 main1.py 里替换 _get_named_shops 的 AI 调用）：
    from ai_helper_vision import VisionParser
    parser = VisionParser()
    shops = parser.parse_shops(shot_path, screen_size=(w, h), keyword=keyword)

独立测试：
    python ai_helper_vision.py  <截图路径>
"""
import base64
import json
import os
import re
import sys
import time
from pathlib import Path
from typing import Optional

import requests

# ── 配置 ────────────────────────────────────────────────
# 视觉引擎开关：
#   True  = 本地 HunyuanOCR（llama-server，http://127.0.0.1:8080，OpenAI兼容，免费/离线）
#   False = DeepSeek 官网视觉 API（云端，收费，需 API Key）
USE_LOCAL_HUNYUAN = True

# 本地 HunyuanOCR（需先启动 llama-server：双击 D:\HunyuanOCR\start_gpu.bat）
HY_API_URL = "http://127.0.0.1:8080/v1/chat/completions"
HY_MODEL = "HYVL"

# DeepSeek 官网 API（OpenAI 兼容格式，非 Anthropic 格式）
DS_API_URL = "https://api.deepseek.com/v1/chat/completions"
# 可设环境变量 VISION_MODEL 覆盖 DeepSeek 模型名
DS_VISION_MODEL = os.environ.get("VISION_MODEL", "deepseek-v4-flash-vision-exp")

# 当前生效的后端（由开关决定）
API_URL = HY_API_URL if USE_LOCAL_HUNYUAN else DS_API_URL
VISION_MODEL = HY_MODEL if USE_LOCAL_HUNYUAN else DS_VISION_MODEL

# 图片过大时缩放到这个宽度（视觉模型对超长图支持差，2712px 长图容易漏识别底部卡片）
MAX_IMAGE_WIDTH = 1280

PROJECT_ROOT = Path(__file__).parent


class VisionParser:
    """截图直接送视觉模型，返回结构化店铺列表"""

    # API Key 直接写死在文件里
    _API_KEY = "sk-08f45783a60f4b3183deef41d321474d"

    def __init__(self, api_key: Optional[str] = None):
        if USE_LOCAL_HUNYUAN:
            # 本地 llama-server 不校验 Key，传占位即可
            self.api_key = api_key or "not-needed"
        else:
            self.api_key = api_key or self._API_KEY
            if not self.api_key:
                raise RuntimeError("未设置 API Key")

    # ── 核心：截图 → 视觉模型 → 店铺列表 ───────────────

    def parse_shops(
        self,
        image_path: str,
        screen_size: tuple = (1220, 2712),
        keyword: str = "",
        crop_x: int = 0,
    ) -> list:
        """
        截图送视觉模型，返回结构化店铺列表。

        Args:
            image_path: 截图文件路径
            screen_size: (宽, 高)，截图的原始尺寸，用于坐标校验
            keyword: 当前搜索的商品关键词
            crop_x: 左列商品图片的分界x，>0时只裁右列（文字区）给AI，去掉商品图片

        Returns:
            [[店铺名, 商品标题, 价格, [点击x, 点击y], 店铺名y, 月售量], ...]
            点击坐标由「标题y 和 店铺名y 的中点」算出，不直接让AI返回坐标
        """
        screen_w, screen_h = screen_size

        # 1. 读取图片，裁掉左列商品图片（只留右列文字区），超宽时缩放
        img_data, media_type, scale, crop_w = self._encode_image(image_path, screen_w, crop_x)

        # 2. 构造 prompt
        system_prompt = self._system_prompt(keyword)
        crop_desc = f"已裁掉左侧商品图片，只保留右侧文字区（宽{crop_w}px）" if crop_x else "完整截图"
        user_prompt = (
            f"这是一张饿了么闪购App搜索结果截图（{crop_desc}，高度{screen_h}px）。"
            f"请找出屏幕上所有商品卡片，按格式返回。\n"
            f"当前搜索的商品是「{keyword}」。"
        )

        # 3. 调 API（带图片）
        print(f"[ai-vision] 识别图片: {os.path.basename(image_path)} ({screen_w}x{screen_h}, crop_x={crop_x})...")
        response = self._call(system_prompt, user_prompt, img_data, media_type)

        # 4. 解析返回，点击坐标由标题y和店铺名y的中点算出
        result = self._parse_response(response, screen_w, screen_h, scale, crop_x, crop_w)
        print(f"[ai-vision] 解析到 {len(result)} 个店铺")
        return result

    # ── 图片编码 + 缩放 ──────────────────────────────

    def _encode_image(self, image_path: str, screen_w: int, crop_x: int = 0):
        """读取图片，裁掉左列商品图片（crop_x>0时只留右列），超宽时缩放。
        返回 (base64_data, media_type, scale, crop_width)"""
        global MAX_IMAGE_WIDTH
        try:
            from PIL import Image
            import io

            img = Image.open(image_path)
            w, h = img.size

            # 裁掉左列商品图片，只留右列文字区
            if crop_x and crop_x > 0 and crop_x < w:
                img = img.crop((crop_x, 0, w, h))
                w = img.size[0]
                print(f"[ai-vision] 裁掉左列商品图片: 只留右列文字区 (宽{w}px)")

            scale = 1.0
            if w > MAX_IMAGE_WIDTH:
                scale = MAX_IMAGE_WIDTH / w
                new_w = MAX_IMAGE_WIDTH
                new_h = int(h * scale)
                img = img.resize((new_w, new_h), Image.LANCZOS)
                print(f"[ai-vision] 图片缩放: {w}x{h} → {new_w}x{new_h} (scale={scale:.3f})")

            # 转 PNG（统一格式，避免 JPEG 压缩干扰文字识别）
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            data = base64.b64encode(buf.getvalue()).decode("utf-8")
            return data, "image/png", scale, w
        except ImportError:
            # 没有 PIL，直接读原始文件（无法裁剪）
            with open(image_path, "rb") as f:
                data = base64.b64encode(f.read()).decode("utf-8")
            ext = os.path.splitext(image_path)[1].lower()
            media = {"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg"}.get(ext.lstrip("."), "image/png")
            return data, media, 1.0, screen_w

    # ── system prompt ────────────────────────────────

    @staticmethod
    def _system_prompt(keyword: str) -> str:
        return f"""你是手机截图分析助手。收到一张饿了么闪购App的搜索结果截图（已裁掉左侧商品图片，只剩右侧文字区），需要找出每个商品卡片的信息。

【重要：当前搜索的商品】
用户正在搜索「{keyword}」。请严格遵守：
- 商品标题优先取与「{keyword}」匹配的文字行（含品牌名和规格）
- 功能描述（如"养阴敛肺""解热镇痛""清热解毒""止咳化痰"等）字数少、无品牌名、在多个卡片中重复 → 绝对不能作为商品标题
- 找不到匹配时，商品标题取卡片中最像完整商品名的文字（含品牌+药名+规格结构）
- 如果整个页面与「{keyword}」无关（如显示快餐/外卖商家），返回空数组 []

【卡片布局 — 从上到下】
1. 商品标题（最上方，通常以品牌标签开头，如[理洫王]、[999]、[以岭]、[康恩贝]，后面跟药名+规格）
2. 功能描述（短句，如"养阴敛肺"，可忽略）
3. 价格（￥或¥开头）
4. 月售/已售（如"月售6000""已售92"）
5. 店铺名（以"店/药房/大药房/旗舰店"结尾）

【字段规则】
1. 店铺名 = 卡片最下方像店名的文字。含"（"只取"（"之前部分。
2. 商品标题 = 卡片最上方的完整商品名称，**必须包含开头的品牌标签**（如"[理洫王]血塞通软胶囊0.33g*24粒/盒"，不能把[理洫王]丢掉）。
   品牌标签和后面的药名+规格要完整拼接在一起，不要截断、不要去掉品牌。
3. 价格 = 卡片中带￥或¥的文字，只取数值+货币符号（如"￥12.8"）。
4. 点击坐标[x,y]：
   - x = 卡片文字区的水平中间位置（图片宽度的中间左右）
   - y = **该卡片「商品标题文字」和「店铺名文字」的垂直中点**（标题在卡片上方，店铺名在卡片下方，取两者y坐标的中间值）
   - 这个坐标用于点击进入该商品，请尽量准确
5. 店铺名y = 店铺名文字所在位置的y坐标。
6. 月售量 = 含"月售"或"已售"文字里的数字部分。如"月售6000"→"6000"、"已售92"→"92"。找不到留空字符串""。

【约束】
- 坐标基于你看到的图片实际像素尺寸（图片宽度即文字区宽度）
- 店铺名和商品标题不能是同一个文字（如果是，说明没有独立标题，跳过该卡片）
- 有多少个卡片就输出多少个，不要遗漏，不要多加
- 输出纯JSON，不要任何多余文字
- **所有字符串值必须用双引号包裹**，如"灵露大药房"不能写成灵露大药房

返回格式（JSON对象数组，所有字符串加双引号）:
[{{"shop":"店铺名(去括号)", "title":"商品标题(完整含品牌)", "price":"价格", "click":{{"x":点击x, "y":点击y}}, "shop_y":店铺名y, "sales":"月售量"}}, ...]

示例:
[{{"shop":"灵露大药房", "title":"[理洫王]血塞通软胶囊0.33g*60粒/盒", "price":"￥42.6", "click":{{"x":368,"y":285}}, "shop_y":439, "sales":""}}]"""

    # ── API 调用（带图片）──────────────────────────────

    def _call(
        self,
        system_prompt: str,
        user_text: str,
        img_data: str,
        media_type: str,
        timeout: int = 300 if USE_LOCAL_HUNYUAN else 120,
        retries: int = 3,
    ) -> str:
        """调用视觉 API（本地 HunyuanOCR 或 DeepSeek，均为 OpenAI 兼容格式，content 含 image_url 块）"""
        # OpenAI 兼容格式：data URL 内嵌 base64 图片
        data_url = f"data:{media_type};base64,{img_data}"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "model": VISION_MODEL,
            "max_tokens": 8192,
            "messages": [
                {"role": "system", "content": system_prompt},
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": user_text},
                        {"type": "image_url", "image_url": {"url": data_url}},
                    ],
                },
            ],
        }

        last_err = None
        for attempt in range(retries):
            try:
                resp = requests.post(API_URL, json=payload, headers=headers, timeout=timeout)
                if resp.status_code != 200:
                    # 打印HTTP错误码和响应体，方便排查
                    print(f"    [ai-vision] HTTP {resp.status_code}: {resp.text[:300]}")
                    resp.raise_for_status()
                data = resp.json()
                # 检查API是否返回了错误（如模型不存在、限流等）
                if "error" in data:
                    err_msg = data["error"].get("message", str(data["error"])) if isinstance(data["error"], dict) else str(data["error"])
                    print(f"    [ai-vision] API返回错误: {err_msg[:300]}")
                    last_err = Exception(err_msg)
                    if attempt < retries - 1:
                        time.sleep(2)
                    continue
                # OpenAI 格式：choices[0].message.content
                content = data.get("choices", [{}])[0].get("message", {}).get("content", "") or ""
                if not content:
                    print(f"    [ai-vision] 返回内容为空，完整响应: {str(data)[:300]}")
                return content
            except Exception as e:
                last_err = e
                if attempt < retries - 1:
                    print(f"    [ai-vision] API调用失败(第{attempt+1}次)，重试... ({e})")
                    time.sleep(2)
        print(f"    [ai-vision] API调用{retries}次均失败，跳过: {last_err}")
        return ""

    # ── 解析返回 ────────────────────────────────────

    @staticmethod
    def _parse_response(response: str, screen_w: int, screen_h: int,
                         scale: float = 1.0, crop_x: int = 0, crop_w: int = 0) -> list:
        """解析 AI 返回的 JSON，坐标按缩放比例还原 + 加裁剪偏移"""
        if not response.strip():
            print("[ai-vision] 返回为空")
            return []

        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()

        try:
            data = json.loads(cleaned)
        except json.JSONDecodeError:
            match = re.search(r'\[[\s\S]*\]', cleaned)
            if match:
                try:
                    data = json.loads(match.group(0))
                except json.JSONDecodeError:
                    # JSON解析失败（AI没给字符串加引号），直接按字符串分割提取
                    print(f"[ai-vision] JSON解析失败，改用字符串分割提取...")
                    return VisionParser._loose_parse(cleaned, screen_w, screen_h, scale, crop_x)
            else:
                print(f"[ai-vision] 未找到 JSON 数组")
                return []

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

        # 缩放还原：模型看到的是缩放后的图，坐标要除以scale还原
        # x 还要加 crop_x 偏移（因为图被裁掉了左列，AI的x是相对裁剪后的）
        def _restore(v):
            return int(round(v / scale)) if scale and scale != 1.0 else int(v)

        validated = []
        for item in data:
            # 兼容 dict 格式
            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("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")
                sales = str(item.get("sales") or item.get("month_sales") or "")
                item = [shop, title, price, coords, shop_name_y, sales]
            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:
                # x: 先还原缩放，再加裁剪偏移 → 绝对屏幕坐标
                cx = _restore(coords[0]) + crop_x
                # y: 只还原缩放（裁剪是水平的，y不变）
                cy = _restore(coords[1])
                # 坐标校验（还原后在原始屏幕尺寸范围内）
                if 0 <= cx <= screen_w and 0 <= cy <= screen_h:
                    shop_name_y = _restore(item[4]) if len(item) >= 5 and isinstance(item[4], (int, float)) else min(cy + 200, screen_h)
                    sales = str(item[5]) if len(item) >= 6 and item[5] not in (None, "") else ""
                    validated.append([shop, title, price, [cx, cy], shop_name_y, sales])
        return validated

    @staticmethod
    def _loose_parse(text: str, screen_w: int, screen_h: int,
                      scale: float = 1.0, crop_x: int = 0) -> list:
        """JSON解析失败时，用字符串分割提取店铺数据。
        格式固定: [[店铺名, 标题, 价格, [x, y], shop_y, 月售], ...]
        不用正则，纯字符串操作。"""
        text = text.strip()
        # 去掉最外层方括号
        if text.startswith('['):
            text = text[1:]
        if text.endswith(']'):
            text = text[:-1]

        # 按 ], [ 分割成每个店铺条目
        entries = text.split('], [')

        def _restore(v):
            return int(round(v / scale)) if scale and scale != 1.0 else int(v)

        shops = []
        for entry in entries:
            entry = entry.strip().lstrip('[').rstrip(']')
            # 找价格标记(￥或¥)，用它定位标题和价格的分界
            price_idx = -1
            for marker in ('￥', '¥'):
                idx = entry.find(marker)
                if idx >= 0:
                    price_idx = idx
                    break
            if price_idx < 0:
                continue

            # 价格之前: "店铺名, 标题"
            before = entry[:price_idx].strip().rstrip(',').strip()
            comma_idx = before.find(',')
            if comma_idx < 0:
                continue
            shop = before[:comma_idx].strip().strip('"')
            title = before[comma_idx + 1:].strip().strip('"')

            # 价格之后: "￥42.6, [368, 285], 439, """
            rest = entry[price_idx:]
            parts = rest.split(',')
            # parts: ['￥42.6', ' [368', ' 285]', ' 439', ' ""']
            if len(parts) < 5:
                continue
            try:
                price = parts[0].strip()
                x = _restore(int(parts[1].strip().lstrip('[')))
                y = _restore(int(parts[2].strip().rstrip(']')))
                shop_y = _restore(int(parts[3].strip()))
                sales = parts[4].strip().strip('"') if len(parts) > 4 else ""
            except (ValueError, IndexError):
                continue
            x += crop_x
            if 0 <= x <= screen_w and 0 <= y <= screen_h:
                shops.append([shop, title, price, [x, y], shop_y, sales])

        print(f"[ai-vision] 字符串分割提取到 {len(shops)} 个店铺")
        return shops


# ── 独立测试 ───────────────────────────────────────────
if __name__ == "__main__":
    # 用法: python ai_helper_vision.py <截图路径> [关键词]
    if len(sys.argv) < 2:
        print("用法: python ai_helper_vision.py <截图路径> [关键词]")
        print("示例: python ai_helper_vision.py D:\\drug\\sg\\screenshots\\step3_b0.png 强力枇杷露")
        sys.exit(1)

    img_path = sys.argv[1]
    kw = sys.argv[2] if len(sys.argv) > 2 else "D:\drug\sg\screenshots\SK4T6XZH4PEUOZ99\misc\check_pos.png"

    if not os.path.exists(img_path):
        print(f"文件不存在: {img_path}")
        sys.exit(1)

    # 获取图片原始尺寸
    try:
        from PIL import Image
        w, h = Image.open(img_path).size
    except ImportError:
        w, h = 1220, 2712

    parser = VisionParser()
    shops = parser.parse_shops(img_path, screen_size=(w, h), keyword=kw)

    print("\n" + "=" * 60)
    print(f"  识别到 {len(shops)} 个店铺")
    print("=" * 60)
    for i, s in enumerate(shops, 1):
        print(f"  [{i}] 店铺: {s[0]}")
        print(f"      商品: {s[1]}")
        print(f"      价格: {s[2]}")
        print(f"      点击: ({s[3][0]}, {s[3][1]})")
        print(f"      月售: {s[5]}")
        print()