metadata_fetcher.py - dev - vaskaz/flexin_flow

metadata_fetcher.py
315 строк10 KB
vaskaz
newborn
10 июн 2026, 00:16
10 июн 2026, 00:16dc9da65
"""Metadata fetcher — enriches ROMs with descriptions and cover art via Wikipedia API."""

import os
import re
import json
import time
import urllib.request
import urllib.parse
import urllib.error
import hashlib
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

from database import get_all_roms, get_rom_by_id, search_roms_by_filters, get_conn

COVERS_DIR = Path(__file__).parent / "covers_cache"
COVERS_DIR.mkdir(exist_ok=True)

WIKI_BASE = "https://en.wikipedia.org/api/rest_v1"
WIKI_COMMONS = "https://commons.wikimedia.org/w/api.php"
USER_AGENT = "RetroLibrary/1.0 (local metadata fetcher)"

# Platform → Wikipedia category mapping for better search
PLATFORM_WIKI_TERMS = {
    "nes": "Nintendo Entertainment System",
    "snes": "Super Nintendo Entertainment System",
    "n64": "Nintendo 64",
    "gb": "Game Boy",
    "gbc": "Game Boy Color",
    "gba": "Game Boy Advance",
    "nds": "Nintendo DS",
    "3ds": "Nintendo 3DS",
    "psx": "PlayStation",
    "ps2": "PlayStation 2",
    "psp": "PlayStation Portable",
    "genesis": "Sega Genesis",
    "saturn": "Sega Saturn",
    "dreamcast": "Dreamcast",
    "sms": "Sega Master System",
    "gg": "Game Gear",
    "pce": "TurboGrafx-16",
    "neogeo": "Neo Geo",
    "atari2600": "Atari 2600",
    "atari5200": "Atari 5200",
    "atari7800": "Atari 7800",
    "lynx": "Atari Lynx",
    "jaguar": "Atari Jaguar",
}

# ── Wikipedia API helpers ──────────────────────────────────────────────

def _wiki_fetch(path: str) -> dict | None:
    """Fetch from Wikipedia REST API with retries."""
    url = f"{WIKI_BASE}{path}"
    req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
    for attempt in range(3):
        try:
            with urllib.request.urlopen(req, timeout=10) as resp:
                return json.loads(resp.read())
        except urllib.error.HTTPError as e:
            if e.code == 404:
                return None
            if e.code == 429:  # rate limit
                time.sleep(2 ** attempt)
                continue
            return None
        except (urllib.error.URLError, OSError, json.JSONDecodeError):
            time.sleep(1)
            continue
    return None


def _wiki_search(query: str) -> list[dict]:
    """Search Wikipedia pages by title."""
    params = urllib.parse.urlencode({
        "action": "query",
        "list": "search",
        "srsearch": query,
        "srlimit": 5,
        "format": "json",
    })
    url = f"{WIKI_BASE.replace('/api/rest_v1', '/w/api.php')}?{params}"
    req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = json.loads(resp.read())
            return data.get("query", {}).get("search", [])
    except Exception:
        return []


def _download_image(url: str, filename: str) -> str | None:
    """Download an image to covers_cache. Returns filename or None."""
    local_name = hashlib.md5(filename.encode()).hexdigest() + ".jpg"
    local_path = COVERS_DIR / local_name
    if local_path.exists():
        return local_name  # already cached

    req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = resp.read()
            if len(data) < 500:  # too small to be a real image
                return None
            local_path.write_bytes(data)
            return local_name
    except Exception:
        return None


def _wiki_image_url(page_title: str) -> str | None:
    """Get the best thumbnail URL for a Wikipedia page."""
    data = _wiki_fetch(f"/page/summary/{urllib.parse.quote(page_title)}")
    if data and data.get("thumbnail"):
        return data["thumbnail"]["source"]
    return None


def _build_search_terms(title: str, platform_id: str) -> list[str]:
    """Build search queries for Wikipedia, from specific to broad."""
    platform_name = PLATFORM_WIKI_TERMS.get(platform_id, "")
    terms = [title]

    if platform_name:
        terms.append(f"{title} ({platform_name})")
        terms.append(f"{title} ({platform_name} video game)")
        terms.append(f"{title} video game")

    # Handle "The" prefix
    if title.startswith("The "):
        base = title[4:]
        terms.append(base)
        if platform_name:
            terms.append(f"{base} ({platform_name})")

    return terms


def _wiki_search_best(title: str, platform_id: str) -> dict:
    """Search Wikipedia for the best matching game page."""
    search_terms = _build_search_terms(title, platform_id)

    for term in search_terms:
        # Try exact page first
        data = _wiki_fetch(f"/page/summary/{urllib.parse.quote(term)}")
        if data and data.get("type") not in ("disambiguation",):
            return data

        # Try search
        results = _wiki_search(term)
        for r in results:
            page_title = r.get("title", "")
            data = _wiki_fetch(f"/page/summary/{urllib.parse.quote(page_title)}")
            if data and data.get("type") not in ("disambiguation",):
                return data

    return {}


# ── Core enrichment ───────────────────────────────────────────────────

def enrich_game(rom_id: int) -> dict:
    """Fetch metadata for a single ROM from Wikipedia and store it."""
    rom = get_rom_by_id(rom_id)
    if not rom:
        return {"error": "ROM not found"}

    title = rom.get("title", "")
    platform = rom.get("platform", "")
    if not title:
        return {"error": "No title"}

    result = {"rom_id": rom_id, "title": title}

    # Skip if already enriched
    if rom.get("description") and rom.get("cover_url"):
        result["status"] = "already_enriched"
        result["cover_url"] = rom["cover_url"]
        result["description"] = rom["description"]
        return result

    # Search Wikipedia
    wiki_data = _wiki_search_best(title, platform)
    if not wiki_data:
        result["status"] = "not_found"
        return result

    wiki_title = wiki_data.get("title", title)
    description = wiki_data.get("extract", "")

    # Clean up description — remove excessive newlines
    if description:
        description = re.sub(r"\n+", "\n", description).strip()
        # Truncate to reasonable length
        if len(description) > 2000:
            description = description[:1997] + "..."

    # Download cover image
    cover_filename = ""
    thumbnail = wiki_data.get("thumbnail", {}).get("source", "")
    if thumbnail:
        # Wikipedia thumbnails look like: https://upload.wikimedia.org/.../thumb/foo.jpg/300px-foo.jpg
        # Get the full-resolution version by removing /thumb/ and the size suffix
        full_url = re.sub(r"/thumb(/.*?)/\d+px-", r"\1/", thumbnail)
        # Remove size suffix for images not in thumb format
        if not full_url.startswith("http"):
            full_url = thumbnail  # fallback to thumbnail
        cover_filename = _download_image(full_url, title)
        if not cover_filename:
            # Try thumbnail version
            cover_filename = _download_image(thumbnail, title + "_thumb")

    # Save to DB
    conn = get_conn()
    try:
        conn.execute(
            "UPDATE roms SET description = ?, cover_url = ? WHERE id = ?",
            (description, cover_filename or "", rom_id),
        )
        conn.commit()
    finally:
        conn.close()

    result["status"] = "enriched"
    result["cover_url"] = cover_filename or ""
    result["description"] = description
    return result


def enrich_all_games(limit: int = 0, progress_callback=None) -> dict:
    """Enrich all games that don't have metadata yet."""
    all_roms = get_all_roms()
    to_enrich = [r for r in all_roms if not r.get("description") or not r.get("cover_url")]

    if limit:
        to_enrich = to_enrich[:limit]

    stats = {"total": len(all_roms), "enriched": 0, "not_found": 0, "errors": 0}

    if not to_enrich:
        stats["message"] = "All games already enriched"
        return stats

    def enrich_one(rom: dict) -> str:
        try:
            result = enrich_game(rom["id"])
            if result.get("status") == "enriched":
                return "ok"
            elif result.get("status") == "not_found":
                return "not_found"
            else:
                return "error"
        except Exception:
            return "error"

    with ThreadPoolExecutor(max_workers=4) as pool:
        futures = {pool.submit(enrich_one, r): r for r in to_enrich}
        done = 0
        total = len(to_enrich)
        for future in as_completed(futures):
            status = future.result()
            done += 1
            if status == "ok":
                stats["enriched"] += 1
            elif status == "not_found":
                stats["not_found"] += 1
            else:
                stats["errors"] += 1
            if progress_callback and done % max(1, total // 10) == 0:
                progress_callback(done, total)

    stats["total_attempted"] = len(to_enrich)
    return stats


# ── Cover download for existing entries ────────────────────────────────

def download_cover(rom_id: int) -> str | None:
    """Download cover for a ROM that already has metadata but missing cover."""
    rom = get_rom_by_id(rom_id)
    if not rom:
        return None
    if rom.get("cover_url"):
        return rom["cover_url"]  # already have one

    title = rom.get("title", "")
    platform = rom.get("platform", "")
    if not title:
        return None

    wiki_data = _wiki_search_best(title, platform)
    if not wiki_data:
        return None

    thumbnail = wiki_data.get("thumbnail", {}).get("source", "")
    if not thumbnail:
        return None

    full_url = re.sub(r"/thumb(/.*?)/\d+px-", r"\1/", thumbnail)
    if not full_url.startswith("http"):
        full_url = thumbnail

    cover_filename = _download_image(full_url, title)
    if cover_filename:
        conn = get_conn()
        try:
            conn.execute("UPDATE roms SET cover_url = ? WHERE id = ?", (cover_filename, rom_id))
            conn.commit()
        finally:
            conn.close()
        return cover_filename
    return None


__all__ = ["enrich_game", "enrich_all_games", "download_cover", "COVERS_DIR"]