#!/usr/bin/env python3
"""
Dump j63_recipe data paths for migration.

This script:
1. Reads VILA registry configs (cs-oci-ord.yaml, mixtures.yaml)
2. Parses j63_recipe mixture to get all dataset names
3. Collects all data_path and media_dir for each dataset
4. Deduplicates and merges parent directories
5. Outputs minimal transfer set with rsync commands

Usage:
    python chengyue_script/dvlm/utils/dump_j63_paths.py
    python chengyue_script/dvlm/utils/dump_j63_paths.py --mixture j63_math
    python chengyue_script/dvlm/utils/dump_j63_paths.py --check-size
"""

import os
import sys
import argparse
import subprocess
from pathlib import Path
from datetime import datetime
from collections import defaultdict

try:
    import yaml
except ImportError:
    print("ERROR: PyYAML not installed. Run: pip install pyyaml")
    sys.exit(1)

# Project paths
SCRIPT_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = SCRIPT_DIR.parents[3]  # LMFlow root
VILA_ROOT = PROJECT_ROOT / "VILA"
REGISTRY_DIR = VILA_ROOT / "llava" / "data" / "registry"


def load_yaml(path: Path) -> dict:
    """Load a YAML file."""
    with open(path, "r") as f:
        return yaml.safe_load(f) or {}


def load_mixtures() -> dict:
    """Load mixtures.yaml."""
    mixtures_path = REGISTRY_DIR / "mixtures.yaml"
    if not mixtures_path.exists():
        print(f"ERROR: mixtures.yaml not found at {mixtures_path}")
        sys.exit(1)
    return load_yaml(mixtures_path)


def load_datasets(env: str = "cs-oci-ord") -> dict:
    """Load datasets registry for the specified environment."""
    datasets_path = REGISTRY_DIR / "datasets" / f"{env}.yaml"
    if not datasets_path.exists():
        print(f"ERROR: {env}.yaml not found at {datasets_path}")
        sys.exit(1)
    return load_yaml(datasets_path)


def parse_mixture(mixture_name: str, mixtures: dict) -> list:
    """
    Parse a mixture name into a list of dataset names.
    Handles nested mixtures and + syntax.
    """
    if "+" in mixture_name:
        names = mixture_name.split("+")
    else:
        names = [mixture_name]
    
    # Expand nested mixtures
    result = []
    for name in names:
        # Handle @ratio suffix (e.g., sharegpt4v_sft@0.3)
        base_name = name.split("@")[0]
        
        if base_name in mixtures:
            # It's a mixture, expand it
            expanded = parse_mixture("+".join(mixtures[base_name]), mixtures)
            result.extend(expanded)
        else:
            # It's a dataset
            result.append(base_name)
    
    return sorted(set(result))


def get_dataset_paths(dataset_name: str, datasets: dict) -> dict:
    """
    Get data_path and media_dir for a dataset.
    Returns dict with 'data_path' and 'media_dir' keys.
    """
    if dataset_name not in datasets:
        return {"data_path": None, "media_dir": None, "error": f"Dataset '{dataset_name}' not found in registry"}
    
    config = datasets[dataset_name]
    
    # Handle different config key names
    data_path = config.get("data_path")
    media_dir = config.get("media_dir") or config.get("image_path") or config.get("video_dir") or config.get("image_folder")
    
    return {
        "data_path": data_path,
        "media_dir": media_dir,
        "config": config
    }


def get_dir_size(path: str) -> str:
    """Get directory size using du command. Returns human-readable string."""
    if not os.path.exists(path):
        return "N/A"
    
    try:
        result = subprocess.run(
            ["du", "-sh", path],
            capture_output=True,
            text=True,
            timeout=60
        )
        if result.returncode == 0:
            return result.stdout.split()[0]
    except (subprocess.TimeoutExpired, Exception):
        pass
    
    return "?"


def find_common_parents(paths: list) -> list:
    """
    Find minimal set of parent directories that cover all paths.
    If path A is a parent of path B, only keep A.
    """
    if not paths:
        return []
    
    # Normalize and sort paths
    normalized = []
    for p in paths:
        if p:
            # Ensure trailing slash for directory comparison
            norm = os.path.normpath(p)
            normalized.append(norm)
    
    normalized = sorted(set(normalized))
    
    # Remove paths that are children of other paths
    result = []
    for path in normalized:
        is_child = False
        for other in normalized:
            if path != other and path.startswith(other + "/"):
                is_child = True
                break
        if not is_child:
            result.append(path)
    
    return result


def identify_shared_dirs(dataset_paths: dict) -> dict:
    """
    Identify which directories are shared by multiple datasets.
    Returns dict: dir_path -> list of dataset names
    """
    dir_to_datasets = defaultdict(list)
    
    for ds_name, paths in dataset_paths.items():
        if paths.get("media_dir"):
            dir_to_datasets[paths["media_dir"]].append(ds_name)
    
    # Only return dirs shared by multiple datasets
    return {k: v for k, v in dir_to_datasets.items() if len(v) > 1}


def main():
    parser = argparse.ArgumentParser(description="Dump j63_recipe data paths for migration")
    parser.add_argument(
        "--mixture",
        type=str,
        default="j63_recipe",
        help="Mixture name to analyze (default: j63_recipe)"
    )
    parser.add_argument(
        "--env",
        type=str,
        default="cs-oci-ord",
        help="Dataset environment/registry (default: cs-oci-ord)"
    )
    parser.add_argument(
        "--check-size",
        action="store_true",
        help="Check directory sizes (slow, uses du)"
    )
    parser.add_argument(
        "--dest",
        type=str,
        default="$SHARE_OUTPUT/data",
        help="Destination path for rsync commands"
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Output file path (default: stdout)"
    )
    args = parser.parse_args()
    
    # Load configs
    print(f"Loading VILA registry from: {REGISTRY_DIR}", file=sys.stderr)
    mixtures = load_mixtures()
    datasets = load_datasets(args.env)
    
    # Parse mixture
    dataset_names = parse_mixture(args.mixture, mixtures)
    print(f"Mixture '{args.mixture}' contains {len(dataset_names)} datasets", file=sys.stderr)
    
    # Collect paths for each dataset
    dataset_paths = {}
    all_data_paths = []
    all_media_dirs = []
    missing_datasets = []
    
    for ds_name in dataset_names:
        paths = get_dataset_paths(ds_name, datasets)
        dataset_paths[ds_name] = paths
        
        if paths.get("error"):
            missing_datasets.append(ds_name)
        else:
            if paths.get("data_path"):
                all_data_paths.append(paths["data_path"])
            if paths.get("media_dir"):
                all_media_dirs.append(paths["media_dir"])
    
    # Find shared directories
    shared_dirs = identify_shared_dirs(dataset_paths)
    
    # Combine all paths and find minimal set
    all_paths = all_data_paths + all_media_dirs
    
    # Get parent directories of data_path files
    data_path_dirs = [os.path.dirname(p) for p in all_data_paths if p]
    
    # Combine media_dirs and data_path parent dirs
    all_dirs = list(set(all_media_dirs + data_path_dirs))
    
    # Find minimal transfer set
    minimal_dirs = find_common_parents(all_dirs)
    
    # Generate output
    output_lines = []
    output_lines.append("=" * 80)
    output_lines.append(f"{args.mixture} 数据迁移路径报告")
    output_lines.append(f"生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    output_lines.append(f"环境: {args.env}")
    output_lines.append("=" * 80)
    output_lines.append("")
    
    # Minimal transfer directories
    output_lines.append("【最小传输目录】(去重合并后，直接用于 rsync)")
    output_lines.append("-" * 80)
    
    for i, dir_path in enumerate(minimal_dirs, 1):
        exists = "✓" if os.path.exists(dir_path) else "✗"
        size = get_dir_size(dir_path) if args.check_size else ""
        size_str = f" ({size})" if size else ""
        output_lines.append(f"{i}. [{exists}] {dir_path}{size_str}")
    
    output_lines.append("")
    
    # rsync commands
    output_lines.append("【rsync 命令】")
    output_lines.append("-" * 80)
    for dir_path in minimal_dirs:
        # Get relative path from common prefix
        rel_path = Path(dir_path).name
        output_lines.append(f"rsync -avP --progress {dir_path}/ {args.dest}/{rel_path}/")
    output_lines.append("")
    
    # Shared directories warning
    if shared_dirs:
        output_lines.append("【共享目录】(多个数据集共用)")
        output_lines.append("-" * 80)
        for dir_path, ds_names in shared_dirs.items():
            output_lines.append(f"  {dir_path}")
            output_lines.append(f"    -> 被以下数据集共用: {', '.join(ds_names)}")
        output_lines.append("")
    
    # Missing datasets
    if missing_datasets:
        output_lines.append("【警告: 未找到的数据集】")
        output_lines.append("-" * 80)
        for ds_name in missing_datasets:
            output_lines.append(f"  - {ds_name}")
        output_lines.append("")
    
    # Detailed per-dataset breakdown
    output_lines.append("【详细清单】(per-dataset)")
    output_lines.append("-" * 80)
    
    for i, ds_name in enumerate(dataset_names, 1):
        paths = dataset_paths[ds_name]
        output_lines.append(f"{i}. {ds_name}")
        
        if paths.get("error"):
            output_lines.append(f"   ⚠️  {paths['error']}")
        else:
            # data_path
            dp = paths.get("data_path")
            if dp:
                exists = "✓" if os.path.exists(dp) else "✗"
                output_lines.append(f"   data_path: [{exists}] {dp}")
            else:
                output_lines.append(f"   data_path: (none)")
            
            # media_dir
            md = paths.get("media_dir")
            if md:
                exists = "✓" if os.path.exists(md) else "✗"
                shared_note = ""
                if md in shared_dirs:
                    other_ds = [n for n in shared_dirs[md] if n != ds_name]
                    if other_ds:
                        shared_note = f" [共享: {', '.join(other_ds)}]"
                output_lines.append(f"   media_dir: [{exists}] {md}{shared_note}")
            else:
                output_lines.append(f"   media_dir: (none)")
        
        output_lines.append("")
    
    # Raw path list for scripting
    output_lines.append("【原始路径列表】(用于脚本处理)")
    output_lines.append("-" * 80)
    output_lines.append("# 最小传输目录:")
    for dir_path in minimal_dirs:
        output_lines.append(dir_path)
    output_lines.append("")
    output_lines.append("# 所有 data_path 文件:")
    for dp in sorted(set(all_data_paths)):
        if dp:
            output_lines.append(dp)
    output_lines.append("")
    output_lines.append("# 所有 media_dir 目录:")
    for md in sorted(set(all_media_dirs)):
        if md:
            output_lines.append(md)
    
    # Output
    output_text = "\n".join(output_lines)
    
    if args.output:
        with open(args.output, "w") as f:
            f.write(output_text)
        print(f"Report saved to: {args.output}", file=sys.stderr)
    else:
        print(output_text)
    
    # Summary to stderr
    print(f"\nSummary:", file=sys.stderr)
    print(f"  Datasets: {len(dataset_names)}", file=sys.stderr)
    print(f"  Missing: {len(missing_datasets)}", file=sys.stderr)
    print(f"  Unique data_paths: {len(set(all_data_paths))}", file=sys.stderr)
    print(f"  Unique media_dirs: {len(set(all_media_dirs))}", file=sys.stderr)
    print(f"  Minimal transfer dirs: {len(minimal_dirs)}", file=sys.stderr)


if __name__ == "__main__":
    main()

