#!/usr/bin/env python3
"""Build a self-contained dataset_info.json from 3Deditverse_data.

The generated info is compatible with the existing 3DEditFormer editing
loader, while the sample set is decided only by files present in the clean
encoded/condition-image tree.
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import Dict, Iterable, Optional, Tuple

import numpy as np


IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--clean-root",
        type=Path,
        default=Path("/mnt/zsn/zsn_workspace/3DEditFormer/3DEditVerse/3Deditverse_data"),
        help="Path to 3Deditverse_data.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="Output dataset_info.json path. Defaults to <clean-root>/dataset_info.json.",
    )
    parser.add_argument(
        "--indent",
        type=int,
        default=2,
        help="JSON indentation. Use 0 for compact JSON.",
    )
    parser.add_argument(
        "--voxel-count-mode",
        choices=("constant", "slat"),
        default="constant",
        help=(
            "How to fill ori_voxel_num/edit_voxel_num. "
            "constant is fast and sufficient for non-sparse SS training; "
            "slat opens every slat_latent.npz and can be slow."
        ),
    )
    return parser.parse_args()


def iter_dir_names(path: Path) -> Iterable[str]:
    if not path.is_dir():
        return []
    with os.scandir(path) as entries:
        return sorted(entry.name for entry in entries if entry.is_dir(follow_symlinks=False))


def first_existing(paths: Iterable[Path]) -> Optional[Path]:
    for path in paths:
        if path.is_file():
            return path
    return None


def first_extra_image(path: Path, exclude: set[str]) -> Optional[Path]:
    if not path.is_dir():
        return None
    with os.scandir(path) as entries:
        candidates = sorted(
            Path(entry.path)
            for entry in entries
            if entry.is_file(follow_symlinks=False)
            and Path(entry.name).suffix.lower() in IMAGE_EXTS
            and entry.name not in exclude
        )
    return candidates[0] if candidates else None


def clean_condition_pair(clean_root: Path, branch: str, key: str) -> Optional[Tuple[Path, Path]]:
    cond_dir = clean_root / "condition-image" / branch / key
    ori = first_existing([cond_dir / "ori_image.png"])
    if branch == "alpaca":
        edit = first_existing([cond_dir / "after_edited_Flux.png"])
    else:
        edit = first_existing([cond_dir / "edited_0.png"])
        if edit is None:
            edit = first_extra_image(cond_dir, {"ori_image.png"})
    if ori is None or edit is None:
        return None
    return ori, edit


def latent_pair(encoded_dir: Path, key: str, name: str) -> Optional[Tuple[Path, Path]]:
    ori = encoded_dir / key / "ori" / name
    edit = encoded_dir / key / "edit" / name
    if ori.is_file() and edit.is_file():
        return ori, edit
    return None


def voxel_count(slat_path: Path) -> int:
    with np.load(slat_path) as data:
        if "coords" in data:
            return int(data["coords"].shape[0])
        if "slat_coords" in data:
            return int(data["slat_coords"].shape[0])
    raise KeyError(f"Expected coords or slat_coords in {slat_path}")


def rel(path: Path, clean_root: Path) -> str:
    return path.relative_to(clean_root).as_posix()


def build_branch(clean_root: Path, branch: str, encoded_name: str, voxel_count_mode: str) -> Tuple[Dict[str, dict], Dict[str, int]]:
    encoded_dir = clean_root / "encoded_ouput" / encoded_name
    condition_dir = clean_root / "condition-image" / branch
    encoded_keys = set(iter_dir_names(encoded_dir))
    condition_keys = set(iter_dir_names(condition_dir))

    records: Dict[str, dict] = {}
    stats = {
        "encoded_keys": len(encoded_keys),
        "condition_keys": len(condition_keys),
        "missing_condition": 0,
        "missing_ss_latent": 0,
        "missing_slat_latent": 0,
        "voxel_read_errors": 0,
    }

    for key in sorted(encoded_keys & condition_keys):
        cond_pair = clean_condition_pair(clean_root, branch, key)
        if cond_pair is None:
            stats["missing_condition"] += 1
            continue

        ss_pair = latent_pair(encoded_dir, key, "ss_latent.npz")
        if ss_pair is None:
            stats["missing_ss_latent"] += 1
            continue

        slat_pair = latent_pair(encoded_dir, key, "slat_latent.npz")
        if slat_pair is None:
            stats["missing_slat_latent"] += 1
            continue

        if voxel_count_mode == "slat":
            try:
                ori_voxel_num = voxel_count(slat_pair[0])
                edit_voxel_num = voxel_count(slat_pair[1])
            except Exception as exc:  # keep bad files out of training info
                stats["voxel_read_errors"] += 1
                print(f"[WARN] skip {branch}/{key}: failed to read voxel count: {exc}")
                continue
        else:
            ori_voxel_num = 1
            edit_voxel_num = 1

        ori_img, edit_img = cond_pair
        records[key] = {
            "ori_ss_latents_path": rel(ss_pair[0], clean_root),
            "edit_ss_latents_path": rel(ss_pair[1], clean_root),
            "ori_latents_path": rel(slat_pair[0], clean_root),
            "edit_latents_path": rel(slat_pair[1], clean_root),
            "ori_img_path": rel(ori_img, clean_root),
            "edit_img_path": rel(edit_img, clean_root),
            "ori_voxel_num": ori_voxel_num,
            "edit_voxel_num": edit_voxel_num,
        }

    stats["records"] = len(records)
    stats["encoded_without_condition_key"] = len(encoded_keys - condition_keys)
    stats["condition_without_encoded_key"] = len(condition_keys - encoded_keys)
    return records, stats


def main() -> None:
    args = parse_args()
    clean_root = args.clean_root.resolve()
    output = args.output.resolve() if args.output is not None else clean_root / "dataset_info.json"

    if not clean_root.is_dir():
        raise FileNotFoundError(f"clean root not found: {clean_root}")

    flux, flux_stats = build_branch(clean_root, "flux", "flux_encode", args.voxel_count_mode)
    alpaca, alpaca_stats = build_branch(clean_root, "alpaca", "alpaca_encode", args.voxel_count_mode)

    info = {
        "mixamo": {},
        "flux_edit": flux,
        "alpaca": alpaca,
        "_meta": {
            "clean_root": str(clean_root),
            "source": "generated_from_3Deditverse_data",
            "voxel_count_mode": args.voxel_count_mode,
            "branches": {
                "flux_edit": flux_stats,
                "alpaca": alpaca_stats,
            },
            "total_editverse_records": len(flux) + len(alpaca),
        },
    }

    output.parent.mkdir(parents=True, exist_ok=True)
    indent = None if args.indent == 0 else args.indent
    output.write_text(json.dumps(info, indent=indent, sort_keys=True) + "\n")

    print(f"Wrote: {output}")
    print(f"Flux records: {len(flux)}")
    print(f"Alpaca records: {len(alpaca)}")
    print(f"Total records: {len(flux) + len(alpaca)}")
    print("Stats:")
    print(json.dumps(info["_meta"]["branches"], indent=2, sort_keys=True))


if __name__ == "__main__":
    main()