#!/usr/bin/env python3
"""
Apply image_descriptions.jsonl captions into with_images.csv.

Output format for each with-images row:
  img1(file): ... || img2(file): ...

Behavior:
- Append-only jsonl; last description for a file wins.
- Description text is normalized to one line.
- If a downloaded image is missing in jsonl, writes [MISSING DESCRIPTION].
- Avoids clobbering existing non-empty image_description when zero mapped captions
  exist for that row (safety for manual runs).
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import sys
from pathlib import Path
from typing import Dict, Iterable, List, Tuple

ROOT = Path(__file__).resolve().parents[2]


def single_line(text: str) -> str:
    return " ".join((text or "").replace("\r", " ").replace("\n", " ").split()).strip()


def iter_jsonl(path: Path) -> Iterable[dict]:
    if not path.exists():
        return
    with path.open("r", encoding="utf-8") as f:
        for idx, raw_line in enumerate(f, 1):
            line = raw_line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except json.JSONDecodeError as exc:
                print(f"Warning: skipping invalid jsonl line {idx}: {exc}", file=sys.stderr)
                continue
            if not isinstance(obj, dict):
                print(f"Warning: skipping non-object jsonl line {idx}", file=sys.stderr)
                continue
            yield obj


def load_caption_map(path: Path) -> Dict[str, str]:
    out: Dict[str, str] = {}
    for obj in iter_jsonl(path):
        file_name = single_line(str(obj.get("file", "")))
        description = single_line(str(obj.get("description", "")))
        if not file_name or not description:
            continue
        out[file_name] = description
    return out


def split_downloaded_images(value: str) -> List[str]:
    return [x.strip() for x in (value or "").split("|") if x.strip()]


def build_description(downloaded_images: str, caption_map: Dict[str, str]) -> Tuple[str, int]:
    """
    Returns (combined_description, mapped_count).
    mapped_count tracks how many files had captions in jsonl.
    """
    files = split_downloaded_images(downloaded_images)
    parts: List[str] = []
    mapped = 0
    for idx, file_name in enumerate(files, 1):
        caption = caption_map.get(file_name, "")
        if caption:
            mapped += 1
            parts.append(f"img{idx}({file_name}): {caption}")
        else:
            parts.append(f"img{idx}({file_name}): [MISSING DESCRIPTION]")
    return " || ".join(parts), mapped


def apply_for_date(date_iso: str) -> int:
    day_dir = ROOT / "PHASES" / "PHASE-01-INGEST" / "tweets" / "archive" / date_iso
    with_images_path = day_dir / "with_images.csv"
    jsonl_path = day_dir / "image_descriptions.jsonl"

    if not with_images_path.exists():
        print(f"ERROR: missing {with_images_path}")
        return 2

    caption_map = load_caption_map(jsonl_path)
    print(f"Loaded captions: {len(caption_map)} file entries from {jsonl_path}")

    backup_path = day_dir / "with_images.before_apply.csv"
    if not backup_path.exists():
        backup_path.write_bytes(with_images_path.read_bytes())

    with with_images_path.open("r", newline="", encoding="utf-8-sig") as f:
        reader = csv.DictReader(f)
        if not reader.fieldnames:
            print("ERROR: with_images.csv has no header")
            return 2
        fieldnames = list(reader.fieldnames)
        if "image_description" not in fieldnames:
            fieldnames.append("image_description")
        rows = list(reader)

    updated_rows = 0
    rows_with_placeholders = 0
    for row in rows:
        downloaded = single_line(str(row.get("downloaded_images", "")))
        if not downloaded:
            continue

        existing = single_line(str(row.get("image_description", "")))
        new_desc, mapped_count = build_description(downloaded, caption_map)

        # Safety: if this row has no mapped captions and already has a real value,
        # do not clobber it with placeholders.
        if mapped_count == 0 and existing:
            continue

        if "[MISSING DESCRIPTION]" in new_desc:
            rows_with_placeholders += 1

        if existing != new_desc:
            row["image_description"] = new_desc
            updated_rows += 1

    tmp_path = day_dir / ".with_images.csv.tmp"
    with tmp_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=fieldnames,
            quoting=csv.QUOTE_MINIMAL,
            extrasaction="ignore",
        )
        writer.writeheader()
        writer.writerows(rows)

    os.replace(tmp_path, with_images_path)
    print(
        "OK: applied captions "
        f"(updated_rows={updated_rows}, rows_with_placeholders={rows_with_placeholders})"
    )
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Apply image_descriptions.jsonl into with_images.csv"
    )
    parser.add_argument("date", help="Date in YYYY-MM-DD")
    args = parser.parse_args()
    return apply_for_date(args.date)


if __name__ == "__main__":
    raise SystemExit(main())
