#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模板匹配：在截图中查找模板图片的位置
用法1: python image-match.py <screenshot_path> <template_path> [threshold]
用法2: python image-match.py --adb <adb_path> --device <device_id> --screenshot <out_path> --template <template_path> [--threshold 0.8]
      用法2 会在 Python 内执行 adb 截图，避免 Node 处理二进制数据导致的兼容性问题
输出: JSON 到 stdout
"""

import sys
import os
import json
import subprocess

try:
    import cv2
    import numpy as np
except ImportError as e:
    print(json.dumps({"success": False, "error": f"OpenCV 导入失败: {e}。请安装: pip install opencv-python numpy"}))
    sys.exit(1)

try:
    from PIL import Image as PILImage
    HAS_PIL = True
except ImportError:
    HAS_PIL = False

def run_adb_screencap(adb_path, device, output_path):
    """在 Python 内执行 adb 截图，直接处理二进制流"""
    # Windows 下子进程需要可执行路径，正斜杠也可用
    args = [adb_path.replace('/', os.sep), '-s', device, 'exec-out', 'screencap', '-p']
    try:
        result = subprocess.run(args, capture_output=True, timeout=15)
        if result.returncode != 0:
            return False, (result.stderr or result.stdout or b'').decode('utf-8', errors='replace')
        data = result.stdout
        if not data or len(data) < 100:
            return False, "截图数据为空"
        # 注意：不要对 PNG 数据做 \r\n 替换，会破坏 IDAT 压缩块导致无法解析
        out_dir = os.path.dirname(output_path)
        if out_dir:
            os.makedirs(out_dir, exist_ok=True)
        with open(output_path, 'wb') as f:
            f.write(data)
        return True, output_path
    except subprocess.TimeoutExpired:
        return False, "截图超时"
    except Exception as e:
        return False, str(e)

def load_image(path):
    """从文件路径加载图片，兼容 OpenCV 无法直接读取的 PNG（如部分 Android 截图）"""
    if not os.path.exists(path):
        return None
    with open(path, 'rb') as f:
        data = np.frombuffer(f.read(), dtype=np.uint8)
    img = cv2.imdecode(data, cv2.IMREAD_COLOR)
    if img is not None:
        return img
    img = cv2.imread(path)
    if img is not None:
        return img
    if HAS_PIL:
        try:
            pil_img = PILImage.open(path).convert('RGB')
            img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
            return img
        except Exception:
            pass
    return None

def main():
    screenshot_path = None
    template_path = None
    threshold = 0.8
    adb_path = None
    device = None

    if len(sys.argv) >= 2 and sys.argv[1] == '--adb':
        # 用法2：--adb --device --screenshot --template
        i = 1
        while i < len(sys.argv):
            if sys.argv[i] == '--adb' and i + 1 < len(sys.argv):
                adb_path = sys.argv[i + 1]
                i += 2
            elif sys.argv[i] == '--device' and i + 1 < len(sys.argv):
                device = sys.argv[i + 1]
                i += 2
            elif sys.argv[i] == '--screenshot' and i + 1 < len(sys.argv):
                screenshot_path = sys.argv[i + 1]
                i += 2
            elif sys.argv[i] == '--template' and i + 1 < len(sys.argv):
                template_path = sys.argv[i + 1]
                i += 2
            elif sys.argv[i] == '--threshold' and i + 1 < len(sys.argv):
                threshold = float(sys.argv[i + 1])
                i += 2
            else:
                i += 1
        if adb_path and device and screenshot_path and template_path:
            ok, msg = run_adb_screencap(adb_path, device, screenshot_path)
            if not ok:
                print(json.dumps({"success": False, "error": f"截图失败: {msg}"}))
                sys.exit(1)
        else:
            print(json.dumps({"success": False, "error": "缺少 --adb/--device/--screenshot/--template 参数"}))
            sys.exit(1)
    else:
        # 用法1：位置参数
        if len(sys.argv) < 3:
            print(json.dumps({"success": False, "error": "用法: image-match.py <screenshot_path> <template_path> [threshold]"}))
            sys.exit(1)
        screenshot_path = sys.argv[1]
        template_path = sys.argv[2]
        threshold = float(sys.argv[3]) if len(sys.argv) > 3 else 0.8

    if not os.path.exists(screenshot_path):
        print(json.dumps({"success": False, "error": f"截图文件不存在: {screenshot_path}"}))
        sys.exit(1)

    if not os.path.exists(template_path):
        print(json.dumps({"success": False, "error": f"模板文件不存在: {template_path}"}))
        sys.exit(1)

    screenshot = load_image(screenshot_path)
    template = load_image(template_path)

    if screenshot is None:
        print(json.dumps({"success": False, "error": "无法读取截图（文件损坏或格式不支持）"}))
        sys.exit(1)

    if template is None:
        print(json.dumps({"success": False, "error": f"无法读取模板: {template_path}"}))
        sys.exit(1)

    t_h, t_w = template.shape[:2]
    if t_h > screenshot.shape[0] or t_w > screenshot.shape[1]:
        print(json.dumps({"success": False, "error": "模板尺寸大于截图"}))
        sys.exit(1)

    # 使用 TM_CCOEFF_NORMED 进行模板匹配
    result = cv2.matchTemplate(screenshot, template, cv2.TM_CCOEFF_NORMED)
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)

    if max_val < threshold:
        print(json.dumps({"success": False, "error": f"未找到匹配 (相似度 {max_val:.3f} < {threshold})"}))
        sys.exit(1)

    x, y = int(max_loc[0]), int(max_loc[1])
    center_x = x + t_w // 2
    center_y = y + t_h // 2

    output = {
        "success": True,
        "x": x,
        "y": y,
        "width": t_w,
        "height": t_h,
        "center_x": center_x,
        "center_y": center_y
    }
    print(json.dumps(output))
    sys.exit(0)

if __name__ == "__main__":
    main()