"""Build a Food-101 subset for training.

After ``scripts/download_data.py --cv`` extracted the 10 target classes into
``data/raw/food101_subset/food-101/images/<class>/<image>.jpg``, this script
splits them into train/val/test directories using Food-101's own meta files.

Produces:
    data/processed/cv/train/<class>/*.jpg
    data/processed/cv/val/<class>/*.jpg
    data/processed/cv/test/<class>/*.jpg

Usage:
    python -m src.cv.prepare_data
"""

from __future__ import annotations

import json
import shutil
import sys
from pathlib import Path
from random import Random

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from src.config import CV_CLASSES_PATH, CV_TARGET_CLASSES, PROCESSED_DIR, RAW_DIR  # noqa: E402


CV_RAW = RAW_DIR / "food101_subset" / "food-101"
CV_OUT = PROCESSED_DIR / "cv"


def _read_meta(path: Path) -> list[str]:
    return [line.strip() for line in path.read_text().splitlines() if line.strip()]


def _link_image(src: Path, dst: Path) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    if dst.exists():
        return
    try:
        dst.symlink_to(src)
    except OSError:
        shutil.copy2(src, dst)


def main(val_fraction: float = 0.15, seed: int = 42) -> None:
    train_meta = CV_RAW / "meta" / "train.txt"
    test_meta = CV_RAW / "meta" / "test.txt"
    if not train_meta.exists() or not test_meta.exists():
        raise FileNotFoundError(
            "Food-101 meta files missing. Run 'python scripts/download_data.py --cv'."
        )

    targets = set(CV_TARGET_CLASSES)
    rnd = Random(seed)

    train_entries = [e for e in _read_meta(train_meta) if e.split("/")[0] in targets]
    test_entries = [e for e in _read_meta(test_meta) if e.split("/")[0] in targets]

    val_entries: list[str] = []
    new_train: list[str] = []
    by_class: dict[str, list[str]] = {}
    for entry in train_entries:
        by_class.setdefault(entry.split("/")[0], []).append(entry)
    for cls, items in by_class.items():
        rnd.shuffle(items)
        cut = int(len(items) * val_fraction)
        val_entries.extend(items[:cut])
        new_train.extend(items[cut:])

    splits = {"train": new_train, "val": val_entries, "test": test_entries}
    for split, entries in splits.items():
        for entry in entries:
            src = CV_RAW / "images" / f"{entry}.jpg"
            cls = entry.split("/")[0]
            dst = CV_OUT / split / cls / f"{entry.split('/')[1]}.jpg"
            if src.exists():
                _link_image(src, dst)
        print(f"[cv.prepare_data] {split:>5}: {len(entries):,} images")

    CV_CLASSES_PATH.parent.mkdir(parents=True, exist_ok=True)
    CV_CLASSES_PATH.write_text(json.dumps(sorted(CV_TARGET_CLASSES), indent=2))
    print(f"[cv.prepare_data] wrote class list to {CV_CLASSES_PATH}")


if __name__ == "__main__":
    main()