#!/usr/bin/env python3
"""Generate the synthetic e2e media library.

Every image is rendered with its own relative path, date and scenario
burned into the pixels, so an agent looking at a screenshot can tell
exactly which fixture is on screen. EXIF (dates, GPS, camera) is stamped
in a single exiftool batch. A `manifest.json` next to the media records
what was generated and what the server should derive from it, so tests
can assert against it instead of hardcoding counts.

The harness runs in WSL but the tools (ffmpeg.exe / exiftool.exe) are
Windows binaries, so every path handed to them goes through `wslpath -w`.

Usage: gen_fixtures.py <out_dir> [--import DIR]
  <out_dir>/media/          the pristine library (never served directly)
  <out_dir>/manifest.json   what's in it
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import shutil
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from zoneinfo import ZoneInfo

FFMPEG = shutil.which("ffmpeg.exe") or shutil.which("ffmpeg")
EXIFTOOL = shutil.which("exiftool.exe") or shutil.which("exiftool")
# drawtext needs an explicit font on Windows ffmpeg builds. The colon after
# the drive letter must be escaped inside a filter argument.
FONT = "C\\:/Windows/Fonts/consola.ttf"

# The owner's display zone and the server zone the e2e env pins (TZ=...).
# Naive fixture datetimes without GPS are walls in this zone.
HOME_TZ = ZoneInfo("America/New_York")

PLACES = {
    "santa_monica": (34.0100, -118.4960),
    "paris": (48.8584, 2.2945),
    "yosemite": (37.7456, -119.5936),
    "tokyo": (35.6595, 139.7005),
}

# IANA zone of each place — what the server's GPS rung of the offset cascade
# resolves, so the expected true instant is the wall read in this zone.
PLACE_ZONES = {
    "santa_monica": "America/Los_Angeles",
    "paris": "Europe/Paris",
    "yosemite": "America/Los_Angeles",
    "tokyo": "Asia/Tokyo",
}

PALETTE = [
    "0x1f4e79", "0x7a3e9d", "0x2e7d32", "0xb71c1c", "0xe65100",
    "0x006064", "0x4e342e", "0x283593", "0x827717", "0xad1457",
]


@dataclass
class Item:
    rel: str  # path relative to the library root, forward slashes
    kind: str  # "image" | "video"
    size: tuple[int, int] = (1600, 1200)
    # Capture wall clock: EXIF DateTimeOriginal as-is; for videos, the wall in
    # HOME_TZ, stamped as UTC creation_time the way phones write containers.
    taken: dt.datetime | None = None
    gps: str | None = None  # key into PLACES
    camera: tuple[str, str] | None = None  # (Make, Model)
    mtime: dt.datetime | None = None  # pinned file mtime (undated fallback)
    duration: int = 0  # video seconds
    tags: list[str] = field(default_factory=list)
    scenario: str = ""
    copy_of: str | None = None  # byte-identical duplicate of another item

    def label_date(self) -> str:
        if self.taken:
            return self.taken.strftime("%Y-%m-%d %H:%M")
        return "no EXIF date"


def win(p: Path) -> str:
    return subprocess.check_output(["wslpath", "-w", str(p)], text=True).strip()


def build_items(today: dt.date) -> list[Item]:
    items: list[Item] = []
    canon = ("Canon", "Canon EOS R5")
    iphone = ("Apple", "iPhone 12")
    pixel = ("Google", "Pixel 7")

    # Dated album with GPS — map pins, date grouping, tag search.
    for i in range(6):
        items.append(Item(
            f"2019/2019-07-04 Beach Day/beach_{i + 1:02d}.jpg", "image",
            taken=dt.datetime(2019, 7, 4, 10 + i, 15), gps="santa_monica",
            camera=canon, tags=["beach", "summer"], scenario="dated album + GPS",
        ))
    # Multi-day trip; mix of orientations.
    for i in range(8):
        items.append(Item(
            f"2021/Paris Trip/IMG_{1000 + i}.jpg", "image",
            size=(1200, 1600) if i % 3 == 0 else (1600, 1200),
            taken=dt.datetime(2021, 5, 12 + i // 3, 9 + i, 30), gps="paris",
            camera=iphone, tags=["travel", "paris"], scenario="trip, mixed orientation",
        ))
    for i in range(5):
        items.append(Item(
            f"2022/Hiking/hike_{i + 1:02d}.jpg", "image",
            taken=dt.datetime(2022, 8, 20, 7 + i, 0), gps="yosemite",
            camera=pixel, tags=["outdoors"], scenario="album with a video",
        ))
    items.append(Item(
        "2022/Hiking/summit.mp4", "video", size=(1280, 720), duration=6,
        taken=dt.datetime(2022, 8, 20, 12, 30), tags=["outdoors"],
        scenario="landscape video (HLS)",
    ))
    # No GPS, tagged people-ish album.
    for i in range(6):
        items.append(Item(
            f"2023/Birthday/party_{i + 1:02d}.jpg", "image",
            taken=dt.datetime(2023, 11, 3, 18, 5 * i), camera=iphone,
            tags=["family"], scenario="dated, no GPS",
        ))
    # Enough items to force scrolling / pagination in the grid.
    start = dt.datetime(2020, 1, 3, 12, 0)
    for i in range(60):
        items.append(Item(
            f"Grid Stress/photo_{i + 1:03d}.jpg", "image", size=(800, 600),
            taken=start + dt.timedelta(days=6 * i), camera=pixel,
            scenario="bulk grid / scroll",
        ))
    # Dates that only exist in the filename (the memories.rs regexes).
    items += [
        Item("Screenshots/Screenshot_20220315-101500.png", "image", size=(1080, 2340),
             scenario="date from filename (YYYYMMDD-HHMMSS), PNG, no EXIF"),
        Item("Screenshots/IMG_20210704_123000.jpg", "image",
             scenario="date from filename (YYYYMMDD_HHMMSS), no EXIF"),
        Item("Screenshots/2020-02-29_08-00-00.jpg", "image",
             scenario="date from filename (leap day), no EXIF"),
    ]
    # No date anywhere: falls back to the pinned mtime.
    for i in range(3):
        items.append(Item(
            f"Undated/scan_{i + 1:02d}.jpg", "image",
            mtime=dt.datetime(2015, 3, 1 + i, 12, 0),
            scenario="no EXIF, no filename date (mtime fallback)",
        ))
    # "On this day" — same month/day as generation date in earlier years.
    for year in (2018, 2020, 2024):
        taken = _same_day(today, year, 14)
        items.append(Item(
            f"On This Day/otd_{year}.jpg", "image", taken=taken, gps="tokyo",
            camera=canon, scenario=f"memories: on this day ({year})",
        ))
    # Excluded via EXCLUDED_DIRS=/private — never indexed, must not appear anywhere.
    items.append(Item(
        "private/otd_hidden.jpg", "image", taken=_same_day(today, 2019, 9),
        scenario="EXCLUDED_DIRS: never indexed or shown",
    ))
    # Portrait video, filename-dated.
    items.append(Item(
        "Videos/VID_20230610_150000.mp4", "video", size=(720, 1280), duration=4,
        taken=dt.datetime(2023, 6, 10, 15, 0), scenario="portrait video",
    ))
    # Byte-identical duplicate for duplicate detection.
    items.append(Item(
        "Duplicates/IMG_1000 copy.jpg", "image", copy_of="2021/Paris Trip/IMG_1000.jpg",
        scenario="exact duplicate of 2021/Paris Trip/IMG_1000.jpg",
    ))
    return items


def expected_wall_clock(it: Item) -> str:
    """The wall clock the server should report (`wall_clock`): the filename
    date for the filename class (it outranks EXIF), else the capture wall,
    else the pinned mtime rendered in HOME_TZ."""
    import re
    name = it.rel.rsplit("/", 1)[-1]
    m = (re.search(r"(\d{4})(\d{2})(\d{2})[_-](\d{2})(\d{2})(\d{2})", name)
         or re.search(r"(\d{4})-(\d{2})-(\d{2})_(\d{2})-(\d{2})-(\d{2})", name))
    if m:
        wall = dt.datetime(*map(int, m.groups()))
    else:
        wall = it.taken or it.mtime or dt.datetime(2024, 1, 1)
    return wall.strftime("%Y-%m-%dT%H:%M:%S")


def expected_true_instant(it: Item) -> int:
    """The true instant the server should report: the expected wall read in
    the zone the camera was in (GPS place, else HOME_TZ — the server's
    assumed zone). Filename walls outrank EXIF, but still take GPS evidence."""
    wall = dt.datetime.fromisoformat(expected_wall_clock(it))
    zone = ZoneInfo(PLACE_ZONES[it.gps]) if it.gps else HOME_TZ
    return int(wall.replace(tzinfo=zone).timestamp())


def _same_day(today: dt.date, year: int, hour: int) -> dt.datetime:
    # Feb 29 in a non-leap year would be invalid; fall back to Feb 28.
    try:
        return dt.datetime(year, today.month, today.day, hour, 0)
    except ValueError:
        return dt.datetime(year, today.month, 28, hour, 0)


def render(item: Item, idx: int, media: Path) -> None:
    out = media / item.rel
    out.parent.mkdir(parents=True, exist_ok=True)
    w, h = item.size
    color = PALETTE[idx % len(PALETTE)]
    lines = [item.rel, item.label_date(), item.scenario]
    # Consolas glyphs are ~0.55em wide; shrink so the longest line fits.
    longest = max(len(line) for line in lines)
    fs = max(12, min(min(w, h) // 22, int(w * 0.92 / (longest * 0.55))))
    text = ",".join(
        f"drawtext=fontfile='{FONT}':text='{_esc(line)}':fontcolor=white:"
        f"fontsize={fs}:x=(w-text_w)/2:y=(h/2)+({n - 1}*{int(fs * 1.6)})"
        f":box=1:boxcolor=black@0.45:boxborderw=8"
        for n, line in enumerate(lines)
    )
    if item.kind == "image":
        cmd = [
            FFMPEG, "-hide_banner", "-loglevel", "error", "-y",
            "-f", "lavfi", "-i",
            f"gradients=s={w}x{h}:c0={color}:c1=0x101010:seed={idx}:duration=1",
            "-vf", text, "-frames:v", "1",
        ]
        cmd += ["-q:v", "3"] if out.suffix.lower() == ".jpg" else []
    else:
        cmd = [
            FFMPEG, "-hide_banner", "-loglevel", "error", "-y",
            "-f", "lavfi", "-i", f"testsrc2=s={w}x{h}:r=30:d={item.duration}",
            "-f", "lavfi", "-i", f"sine=frequency={300 + idx * 20}:d={item.duration}",
            "-vf", text, "-c:v", "libx264", "-pix_fmt", "yuv420p",
            "-c:a", "aac", "-shortest",
        ]
        if item.taken:
            utc = item.taken.replace(tzinfo=HOME_TZ).astimezone(dt.timezone.utc)
            cmd += ["-metadata", f"creation_time={utc.strftime('%Y-%m-%dT%H:%M:%S')}Z"]
    subprocess.run(cmd + [win(out)], check=True)


def _esc(s: str) -> str:
    # drawtext text escaping: backslash, colon, quote, percent, comma.
    for ch in ("\\", ":", "'", "%", ","):
        s = s.replace(ch, "\\" + ch)
    return s


def stamp_exif(items: list[Item], media: Path, work: Path) -> None:
    args: list[str] = []
    for it in items:
        if it.kind != "image" or it.copy_of or not (it.taken or it.gps or it.camera):
            continue
        if not it.rel.lower().endswith(".jpg"):
            continue
        if it.taken:
            stamp = it.taken.strftime("%Y:%m:%d %H:%M:%S")
            args += [f"-DateTimeOriginal={stamp}", f"-CreateDate={stamp}"]
        if it.camera:
            args += [f"-Make={it.camera[0]}", f"-Model={it.camera[1]}"]
        if it.gps:
            lat, lon = PLACES[it.gps]
            args += [
                f"-GPSLatitude={abs(lat)}", f"-GPSLatitudeRef={'N' if lat >= 0 else 'S'}",
                f"-GPSLongitude={abs(lon)}", f"-GPSLongitudeRef={'E' if lon >= 0 else 'W'}",
            ]
        args += [win(media / it.rel), "-execute"]
    argfile = work / "exif_args.txt"
    argfile.write_text("\n".join(args) + "\n", encoding="utf-8")
    subprocess.run(
        [EXIFTOOL, "-@", win(argfile), "-common_args", "-q", "-overwrite_original",
         "-charset", "filename=utf8"],
        check=True,
    )
    argfile.unlink()


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("out_dir", type=Path)
    ap.add_argument("--import", dest="import_dir", type=Path,
                    help="also copy this folder (read-only source) into media/imported/")
    a = ap.parse_args()
    if not FFMPEG or not EXIFTOOL:
        print("ffmpeg(.exe) and exiftool(.exe) must be on PATH", file=sys.stderr)
        return 1

    out: Path = a.out_dir.resolve()
    media = out / "media"
    if media.exists():
        shutil.rmtree(media)
    media.mkdir(parents=True)
    (media / "Empty Album").mkdir()  # edge case: folder with no media

    today = dt.date.today()
    items = build_items(today)
    originals = [(i, it) for i, it in enumerate(items) if not it.copy_of]
    with ThreadPoolExecutor(max_workers=8) as pool:
        list(pool.map(lambda p: render(p[1], p[0], media), originals))
    stamp_exif(items, media, out)

    for it in items:
        if it.copy_of:
            (media / it.rel).parent.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(media / it.copy_of, media / it.rel)
    # Pin mtimes last: exiftool rewrites (and so touches) the files.
    for it in items:
        when = it.mtime or it.taken or dt.datetime(2024, 1, 1)
        ts = when.replace(tzinfo=HOME_TZ).timestamp()  # not WSL's local zone
        os.utime(media / it.rel, (ts, ts))

    imported = 0
    if a.import_dir:
        dest = media / "imported"
        shutil.copytree(a.import_dir, dest)  # copy: the source is never touched
        imported = sum(1 for p in dest.rglob("*") if p.is_file())

    by_path = {it.rel: it for it in items}
    tags: dict[str, list[str]] = {}
    for it in items:
        for t in it.tags:
            tags.setdefault(t, []).append(it.rel)
    manifest = {
        "generated_on": today.isoformat(),
        "note": "'On This Day' fixtures match generated_on's month/day; regenerate "
                "(e2e.sh fixtures --force) on another day to keep memories populated.",
        "counts": {
            "images": sum(1 for it in items if it.kind == "image"),
            "videos": sum(1 for it in items if it.kind == "video"),
            "total": len(items),
            # What the server should index: EXCLUDED_DIRS is skipped at ingest.
            "indexed": sum(1 for it in items if not it.rel.startswith("private/")) + imported,
            "imported_files": imported,
        },
        "excluded_dirs": ["/private"],
        "places": {k: {"lat": v[0], "lon": v[1]} for k, v in PLACES.items()},
        "tags": tags,
        "items": [
            {
                "path": it.rel, "kind": it.kind, "scenario": it.scenario,
                "width": it.size[0], "height": it.size[1],
                "taken": it.taken.isoformat() if it.taken else None,
                "gps": it.gps, "camera": it.camera[1] if it.camera else None,
                "tags": it.tags, "duplicate_of": it.copy_of,
                # A duplicate carries its source's bytes, EXIF included.
                "wall_clock": expected_wall_clock(by_path.get(it.copy_of, it)),
                "true_instant": expected_true_instant(by_path.get(it.copy_of, it)),
            }
            for it in items
        ],
    }
    (out / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
    print(f"generated {len(items)} fixtures ({imported} imported) in {media}")
    return 0


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