# app/cache_features.py
import os
import csv
import argparse
from typing import Dict, Tuple, List

import torch
from PIL import Image

from .config import CFG, AppConfig
from .preprocessing import Tiler
from .features import PatchFeatureExtractor
from .utils import load_image, unzip_images, ensure_dir


def _tiles_from_image(path: str, cfg: AppConfig):
    img = load_image(path)  # PIL RGB
    tiler = Tiler(cfg.tiles)
    records = tiler.tile_image(img)
    boxes = [(int(r["y0"]), int(r["y1"]), int(r["x0"]), int(r["x1"])) for r in records]
    tiles = [r["image"] for r in records]
    return tiles, boxes


def _tiles_from_zip(path: str):
    with open(path, "rb") as fh:
        blob = fh.read()
    pairs = unzip_images(blob)  # [(name, PIL)]
    tiles = [img for _, img in pairs]
    # synth coords for visualization; not used in training
    boxes = [(0, 224, 0, 224)] * len(tiles)
    return tiles, boxes


def cache_from_csv(
    csv_path: str,
    out_dir: str = "features",
    cfg: AppConfig = CFG,
    limit: int = None,
) -> None:
    """
    CSV schema (no header):
        <path_to_image_or_zip>,<label_int>
    Example:
        demo_data/slides/slideA.png,0
        demo_data/patches_zip/caseB.zip,1
    """
    ensure_dir(out_dir)
    fx = PatchFeatureExtractor(cfg.feat)

    with open(csv_path, "r", newline="") as f:
        reader = csv.reader(f)
        for i, row in enumerate(reader):
            if limit is not None and i >= limit:
                break
            if not row or len(row) < 2:
                continue
            src, lab = row[0].strip(), int(row[1].strip())
            if not os.path.exists(src):
                print(f"[skip] missing: {src}")
                continue

            stem = os.path.splitext(os.path.basename(src))[0]
            save_path = os.path.join(out_dir, f"{stem}.pt")
            if os.path.exists(save_path):
                print(f"[skip] exists: {save_path}")
                continue

            if src.lower().endswith(".zip"):
                tiles, boxes = _tiles_from_zip(src)
            else:
                tiles, boxes = _tiles_from_image(src, cfg)

            if len(tiles) == 0:
                print(f"[warn] no tiles retained after tissue filter: {src}")
                continue

            feats = fx.encode(tiles, batch_size=cfg.feat.batch_size)  # [N,D]
            blob = {
                "feats": feats,          # FloatTensor [N,D] on CPU
                "label": lab,            # int
                "boxes": boxes,          # for heatmaps if image-based
            }
            torch.save(blob, save_path)
            print(f"[ok] {src} → {save_path}  (tiles={len(tiles)}, dim={feats.shape[-1]})")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", type=str, required=True, help="CSV file with <path,label>")
    ap.add_argument("--out_dir", type=str, default="features")
    ap.add_argument("--limit", type=int, default=None)
    args = ap.parse_args()

    CFG.sync_dims()
    cache_from_csv(args.csv, args.out_dir, CFG, args.limit)