"""从 Wikimedia Commons 下载训练图片（自由版权，可离线使用）。

策略：优先用物种分类目录（图片内容精确），再用全文搜索补充场景/姿态/光线多样性。
输出: datasets/images/<class>/<index>.jpg
串行 + 失败重试 + pHash 视觉去重。
"""
import json
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from threading import Lock

import imagehash
import numpy as np
from PIL import Image

API = "https://commons.wikimedia.org/w/api.php"
OUT = os.path.join(os.path.dirname(__file__), "datasets", "images")

# 更多图片来源
ADDITIONAL_SOURCES = {
    "inaturalist": "https://api.inaturalist.org/v1/observations",
    "flickr": "https://api.flickr.com/services/rest/",
}

SOURCES = {
    "pheasant": {
        "categories": ["Phasianus colchicus"],
        "queries": [
            # 户外真实场景（重点！）
            "pheasant in field",
            "pheasant in grassland",
            "pheasant in meadow",
            "pheasant in farmland",
            "pheasant in countryside",
            "pheasant in wild",
            # 部分遮挡场景（重点！实际使用场景）
            "pheasant hiding in grass",
            "pheasant hiding in bushes",
            "pheasant partially hidden",
            "pheasant behind vegetation",
            "pheasant peeking through grass",
            "pheasant concealed in foliage",
            "pheasant camouflaged",
            "pheasant blending in",
            # 不同距离和角度
            "pheasant distant view",
            "pheasant far away",
            "pheasant small in frame",
            "pheasant side view",
            "pheasant back view",
            "pheasant from behind",
            # 不同姿态和行为
            "pheasant walking in grass",
            "pheasant foraging",
            "pheasant feeding",
            "pheasant running",
            "pheasant flying low",
            # 光线条件（户外真实光线）
            "pheasant natural light",
            "pheasant daylight",
            "pheasant shade",
            "pheasant shadow",
            "pheasant backlit",
            "pheasant overcast",
            # 季节环境
            "pheasant in autumn",
            "pheasant in winter",
            "pheasant in spring",
            "pheasant in summer",
            "pheasant in dry grass",
            "pheasant in green grass",
            "pheasant in snow",
            "pheasant in mud",
        ],
    },
}
PER_CLASS = 80
MIN_PIXEL = 320
RETRY = 2
WORKERS = 1
SLEEP_S = 2.0
RATE_LIMIT_WAIT_S = 60


def api_request(params: dict) -> dict:
    url = API + "?" + urllib.parse.urlencode(params)
    for attempt in range(RETRY + 1):
        try:
            req = urllib.request.Request(url, headers={"User-Agent": "observer-training/1.0"})
            with urllib.request.urlopen(req, timeout=30) as r:
                return json.loads(r.read().decode("utf-8"))
        except Exception:
            if attempt == RETRY:
                raise
            time.sleep(2 * (attempt + 1))


def collect_pages(pages: dict, results: list) -> None:
    for p in pages.values():
        info = (p.get("imageinfo") or [{}])[0]
        thumb = info.get("thumburl")
        if not thumb:
            continue
        w, h = info.get("width", 0), info.get("height", 0)
        if min(w, h) < MIN_PIXEL:
            continue
        results.append({"url": thumb, "w": w, "h": h, "title": p.get("title", "")})


def category_images(category: str, limit: int) -> list[dict]:
    results = []
    params = {
        "action": "query",
        "generator": "categorymembers",
        "gcmtitle": f"Category:{category}",
        "gcmtype": "file",
        "gcmlimit": "50",
        "prop": "imageinfo",
        "iiprop": "url|size",
        "iiurlwidth": "640",
        "format": "json",
    }
    while len(results) < limit:
        data = api_request(params)
        pages = (data.get("query", {}) or {}).get("pages", {})
        collect_pages(pages, results)
        cont = (data.get("continue") or {}).get("gcmcontinue")
        if not cont:
            break
        params["gcmcontinue"] = cont
    return results[:limit]


def search_images(query: str, limit: int) -> list[dict]:
    results = []
    params = {
        "action": "query",
        "generator": "search",
        "gsrsearch": f"filetype:bitmap {query}",
        "gsrnamespace": "6",
        "gsrlimit": str(limit),
        "prop": "imageinfo",
        "iiprop": "url|size",
        "iiurlwidth": "640",
        "format": "json",
    }
    data = api_request(params)
    pages = (data.get("query", {}) or {}).get("pages", {})
    collect_pages(pages, results)
    return results


def is_text_heavy(image_path: str) -> bool:
    """检测图片是否包含大量文字或为图表/标志"""
    try:
        img = Image.open(image_path).convert("L")  # 灰度
        arr = np.array(img, dtype=np.float32)

        # 1. 检测边缘密度（文字产生大量边缘）
        # 简单边缘检测：计算像素梯度
        dx = np.abs(np.diff(arr, axis=1))
        dy = np.abs(np.diff(arr, axis=0))
        edge_density = (np.mean(dx) + np.mean(dy)) / 2

        # 2. 检测颜色方差（图表通常颜色单一）
        # 转回RGB检查
        img_rgb = Image.open(image_path).convert("RGB")
        arr_rgb = np.array(img_rgb, dtype=np.float32)
        color_variance = np.std(arr_rgb)

        # 3. 检测对比度（文字通常有高对比度的边缘）
        contrast = np.std(arr)

        # 判断逻辑：
        # - 高边缘密度 + 低颜色方差 = 可能是图表/标志
        # - 高对比度 + 高边缘密度 = 可能是文字图片
        if edge_density > 30 and color_variance < 50:
            return True
        if edge_density > 40 and contrast > 80:
            return True

        return False
    except Exception:
        return False


def download(url: str, path: str) -> bool:
    for attempt in range(RETRY + 1):
        try:
            req = urllib.request.Request(url, headers={"User-Agent": "observer-training/1.0"})
            with urllib.request.urlopen(req, timeout=60) as r:
                data = r.read()
            if "image/" not in r.headers.get("Content-Type", ""):
                return False
            if len(data) < 10_000:
                return False
            with open(path, "wb") as f:
                f.write(data)
            return True
        except urllib.error.HTTPError as e:
            if e.code == 429:
                # 限流：冷却后重试
                print(f"  ⏳ 限流(429)，冷却 {RATE_LIMIT_WAIT_S}s", flush=True)
                time.sleep(RATE_LIMIT_WAIT_S)
                continue
            if attempt == RETRY:
                return False
            time.sleep(2 * (attempt + 1))
        except Exception:
            if attempt == RETRY:
                return False
            time.sleep(2 * (attempt + 1))


def main():
    for cls, src in SOURCES.items():
        cls_dir = os.path.join(OUT, cls)
        os.makedirs(cls_dir, exist_ok=True)
        existing = len([f for f in os.listdir(cls_dir) if f.endswith(".jpg")])
        if existing >= PER_CLASS:
            print(f"[skip] {cls}: 已有 {existing} 张", flush=True)
            continue

        # 重建去重池（续跑兼容）
        seen_hashes = []
        for f in os.listdir(cls_dir):
            if f.endswith(".jpg"):
                try:
                    seen_hashes.append(imagehash.phash(Image.open(os.path.join(cls_dir, f))))
                except Exception:
                    pass

        candidates = []
        seen_titles = set()
        for cat in src["categories"]:
            print(f"[category] {cls} <- {cat}", flush=True)
            try:
                for item in category_images(cat, 80):
                    if item["title"] in seen_titles:
                        continue
                    seen_titles.add(item["title"])
                    candidates.append(item)
            except Exception as e:
                print(f"  ! category 失败: {e}", flush=True)
        for query in src["queries"]:
            if len(candidates) >= PER_CLASS * 2:
                break
            print(f"[fetch] {cls} <- \"{query}\"", flush=True)
            try:
                for item in search_images(query, 80):
                    if item["title"] in seen_titles:
                        continue
                    seen_titles.add(item["title"])
                    candidates.append(item)
            except Exception as e:
                print(f"  ! 搜索失败: {e}", flush=True)
        print(f"[download] {cls}: 候选 {len(candidates)} 张", flush=True)

        saved = existing
        fail = 0
        dup = 0
        for item in candidates:
            if saved >= PER_CLASS:
                break
            saved += 1
            path = os.path.join(cls_dir, f"{saved:03d}.jpg")
            if not download(item["url"], path):
                fail += 1
                saved -= 1
                continue

            # 过滤文字/图表类图片
            if is_text_heavy(path):
                os.remove(path)
                saved -= 1
                continue

            try:
                h = imagehash.phash(Image.open(path))
                if any(h - other <= 8 for other in seen_hashes):
                    os.remove(path)
                    dup += 1
                    saved -= 1
                    continue
                seen_hashes.append(h)
            except Exception:
                pass
            if saved % 10 == 0:
                print(f"  + {cls}: {saved}", flush=True)
            time.sleep(SLEEP_S)
        print(f"[done] {cls}: {saved} 张（失败 {fail}，重复 {dup}）", flush=True)


if __name__ == "__main__":
    sys.exit(main())