dl-organizer/src/vlm/enrichment.py

"""Identity enrichment pipeline.

Adds bilingual titles and reputation signals with incremental SQLite caching.
"""

from __future__ import annotations

import hashlib
import json
from typing import Callable, Optional
from urllib.request import urlopen, Request

from vlm.cache import EnrichmentCache
from vlm.config import Config
from vlm.providers import ProviderResult, TMDBAuthError, TMDBProvider, TMDBProviderError
from vlm.parser import normalize_title

RefreshMode = str
ProgressCallback = Callable[[int, int, dict[str, int]], None]


def enrich_identities_data(
    identities_data: dict,
    config: Config,
    *,
    refresh_mode: RefreshMode = "incremental",
    request_timeout: int = 6,
    retries: int = 2,
    logger=None,
    progress_callback: Optional[ProgressCallback] = None,
) -> tuple[dict, dict[str, int | list[dict[str, str]]]]:
    """Enrich parsed identities in memory and return updated data + stats.

    refresh_mode:
        - incremental: default, uses fingerprint cache checks
        - refresh_changed_only: semantic alias of incremental mode
        - refresh_all: bypasses cache and re-fetches all records
    """
    cache = EnrichmentCache(config.enrichment_cache_db)
    providers = _build_providers(
        config,
        request_timeout=request_timeout,
        retries=retries,
    )

    total_records = (
        len(identities_data.get("movies", []))
        + len(identities_data.get("series", []))
        + len(identities_data.get("anime", []))
    )

    stats: dict[str, int | list[dict[str, str]]] = {
        "total": total_records,
        "processed": 0,
        "enriched": 0,
        "cache_hits": 0,
        "skipped": 0,
        "needs_review": 0,
        "failed": 0,
        "api_calls": 0,
        "failed_items": [],
        "skip_reasons": {},
    }

    refresh_all = refresh_mode == "refresh_all"

    for section, media_type in (("movies", "movie"), ("series", "series"), ("anime", "anime")):
        records = identities_data.get(section, [])
        for record in records:
            title = record.get("title") or _fallback_title_from_filename(record.get("filename"))
            if not title:
                stats["skipped"] = int(stats["skipped"]) + 1
                _increment_skip_reason(stats, "invalid_input")
                stats["processed"] = int(stats["processed"]) + 1
                _emit_progress(stats, progress_callback)
                continue

            if "title" not in record:
                record["title"] = title

            fingerprint = _fingerprint(record, media_type)
            identity_key = _identity_key(record, media_type)

            cached = None
            if not refresh_all:
                cached = cache.get_identity(identity_key, fingerprint)

            if cached:
                _apply_payload(record, cached)
                stats["cache_hits"] = int(stats["cache_hits"]) + 1
                if record.get("needs_review"):
                    stats["needs_review"] = int(stats["needs_review"]) + 1
                stats["processed"] = int(stats["processed"]) + 1
                _emit_progress(stats, progress_callback)
                continue

            payload, api_calls, failures, skip_reason = _enrich_record(
                record,
                media_type,
                providers,
                config,
                request_timeout=request_timeout,
                retries=retries,
            )

            _apply_payload(record, payload)
            cache.put_identity(identity_key, fingerprint, payload)
            _update_stats_after_enrich(
                stats, payload, failures, api_calls, record, progress_callback, skip_reason
            )

    metadata = identities_data.setdefault("metadata", {})
    metadata["enriched"] = True
    metadata["enrichment_policy"] = "incremental" if refresh_mode != "refresh_all" else "full"
    metadata["enrichment_refresh_mode"] = refresh_mode

    if logger:
        logger.info(
            "Enrichment completed: total=%s enriched=%s cache_hits=%s skipped=%s failed=%s api_calls=%s",
            stats["total"],
            stats["enriched"],
            stats["cache_hits"],
            stats["skipped"],
            stats["failed"],
            stats["api_calls"],
        )

    return identities_data, stats


def _emit_progress(stats: dict[str, int | list[dict[str, str]]], progress_callback: Optional[ProgressCallback]) -> None:
    if not progress_callback:
        return
    progress_callback(
        int(stats["processed"]),
        int(stats["total"]),
        {
            "cache_hits": int(stats["cache_hits"]),
            "api_calls": int(stats["api_calls"]),
            "failed": int(stats["failed"]),
        },
    )


def _update_stats_after_enrich(
    stats: dict[str, int | list[dict[str, str]]],
    payload: dict,
    failures: list[dict[str, str]],
    api_calls: int,
    record: dict,
    progress_callback: Optional[ProgressCallback],
    skip_reason: str,
) -> None:
    """Update stats and emit progress after enriching a single record."""
    stats["api_calls"] = int(stats["api_calls"]) + api_calls
    if failures:
        failed_items = stats["failed_items"]
        assert isinstance(failed_items, list)
        failed_items.extend(failures)
        stats["failed"] = int(stats["failed"]) + len(failures)
    if payload.get("enriched"):
        stats["enriched"] = int(stats["enriched"]) + 1
    else:
        stats["skipped"] = int(stats["skipped"]) + 1
        _increment_skip_reason(stats, skip_reason or "no_match")
    if record.get("needs_review"):
        stats["needs_review"] = int(stats["needs_review"]) + 1
    stats["processed"] = int(stats["processed"]) + 1
    _emit_progress(stats, progress_callback)


def _build_providers(config: Config, *, request_timeout: int, retries: int) -> list:
    providers = []
    unsupported: list[str] = []
    for name in config.enrichment_providers:
        key = name.lower()
        if key == "tmdb":
            providers.append(
                TMDBProvider(
                    config.tmdb_api_key,
                    bearer_token=config.tmdb_bearer_token,
                    language=config.tmdb_language,
                    region=config.tmdb_region,
                    include_adult=config.tmdb_include_adult,
                    timeout_seconds=request_timeout,
                    retries=retries,
                    min_interval_seconds=0.25,
                )
            )
        else:
            unsupported.append(name)

    if unsupported:
        raise ValueError(
            f"Unsupported enrichment providers: {unsupported}. Supported providers: ['tmdb']"
        )

    return providers


def _enrich_record(
    record: dict,
    media_type: str,
    providers: list,
    config: Config,
    *,
    request_timeout: int,
    retries: int,
) -> tuple[dict, int, list[dict[str, str]], str]:
    title = record.get("title")
    year = record.get("year") if media_type == "movie" else None

    provider_results: list[ProviderResult] = []
    failures: list[dict[str, str]] = []
    api_calls = 0
    configured_provider_count = 0

    for provider in providers:
        if not _provider_is_configured(provider, config):
            continue

        configured_provider_count += 1
        try:
            result = provider.enrich(title=title, media_type=media_type, year=year)
        except TMDBAuthError as exc:
            raise RuntimeError(str(exc)) from exc
        except TMDBProviderError as exc:
            failures.append(
                {
                    "path": str(record.get("path", "")),
                    "title": str(title),
                    "provider": provider.name,
                    "reason": str(exc),
                }
            )
            api_calls += provider.last_request_count
            continue
        except Exception as exc:
            failures.append(
                {
                    "path": str(record.get("path", "")),
                    "title": str(title),
                    "provider": provider.name,
                    "reason": str(exc),
                }
            )
            api_calls += provider.last_request_count
            continue

        api_calls += provider.last_request_count
        if result:
            provider_results.append(result)

    merged = _merge_provider_results(provider_results)

    # Optional AI fallback for missing translated titles.
    if config.translation_fallback_machine and config.openai_api_key:
        if not merged.get("title_zh"):
            api_calls += 1
            translated = _translate_with_openai(
                title,
                target_language="Chinese (Simplified)",
                api_key=config.openai_api_key,
                timeout_seconds=request_timeout,
                retries=retries,
            )
            if translated:
                merged["title_zh"] = translated
                merged["translation_source"] = merged.get("translation_source") or "openai"

        if not merged.get("title_en"):
            api_calls += 1
            translated = _translate_with_openai(
                title,
                target_language="English",
                api_key=config.openai_api_key,
                timeout_seconds=request_timeout,
                retries=retries,
            )
            if translated:
                merged["title_en"] = translated
                merged["translation_source"] = merged.get("translation_source") or "openai"

    confidence = _enrichment_confidence(merged)
    merged["enrichment_confidence"] = confidence

    review_status = record.get("review_status", "pending")
    needs_review = bool(record.get("needs_review", False))

    if confidence < config.enrichment_min_match_score:
        needs_review = True

    score = merged.get("reputation_score")
    votes = merged.get("reputation_votes") or 0
    if (
        score is not None
        and votes >= config.reputation_min_votes
        and score < config.reputation_low_score_threshold
    ):
        needs_review = True

    merged["review_status"] = review_status
    merged["needs_review"] = needs_review
    merged["enriched"] = bool(provider_results or merged.get("translation_source"))
    merged["display_title"] = _build_display_title(record, merged, config)

    skip_reason = _determine_skip_reason(
        provider_results=provider_results,
        failures=failures,
        configured_provider_count=configured_provider_count,
        api_calls=api_calls,
    )

    return merged, api_calls, failures, skip_reason


def _determine_skip_reason(
    *,
    provider_results: list[ProviderResult],
    failures: list[dict[str, str]],
    configured_provider_count: int,
    api_calls: int,
) -> str:
    if provider_results:
        return ""
    if configured_provider_count == 0 and api_calls == 0:
        return "no_key"
    if failures:
        for failure in failures:
            reason = str(failure.get("reason", "")).lower()
            if "rate limit" in reason or "(429)" in reason:
                return "rate_limited"
            if "authentication failed" in reason or "(401/403)" in reason:
                return "auth_error"
        return "provider_error"
    return "no_match"


def _increment_skip_reason(stats: dict[str, int | list[dict[str, str]]], reason: str) -> None:
    if not reason:
        return
    current = stats.get("skip_reasons")
    if not isinstance(current, dict):
        current = {}
        stats["skip_reasons"] = current
    current[reason] = int(current.get(reason, 0)) + 1


def _provider_is_configured(provider: object, config: Config) -> bool:
    provider_name = provider.name.lower()

    if provider_name == "tmdb":
        return bool(config.tmdb_bearer_token or config.tmdb_api_key)

    return True


def _merge_provider_results(results: list[ProviderResult]) -> dict:
    payload: dict = {
        "canonical_id": None,
        "title_zh": None,
        "title_en": None,
        "translation_source": None,
        "reputation_score": None,
        "reputation_votes": None,
        "reputation_source": None,
        "provider_metadata": {},
    }

    if not results:
        return payload

    first = results[0]
    payload["canonical_id"] = first.canonical_id
    payload["title_zh"] = first.title_zh
    payload["title_en"] = first.title_en
    payload["translation_source"] = first.translation_source

    total_weight = 0
    weighted_score = 0.0
    source_names = []

    for result in results:
        source_names.append(result.provider)
        payload["provider_metadata"][result.provider] = json.dumps(result.raw_metadata, ensure_ascii=False)

        if not payload["title_zh"] and result.title_zh:
            payload["title_zh"] = result.title_zh
            payload["translation_source"] = result.translation_source or result.provider

        if not payload["title_en"] and result.title_en:
            payload["title_en"] = result.title_en

        if result.reputation_score is None:
            continue

        votes = result.reputation_votes if result.reputation_votes and result.reputation_votes > 0 else 1
        weighted_score += result.reputation_score * votes
        total_weight += votes

    if total_weight > 0:
        payload["reputation_score"] = round(weighted_score / total_weight, 3)
        payload["reputation_votes"] = total_weight
        payload["reputation_source"] = "+".join(sorted(set(source_names)))

    return payload


def _enrichment_confidence(payload: dict) -> float:
    score = 0.0
    if payload.get("canonical_id"):
        score += 0.4
    if payload.get("title_zh"):
        score += 0.2
    if payload.get("title_en"):
        score += 0.2
    if payload.get("reputation_score") is not None:
        score += 0.2
    return round(score, 3)


def _translate_with_openai(
    text: str,
    *,
    target_language: str,
    api_key: Optional[str],
    timeout_seconds: int,
    retries: int,
) -> Optional[str]:
    if not api_key:
        return None

    body = {
        "model": "gpt-4o-mini",
        "input": (
            f"Translate the movie or TV title into {target_language}. "
            "Return only the translated title without explanations."
            f"\nTitle: {text}"
        ),
    }
    request = Request(
        "https://api.openai.com/v1/responses",
        data=json.dumps(body).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
        method="POST",
    )

    for _ in range(max(retries + 1, 1)):
        try:
            with urlopen(request, timeout=timeout_seconds) as response:
                payload = json.loads(response.read().decode("utf-8"))
                output_text = payload.get("output_text")
                if isinstance(output_text, str) and output_text.strip():
                    return output_text.strip()
        except Exception:
            continue

    return None


def _apply_payload(record: dict, payload: dict) -> None:
    for key in (
        "canonical_id",
        "title_zh",
        "title_en",
        "translation_source",
        "reputation_score",
        "reputation_votes",
        "reputation_source",
        "review_status",
        "enrichment_confidence",
        "provider_metadata",
        "display_title",
    ):
        if key in payload:
            record[key] = payload[key]

    if "needs_review" in payload:
        record["needs_review"] = payload["needs_review"]


def _identity_key(record: dict, media_type: str) -> str:
    return f"{media_type}:{record.get('path', '')}"


def _fingerprint(record: dict, media_type: str) -> str:
    fields = [
        media_type,
        str(record.get("path", "")),
        str(record.get("title", "")),
        str(record.get("year", "")),
        str(record.get("season", "")),
        json.dumps(record.get("episodes", [])),
    ]
    digest = hashlib.sha256("|".join(fields).encode("utf-8")).hexdigest()
    return digest


def _fallback_title_from_filename(filename: Optional[str]) -> Optional[str]:
    if not filename:
        return None
    base = filename.rsplit(".", 1)[0]
    return normalize_title(base)


def _build_display_title(record: dict, payload: dict, config: Config) -> str:
    fallback_title = record.get("title")
    title_zh = payload.get("title_zh") or fallback_title
    title_en = payload.get("title_en") or fallback_title
    if title_zh and title_en and title_zh == title_en:
        title_en = ""
    try:
        formatted = config.naming_title_format.format(
            title_zh=title_zh or "",
            title_en=title_en or "",
        ).strip()
    except Exception:
        formatted = f"{title_zh or ''} {title_en or ''}".strip()
    return " ".join(formatted.split())