"""
Pre-compute CLIP image embeddings for the training set.

Why a separate step?
--------------------
CLIP forward passes are the expensive part of Stage 2 training. The CLIP
backbone is FROZEN — only the small classification head defined in
`src/deepfake_scanner/detectors/clip_classifier.py` is being trained.
Encoding ~100k images on CPU takes 1-4 hours; training the head on the
cached embeddings then takes seconds per epoch and can be iterated freely
without re-encoding.

Local feasibility
-----------------
Yes — this script runs fine on the dev laptop (i7-1195G7, 32 GB).
With the default batch size of 16, peak RAM is ~1 GB. Throughput is
roughly 50-150 ms per image depending on thread count.

Same script runs ~50-100x faster on a CUDA GPU. Use `--device cuda` only
if you happen to be on a GPU box already (e.g. the Flux generation box,
before tearing it down).

Output format
-------------
One `.npz` file per split, in `--out-dir`:
  - `embeddings`: float32 array, shape (N, 512), L2-normalised
  - `labels`:     int8 array,    shape (N,)        0 = authentic, 1 = ai_generated
  - `paths`:      unicode array, shape (N,)        source image paths
  - `sources`:    unicode array, shape (N,)        manifest source values
  - `generators`: unicode array, shape (N,)        AI generator values or ""
  - `model_families`: unicode array, shape (N,)    generator family values or ""
  - `augmentations`: unicode array, shape (N,)     augmentation version or ""
  - `original_paths`: unicode array, shape (N,)    original path for augmented rows

The extra metadata arrays are ignored by `train_head.py` today, but they make
the cached embeddings auditable and give Stage 3A evaluation code enough
context to separate clean rows from augmented rows.

Resumability
------------
Long runs survive crashes, lid-closes, and Ctrl-C. Each split is encoded in
small shards (`{split}_shards/shard_NNNNNNNN.npz`); on restart we count the
contiguous shards already written from index 0 and resume from there. At
the end the shards are concatenated into the final `{split}.npz` and the
shard directory is removed.

Idempotent: an existing final `.npz` is skipped unless `--force` is passed.
`--force` also wipes any leftover shard directory.

Usage
-----
    python scripts/precompute_embeddings.py \\
        --splits-dir data \\
        --out-dir data/embeddings \\
        --data-root data

Include augmented training rows:

    python scripts/precompute_embeddings.py \\
        --splits-dir data \\
        --out-dir data/embeddings \\
        --data-root data \\
        --train-augment-manifest data/train_augmented.csv

Encode a separate augmented evaluation split:

    python scripts/precompute_embeddings.py \\
        --splits-dir data \\
        --out-dir data/embeddings \\
        --data-root data \\
        --extra-split test_augmented=data/test_augmented.csv
"""
from __future__ import annotations

import argparse
import csv
import shutil
import time
from dataclasses import dataclass
from pathlib import Path

import numpy as np


# Order MUST match the head's output ordering in
# src/deepfake_scanner/detectors/clip_classifier.py (logits[..., 0] = authentic,
# logits[..., 1] = ai_generated). Drift between train and inference here
# would silently invert predictions.
CLASS_TO_LABEL: dict[str, int] = {
    "authentic": 0,
    "ai_generated": 1,
}

# Images per shard. 500 × 512 × 4 bytes = 1 MB on disk; small enough that a
# crash loses at most ~25 seconds of CPU work, large enough that shard I/O
# overhead is negligible.
SHARD_SIZE = 500


@dataclass(frozen=True)
class EmbeddingItem:
    image_path: Path
    label: int
    manifest_path: str
    source: str = ""
    generator: str = ""
    model_family: str = ""
    augmentation: str = ""
    original_path: str = ""


def _read_manifest(path: Path, data_root: Path) -> list[EmbeddingItem]:
    items: list[EmbeddingItem] = []
    with path.open() as fh:
        reader = csv.DictReader(fh)
        for row in reader:
            cls = row["class"]
            if cls not in CLASS_TO_LABEL:
                raise ValueError(f"{path}: unknown class {cls!r}")
            manifest_path = row["path"]
            items.append(
                EmbeddingItem(
                    image_path=data_root / manifest_path,
                    label=CLASS_TO_LABEL[cls],
                    manifest_path=manifest_path,
                    source=row.get("source", ""),
                    generator=row.get("generator", ""),
                    model_family=row.get("model_family", ""),
                    augmentation=row.get("augmentation", ""),
                    original_path=row.get("original_path", ""),
                )
            )
    return items


def _read_split(
    path: Path,
    data_root: Path,
    extra_manifests: list[Path] | None = None,
) -> list[EmbeddingItem]:
    items = _read_manifest(path, data_root)
    for extra_manifest in extra_manifests or []:
        items.extend(_read_manifest(extra_manifest, data_root))

    seen: set[str] = set()
    duplicates: list[str] = []
    for item in items:
        if item.manifest_path in seen:
            duplicates.append(item.manifest_path)
        seen.add(item.manifest_path)
    if duplicates:
        raise ValueError(
            f"Duplicate image paths in split inputs for {path}: {duplicates[:5]}"
        )

    return items


def _resume_cursor(shard_dir: Path) -> int:
    """How many items have been encoded already, by counting the contiguous
    shard prefix starting at index 0. Corrupt shards are deleted in passing."""
    if not shard_dir.exists():
        return 0
    by_start: dict[int, int] = {}
    for p in sorted(shard_dir.glob("shard_*.npz")):
        try:
            with np.load(p) as z:
                start = int(z["start_idx"])
                end = int(z["end_idx"])
                required = [
                    "embeddings",
                    "labels",
                    "paths",
                    "sources",
                    "generators",
                    "model_families",
                    "augmentations",
                    "original_paths",
                ]
                for key in required:
                    if key not in z:
                        raise KeyError(key)
        except Exception:
            print(f"  removing corrupt shard {p.name}")
            p.unlink()
            continue
        by_start[start] = end
    cursor = 0
    while cursor in by_start:
        cursor = by_start[cursor]
    return cursor


def _save_shard_atomic(
    shard_path: Path,
    embeddings: np.ndarray,
    labels: np.ndarray,
    paths: np.ndarray,
    sources: np.ndarray,
    generators: np.ndarray,
    model_families: np.ndarray,
    augmentations: np.ndarray,
    original_paths: np.ndarray,
    start_idx: int,
    end_idx: int,
) -> None:
    """Write to .tmp first, then atomic rename. Avoids partial-write corruption
    if the process is killed mid-save."""
    tmp = shard_path.with_suffix(".tmp.npz")
    np.savez_compressed(
        tmp,
        embeddings=embeddings,
        labels=labels,
        paths=paths,
        sources=sources,
        generators=generators,
        model_families=model_families,
        augmentations=augmentations,
        original_paths=original_paths,
        start_idx=np.asarray(start_idx),
        end_idx=np.asarray(end_idx),
    )
    tmp.replace(shard_path)


def _concat_shards(shard_dir: Path, out_path: Path, expected_total: int) -> None:
    paths = sorted(shard_dir.glob("shard_*.npz"))
    embeds_list: list[np.ndarray] = []
    labels_list: list[np.ndarray] = []
    paths_list: list[np.ndarray] = []
    sources_list: list[np.ndarray] = []
    generators_list: list[np.ndarray] = []
    model_families_list: list[np.ndarray] = []
    augmentations_list: list[np.ndarray] = []
    original_paths_list: list[np.ndarray] = []
    for p in paths:
        with np.load(p) as z:
            embeds_list.append(z["embeddings"].copy())
            labels_list.append(z["labels"].copy())
            paths_list.append(z["paths"].copy())
            sources_list.append(z["sources"].copy())
            generators_list.append(z["generators"].copy())
            model_families_list.append(z["model_families"].copy())
            augmentations_list.append(z["augmentations"].copy())
            original_paths_list.append(z["original_paths"].copy())
    embeddings = np.concatenate(embeds_list, axis=0)
    labels = np.concatenate(labels_list, axis=0)
    manifest_paths = np.concatenate(paths_list, axis=0)
    sources = np.concatenate(sources_list, axis=0)
    generators = np.concatenate(generators_list, axis=0)
    model_families = np.concatenate(model_families_list, axis=0)
    augmentations = np.concatenate(augmentations_list, axis=0)
    original_paths = np.concatenate(original_paths_list, axis=0)
    if len(embeddings) != expected_total:
        raise RuntimeError(
            f"Shard concat produced {len(embeddings)} rows but split expected "
            f"{expected_total}. Refusing to write {out_path}."
        )
    np.savez_compressed(
        out_path,
        embeddings=embeddings,
        labels=labels,
        paths=manifest_paths,
        sources=sources,
        generators=generators,
        model_families=model_families,
        augmentations=augmentations,
        original_paths=original_paths,
    )
    shutil.rmtree(shard_dir)


def _process_split(
    split: str,
    split_csv: Path,
    out_dir: Path,
    data_root: Path,
    batch_size: int,
    force: bool,
    extra_manifests: list[Path],
    model,
    processor,
    device: str,
    torch_module,
    Image,
) -> None:
    out_path = out_dir / f"{split}.npz"
    shard_dir = out_dir / f"{split}_shards"

    if force:
        if out_path.exists():
            out_path.unlink()
        if shard_dir.exists():
            shutil.rmtree(shard_dir)

    if out_path.exists():
        print(
            f"Skipping {split} — {out_path} already exists. "
            "Use --force to recompute."
        )
        return

    if not split_csv.exists():
        print(f"Skipping {split} — {split_csv} not found.")
        return

    items = _read_split(split_csv, data_root, extra_manifests)
    n = len(items)
    print(f"\n{split}: {n} images")
    if extra_manifests:
        print(
            f"  includes {len(extra_manifests)} extra manifest(s): "
            + ", ".join(str(p) for p in extra_manifests)
        )

    cursor = _resume_cursor(shard_dir)
    if cursor > 0:
        print(f"  resuming from item {cursor}/{n} ({cursor / n:.1%} done)")
    shard_dir.mkdir(parents=True, exist_ok=True)

    start_t = time.time()
    items_at_start = cursor

    while cursor < n:
        shard_end = min(cursor + SHARD_SIZE, n)
        shard_n = shard_end - cursor
        shard_embeds = np.empty((shard_n, 512), dtype=np.float32)
        shard_items = items[cursor:shard_end]
        shard_labels = np.asarray(
            [item.label for item in shard_items],
            dtype=np.int8,
        )
        shard_paths = np.asarray([item.manifest_path for item in shard_items])
        shard_sources = np.asarray([item.source for item in shard_items])
        shard_generators = np.asarray([item.generator for item in shard_items])
        shard_model_families = np.asarray(
            [item.model_family for item in shard_items]
        )
        shard_augmentations = np.asarray(
            [item.augmentation for item in shard_items]
        )
        shard_original_paths = np.asarray(
            [item.original_path for item in shard_items]
        )

        pos = 0
        i = cursor
        while i < shard_end:
            batch_end = min(i + batch_size, shard_end)
            batch_paths = [item.image_path for item in items[i:batch_end]]
            batch_imgs = []
            for p in batch_paths:
                with Image.open(p) as img:
                    img_rgb = img.convert("RGB")
                    img_rgb.load()
                    batch_imgs.append(img_rgb.copy())

            inputs = processor(images=batch_imgs, return_tensors="pt").to(device)
            with torch_module.no_grad():
                feats = model.get_image_features(**inputs)
            # L2-normalise to match the inference path
            # (clip_classifier.py:94). Drift here breaks predictions.
            feats = feats / feats.norm(p=2, dim=-1, keepdim=True)
            shard_embeds[pos : pos + len(batch_imgs)] = (
                feats.cpu().numpy().astype(np.float32)
            )
            pos += len(batch_imgs)
            i = batch_end

        shard_path = shard_dir / f"shard_{cursor:08d}.npz"
        _save_shard_atomic(
            shard_path,
            shard_embeds,
            shard_labels,
            shard_paths,
            shard_sources,
            shard_generators,
            shard_model_families,
            shard_augmentations,
            shard_original_paths,
            cursor,
            shard_end,
        )
        cursor = shard_end

        elapsed = time.time() - start_t
        done_this_run = cursor - items_at_start
        rate = done_this_run / max(elapsed, 1e-6)
        eta = (n - cursor) / max(rate, 1e-6)
        print(
            f"  {cursor}/{n}  ({rate:.1f} img/s, "
            f"eta {eta / 60:.1f} min, shard saved)"
        )

    print(f"  concatenating shards into {out_path}...")
    _concat_shards(shard_dir, out_path, expected_total=n)
    print(f"  wrote {out_path}  ({out_path.stat().st_size / 1e6:.1f} MB)")


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--splits-dir", type=Path, required=True,
                        help="Directory containing train.csv / val.csv / test.csv")
    parser.add_argument("--out-dir", type=Path, required=True,
                        help="Where to write {train,val,test}.npz")
    parser.add_argument("--data-root", type=Path, required=True,
                        help="Root that the manifest 'path' column is relative to")
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--model", default="openai/clip-vit-base-patch32",
                        help="HF model name. MUST match the inference config.")
    parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
    parser.add_argument("--force", action="store_true",
                        help="Recompute even if output .npz / shards already exist")
    parser.add_argument(
        "--train-augment-manifest",
        type=Path,
        action="append",
        default=[],
        help=(
            "Augmented manifest to append to train.csv only. Can be passed "
            "multiple times."
        ),
    )
    parser.add_argument(
        "--extra-split",
        action="append",
        default=[],
        metavar="NAME=CSV",
        help=(
            "Additional explicit split to encode, for example "
            "test_augmented=data/test_augmented.csv. Can be passed multiple times."
        ),
    )
    args = parser.parse_args()

    args.out_dir.mkdir(parents=True, exist_ok=True)

    print(f"Loading CLIP model {args.model} on {args.device}...")
    import torch
    from PIL import Image
    from transformers import CLIPModel, CLIPProcessor

    processor = CLIPProcessor.from_pretrained(args.model)
    model = CLIPModel.from_pretrained(args.model).to(args.device)
    model.eval()
    for p in model.parameters():
        p.requires_grad = False

    extra_splits: list[tuple[str, Path]] = []
    for spec in args.extra_split:
        if "=" not in spec:
            raise ValueError(f"--extra-split must be NAME=CSV, got {spec!r}")
        name, csv_path = spec.split("=", 1)
        if not name or "/" in name or "\\" in name:
            raise ValueError(f"Invalid extra split name: {name!r}")
        extra_splits.append((name, Path(csv_path)))

    split_specs = [
        ("train", args.splits_dir / "train.csv", args.train_augment_manifest),
        ("val", args.splits_dir / "val.csv", []),
        ("test", args.splits_dir / "test.csv", []),
        *[(name, path, []) for name, path in extra_splits],
    ]

    for split, split_csv, extra_manifests in split_specs:
        _process_split(
            split=split,
            split_csv=split_csv,
            out_dir=args.out_dir,
            data_root=args.data_root,
            batch_size=args.batch_size,
            force=args.force,
            extra_manifests=extra_manifests,
            model=model,
            processor=processor,
            device=args.device,
            torch_module=torch,
            Image=Image,
        )

    print("\nDone.")


if __name__ == "__main__":
    main()