AIChatExporter/src/media.py

"""Media resolution — download conversation assets next to the export files.

Walks a normalized conversation's placeholder blocks, downloads each asset via
the provider, saves it under a ``media/`` directory beside the conversation's
Markdown file, and annotates the block with a relative ``local_path`` so the
renderer emits a real inline image / file link instead of a placeholder.

Policy (``EXPORTER_DOWNLOAD_MEDIA``):
  images (default) — image_placeholder blocks only
  all              — also file_placeholder blocks (voice-mode audio etc.;
                     measured 2026-06-12: 556 clips ≈ 162MB whose transcripts
                     are already in the exports as text)
  off              — leave placeholders untouched

Failures (expired assets, network) keep the existing placeholder and are
counted in the run summary — never fatal to the export.
"""

import logging
import os
import tempfile
from pathlib import Path

from src.blocks import BLOCK_TYPE_FILE_PLACEHOLDER, BLOCK_TYPE_IMAGE_PLACEHOLDER
from src.loss_report import LossReport
from src.providers.base import ProviderError
from src.utils import build_export_path, generate_filename

logger = logging.getLogger(__name__)

MEDIA_IMAGES = "images"
MEDIA_ALL = "all"
MEDIA_OFF = "off"
VALID_MEDIA_POLICIES = {MEDIA_IMAGES, MEDIA_ALL, MEDIA_OFF}

_EXT_BY_MIME = {
    "image/png": ".png",
    "image/jpeg": ".jpg",
    "image/gif": ".gif",
    "image/webp": ".webp",
    "audio/wav": ".wav",
    "audio/x-wav": ".wav",
    "audio/mpeg": ".mp3",
    "audio/mp4": ".m4a",
    "audio/aac": ".aac",
    "application/pdf": ".pdf",
}


def resolve_media_policy() -> str:
    """Read EXPORTER_DOWNLOAD_MEDIA from the environment, defaulting to images."""
    value = os.getenv("EXPORTER_DOWNLOAD_MEDIA", "").strip().lower()
    if not value:
        return MEDIA_IMAGES
    if value not in VALID_MEDIA_POLICIES:
        logger.warning(
            "EXPORTER_DOWNLOAD_MEDIA=%r is invalid (expected images|all|off) "
            "— using 'images'.",
            value,
        )
        return MEDIA_IMAGES
    return value


def resolve_media(
    normalized: dict,
    provider,
    export_base: Path,
    structure: str,
    policy: str,
    report: LossReport,
) -> int:
    """Download assets for one normalized conversation. Returns download count.

    Mutates placeholder blocks in place (adds ``local_path``). Idempotent:
    an asset already on disk is annotated without an API call.
    """
    if policy == MEDIA_OFF:
        return 0

    download = getattr(provider, "download_asset", None)
    if download is None:
        # Provider has no remote assets (e.g. claude-code local transcripts).
        return 0

    wanted_types = {BLOCK_TYPE_IMAGE_PLACEHOLDER}
    if policy == MEDIA_ALL:
        wanted_types.add(BLOCK_TYPE_FILE_PLACEHOLDER)

    targets = [
        block
        for message in normalized.get("messages", [])
        for block in message.get("blocks", [])
        if block.get("type") in wanted_types and block.get("ref")
    ]
    if not targets:
        return 0

    # Same path computation the exporter uses, so media/ lands beside the .md.
    filename = generate_filename(
        normalized.get("title", "Untitled"),
        normalized.get("id", ""),
        normalized.get("created_at") or "2000-01-01",
    )
    conv_dir = build_export_path(
        export_base,
        normalized.get("provider", ""),
        normalized.get("project"),
        normalized.get("created_at") or "2000-01-01",
        filename,
        structure,
    ).parent
    media_dir = conv_dir / "media"

    downloaded = 0
    for block in targets:
        ref = block["ref"]
        file_id = _safe_asset_name(provider, ref)
        if not file_id:
            report.record_media_failed("unparseable-ref")
            continue

        existing = _find_existing(media_dir, file_id)
        if existing is not None:
            block["local_path"] = f"media/{existing.name}"
            continue

        try:
            content, mime, file_name = download(ref)
        except ProviderError as e:
            logger.warning(
                "[media] Could not download %s: %s", ref[:60], e.original
            )
            reason = "expired-or-missing" if "404" in str(e.original) or "not found" in str(
                e.original
            ).lower() else "download-error"
            report.record_media_failed(reason)
            continue

        ext = _pick_extension(mime, file_name)
        target = media_dir / f"{file_id}{ext}"
        try:
            _write_atomic(target, content)
        except OSError as e:
            logger.error("[media] Could not write %s: %s", target, e)
            report.record_media_failed("write-error")
            continue

        block["local_path"] = f"media/{target.name}"
        report.record_media_downloaded(len(content))
        downloaded += 1

    return downloaded


def _safe_asset_name(provider, ref: str) -> str | None:
    """A stable, filesystem-safe name for the asset (the provider file ID)."""
    parser = getattr(provider, "parse_asset_file_id", None)
    if parser is None:
        from src.providers.chatgpt import parse_asset_file_id as parser
    file_id = parser(ref)
    if not file_id:
        return None
    return "".join(c if c.isalnum() or c in "-_." else "_" for c in file_id)


def _find_existing(media_dir: Path, file_id: str) -> Path | None:
    """Return an already-downloaded file for this asset, if present."""
    if not media_dir.is_dir():
        return None
    for candidate in media_dir.glob(f"{file_id}.*"):
        if candidate.is_file() and candidate.stat().st_size > 0:
            return candidate
    return None


def _pick_extension(mime: str | None, file_name: str | None) -> str:
    if mime:
        base_mime = mime.split(";")[0].strip().lower()
        if base_mime in _EXT_BY_MIME:
            return _EXT_BY_MIME[base_mime]
    if file_name:
        suffix = Path(file_name).suffix
        if suffix and len(suffix) <= 8:
            return suffix.lower()
    return ".bin"


def _write_atomic(target: Path, content: bytes) -> None:
    target.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp_name = tempfile.mkstemp(dir=target.parent, suffix=".tmp")
    try:
        with os.fdopen(fd, "wb") as fh:
            fh.write(content)
        os.chmod(tmp_name, 0o600)
        os.replace(tmp_name, target)
    except OSError:
        try:
            os.unlink(tmp_name)
        except OSError:
            pass
        raise