#!/usr/bin/env python3
"""
图片搜索脚本（Bing Playwright + 智能筛选）
用 Playwright 渲染 Bing 图片搜索，获取原图URL，多维度质量筛选

筛选维度：
  1. 分辨率（≥400px）
  2. 宽高比（≤3:1）
  3. 文件大小（≥10KB）
  4. 背景纯净度（边缘像素方差）
  5. 角标/Logo检测（四角深色像素块）
  6. 水印检测（半透明灰白像素）
  7. 文字叠加检测（高对比度边缘密度）
  8. 内容占比（非白色像素比例）

用法：
  python3 search_image.py "电力电缆" /tmp/images
  python3 search_image.py "防火门" /tmp/images --count 3
"""

import argparse
import hashlib
import json
import os
import re
import sys

import requests
from PIL import Image

HEADERS = {
    "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
}

# 筛选参数
MIN_WIDTH = 400
MIN_HEIGHT = 400
MIN_FILE_SIZE = 10 * 1024
MAX_ASPECT_RATIO = 3.0
MAX_CANDIDATES = 15

# 质量阈值
MAX_CORNER_NON_WHITE_RATIO = 0.60
MAX_WATERMARK_RATIO = 0.35
MAX_EDGE_DENSITY = 0.35

# 搜索关键词增强
SEARCH_TEMPLATES = [
    "{kw} 建材 工业品",
    "{kw} 实物图",
    "{kw}",
]


def search_bing_fast(keyword, count=10):
    """用 requests 快速搜索 Bing 图片（不依赖 Playwright）"""
    import urllib.parse

    all_urls = []
    for template in SEARCH_TEMPLATES:
        query = template.format(kw=keyword)
        encoded = urllib.parse.quote(query)
        url = f'https://cn.bing.com/images/search?q={encoded}&qft=+filterui:photo-photo&form=IRFLTR&first=1'

        try:
            resp = requests.get(url, headers=HEADERS, timeout=10)
            resp.raise_for_status()
            urls = re.findall(r'murl":"(https?://[^&]+)"', resp.text)
            all_urls.extend(urls)
            if len(all_urls) >= count:
                break
        except Exception:
            continue

    # 去重
    seen = set()
    result = []
    for u in all_urls:
        if u not in seen and len(result) < count:
            seen.add(u)
            result.append(u)

    return result


def download_image(url, output_dir, prefix="img"):
    """下载图片到指定目录"""
    try:
        resp = requests.get(url, headers=HEADERS, timeout=8, stream=True)
        resp.raise_for_status()

        content_type = resp.headers.get("Content-Type", "")
        if "gif" in content_type:
            return None

        content = resp.content
        if len(content) < MIN_FILE_SIZE:
            return None

        if not (content[:3] == b'\xff\xd8\xff' or
                content[:8] == b'\x89PNG\r\n\x1a\n' or
                content[:4] == b'RIFF'):
            return None

        md5 = hashlib.md5(content).hexdigest()[:12]
        ext = ".jpg"
        if "png" in content_type or content[:8] == b'\x89PNG\r\n\x1a\n':
            ext = ".png"
        elif "webp" in content_type or content[:4] == b'RIFF':
            ext = ".webp"

        filename = f"{prefix}_{md5}{ext}"
        filepath = os.path.join(output_dir, filename)

        with open(filepath, "wb") as f:
            f.write(content)
        return filepath
    except Exception:
        return None


# ============ 图片质量检测 ============


def check_corner_cleanliness(img, sample_ratio=0.08):
    w, h = img.size
    if img.mode != "RGB":
        img = img.convert("RGB")
    pixels = img.load()
    cx, cy = int(w * sample_ratio), int(h * sample_ratio)
    corners = [(0, 0, cx, cy), (w - cx, 0, w, cy), (0, h - cy, cx, h), (w - cx, h - cy, w, h)]
    total, non_white = 0, 0
    for x1, y1, x2, y2 in corners:
        for y in range(y1, y2):
            for x in range(x1, x2):
                r, g, b = pixels[x, y]
                total += 1
                if r + g + b < 720:
                    non_white += 1
    return non_white / total if total > 0 else 0


def check_watermark_ratio(img):
    w, h = img.size
    if img.mode != "RGB":
        img = img.convert("RGB")
    pixels = img.load()
    total, watermark = 0, 0
    step = 6
    for y in range(step, h - step, step):
        for x in range(step, w - step, step):
            r, g, b = pixels[x, y]
            total += 1
            brightness = (r + g + b) / 3
            if 200 <= brightness <= 254 and abs(r - g) < 15 and abs(g - b) < 15:
                has_nearby_content = False
                for dy in [-step, 0, step]:
                    for dx in [-step, 0, step]:
                        nr, ng, nb = pixels[x + dx, y + dy]
                        if nr + ng + nb < 600:
                            has_nearby_content = True
                            break
                    if has_nearby_content:
                        break
                if has_nearby_content:
                    watermark += 1
    return watermark / total if total > 0 else 0


def check_logo_corners(img, sample_ratio=0.12):
    w, h = img.size
    if img.mode != "RGB":
        img = img.convert("RGB")
    pixels = img.load()
    cx, cy = int(w * sample_ratio), int(h * sample_ratio)
    corners = [(0, 0, cx, cy), (w - cx, 0, w, cy), (0, h - cy, cx, h), (w - cx, h - cy, w, h)]
    max_ratio = 0
    for x1, y1, x2, y2 in corners:
        total, dark = 0, 0
        for y in range(y1, y2):
            for x in range(x1, x2):
                r, g, b = pixels[x, y]
                total += 1
                if r + g + b < 500:
                    dark += 1
        max_ratio = max(max_ratio, dark / total if total > 0 else 0)
    return max_ratio


def check_text_overlay(img):
    w, h = img.size
    if img.mode != "RGB":
        img = img.convert("RGB")
    gray = img.convert("L")
    pixels = gray.load()
    total, edges = 0, 0
    step = 3
    for y in range(1, h - 1, step):
        for x in range(1, w - 1, step):
            total += 1
            gx = abs(pixels[x + 1, y] - pixels[x - 1, y])
            gy = abs(pixels[x, y + 1] - pixels[x, y - 1])
            if gx + gy > 80:
                edges += 1
    return edges / total if total > 0 else 0


def check_background_purity(img):
    w, h = img.size
    if img.mode != "RGB":
        img = img.convert("RGB")
    pixels = img.load()
    edge_pixels = []
    step = 8
    for x in range(0, w, step):
        edge_pixels.append(pixels[x, 0])
        edge_pixels.append(pixels[x, h - 1])
    for y in range(0, h, step):
        edge_pixels.append(pixels[0, y])
        edge_pixels.append(pixels[w - 1, y])
    if len(edge_pixels) < 3:
        return 0.5
    avg_r = sum(p[0] for p in edge_pixels) / len(edge_pixels)
    avg_g = sum(p[1] for p in edge_pixels) / len(edge_pixels)
    avg_b = sum(p[2] for p in edge_pixels) / len(edge_pixels)
    variance = sum(
        ((p[0] - avg_r) ** 2 + (p[1] - avg_g) ** 2 + (p[2] - avg_b) ** 2)
        for p in edge_pixels
    ) / len(edge_pixels)
    if variance < 500:
        return 1.0
    elif variance < 2000:
        return 0.8
    elif variance < 5000:
        return 0.5
    else:
        return 0.2


# ============ 综合评分 ============


def score_image(filepath):
    try:
        img = Image.open(filepath)
        w, h = img.size

        if w < MIN_WIDTH or h < MIN_HEIGHT:
            return -1, "分辨率不足"

        ratio = max(w, h) / min(w, h) if min(w, h) > 0 else 999
        if ratio > MAX_ASPECT_RATIO:
            return -1, "宽高比异常"

        bg_purity = check_background_purity(img)
        if bg_purity < 0.3:
            return -1, f"背景过于复杂({bg_purity:.2f})"

        logo_score = check_logo_corners(img)
        if logo_score > MAX_CORNER_NON_WHITE_RATIO:
            return -1, f"有角标/Logo({logo_score:.2f})"

        wm_ratio = check_watermark_ratio(img)
        if wm_ratio > MAX_WATERMARK_RATIO:
            return -1, f"有水印({wm_ratio:.2f})"

        text_density = check_text_overlay(img)
        if text_density > MAX_EDGE_DENSITY:
            return -1, f"有文字叠加({text_density:.2f})"

        if img.mode == "RGB":
            pixels = list(img.getdata())
            non_white = sum(1 for r, g, b in pixels if r + g + b < 700)
            content_ratio = non_white / len(pixels) if pixels else 0
        else:
            content_ratio = 0.5

        if content_ratio < 0.05:
            return -1, "内容过少(几乎全白)"

        pixel_score = min(w, h) / 1000
        ratio_score = 1.0 / ratio
        size_score = min(os.path.getsize(filepath) / 500000, 1.0)

        total = (pixel_score * 2 + ratio_score * 1.5 + size_score * 1 +
                 bg_purity * 2 + content_ratio * 2 -
                 text_density * 3 - logo_score * 2 - wm_ratio * 2)

        return total, "OK"
    except Exception as e:
        return -1, str(e)


def select_best(filepaths):
    scored = []
    for fp in filepaths:
        s, reason = score_image(fp)
        scored.append((s, reason, fp))
    valid = [(s, r, fp) for s, r, fp in scored if s > 0]
    if not valid:
        return None, scored
    valid.sort(reverse=True, key=lambda x: x[0])
    return valid[0][2], scored


def select_top_n(filepaths, n):
    scored = []
    for fp in filepaths:
        s, reason = score_image(fp)
        scored.append((s, reason, fp))
    valid = [(s, r, fp) for s, r, fp in scored if s > 0]
    if not valid:
        return [], scored
    valid.sort(reverse=True, key=lambda x: x[0])
    return [fp for _, _, fp in valid[:n]], scored


def main():
    parser = argparse.ArgumentParser(description="图片搜索")
    parser.add_argument("keyword", help="搜索关键词")
    parser.add_argument("output_dir", help="输出目录")
    parser.add_argument("--count", type=int, default=3, help="返回前N张图片（按质量评分排序）")
    parser.add_argument("--exclude-hash", dest="exclude_hashes", action="append", default=[], help="排除的图片MD5前缀列表（可多次传）")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    # 用 requests 快速搜索（不依赖 Playwright，更可靠）
    urls = search_bing_fast(args.keyword, MAX_CANDIDATES)

    safe_keyword = re.sub(r"[^\w\u4e00-\u9fa5]", "_", args.keyword)[:20]
    exclude_set = set(h.lower() for h in args.exclude_hashes)

    candidates = []
    for i, url in enumerate(urls):
        filepath = download_image(url, args.output_dir, prefix=f"{safe_keyword}_{i}")
        if filepath:
            # 读取 hash 用于后续去重
            try:
                with open(filepath, "rb") as f:
                    content = f.read()
                file_hash = hashlib.md5(content).hexdigest()[:12]
            except Exception:
                file_hash = ""
            if file_hash in exclude_set:
                try:
                    os.remove(filepath)
                except OSError:
                    pass
                continue
            candidates.append((filepath, file_hash))

    results = []
    keep_set = set()
    if candidates:
        fps = [c[0] for c in candidates]
        top, _ = select_top_n(fps, args.count)
        for fp in top:
            results.append({"path": fp, "hash": next((h for p, h in candidates if p == fp), "")})
        keep_set = set(fp for fp in top)
        # 删除未入选的临时文件
        for fp, _ in candidates:
            if fp not in keep_set:
                try:
                    os.remove(fp)
                except OSError:
                    pass

    output = {
        "keyword": args.keyword,
        "searched_urls": urls,
        "candidates": len(candidates),
        "downloaded": [r["path"] for r in results],
        "hashes": [r["hash"] for r in results],
        "count": len(results),
    }
    print(json.dumps(output, ensure_ascii=False))


if __name__ == "__main__":
    main()