#!/usr/bin/env python
# coding=utf-8
"""WebDataset Multi Modal Dataset support for LMFlow.

This module provides support for loading multi-modal datasets from WebDataset format
(.tar files with .wdsmeta.json metadata).
"""
import json
import io
import os
import sys
import hashlib
import getpass
from pathlib import Path
from typing import Optional, List, Tuple, TYPE_CHECKING, Any
from PIL import Image
import torch
from torch.utils.data import Dataset

vila_root = Path(__file__).resolve().parents[2] / "VILA"
if vila_root.exists():
    if str(vila_root) not in sys.path:
        sys.path.append(str(vila_root))
    from llava.wids import ShardListDataset

if TYPE_CHECKING:
    from lmflow.args import DatasetArguments
else:
    DatasetArguments = Any


class CustomWebDatasetMultiModalDataset(Dataset):
    """Dataset for Multi Modal data from WebDataset format.
    
    WebDataset format:
    - Each sample has 3 files: .jpg (image), .json (metadata), .txt (text)
    - Files are grouped by sample ID (e.g., "2001454978075.jpg", "2001454978075.json", "2001454978075.txt")
    - Metadata is stored in .wdsmeta.json file with shard list
    
    This dataset:
    1. Reads .wdsmeta.json to get list of shard URLs
    2. Uses webdataset library to iterate over tar files
    3. Builds conversations from .txt file (simple caption format)
    4. Returns messages format compatible with CustomMultiModalDataset
    """

    def __init__(self, wdsmeta_path: str, data_args: DatasetArguments):
        super(CustomWebDatasetMultiModalDataset, self).__init__()
        self.data_args = data_args
        self.original_wdsmeta_path = Path(wdsmeta_path)
        self.wdsmeta_path = self._resolve_wdsmeta_path(self.original_wdsmeta_path)
        if not self.wdsmeta_path.exists():
            raise FileNotFoundError(
                f"WebDataset metadata file not found: {self.wdsmeta_path}"
            )
        if self.wdsmeta_path != self.original_wdsmeta_path:
            print(
                "[WebDataset] Remapped metadata path "
                f"{self.original_wdsmeta_path} -> {self.wdsmeta_path}",
                flush=True,
            )
        
        # Load metadata
        print(f"[WebDataset] Loading metadata from: {self.wdsmeta_path}", flush=True)
        with self.wdsmeta_path.open('r') as f:
            meta = json.load(f)
        
        # Debug: print metadata structure
        print(f"[WebDataset] Metadata keys: {list(meta.keys())}", flush=True)
        if 'shardlist' not in meta:
            raise ValueError(
                f"Missing 'shardlist' key in wdsmeta.json. "
                f"Available keys: {list(meta.keys())}"
            )
        self._setup_shard_path_resolution()
        
        # Rewrite shard URLs to align with local filesystem if needed.
        rewritten_shards = []
        for shard in meta["shardlist"]:
            url = shard.get("url")
            if url is None:
                continue
            rewritten_url = self._maybe_rewrite_shard_path(url)
            if rewritten_url != url:
                shard = dict(shard)
                shard["url"] = rewritten_url
            rewritten_shards.append(shard)
        meta["shardlist"] = rewritten_shards
        print(f"[WebDataset] Found {len(rewritten_shards)} total shard URLs", flush=True)

        # Instantiate the VILA ShardListDataset using the rewritten metadata.
        cache_root = Path(os.environ.get("LMFLOW_WDS_CACHE", "~/.cache/_wids_cache")).expanduser()
        cache_root.mkdir(parents=True, exist_ok=True)
        cache_uuid = hashlib.sha256(str(self.wdsmeta_path).encode()).hexdigest()[:8]
        cache_dir = cache_root / f"{getpass.getuser()}-{cache_uuid}"
        # In multi-node settings, ensure ranks don't collide on the same cache dir if they write to it
        # VILA's ShardListDataset writes a UUID-based cache, but just in case
        rank = os.environ.get("RANK", "0")
        if rank != "0":
             cache_dir = cache_dir / f"rank_{rank}"
        cache_dir.mkdir(parents=True, exist_ok=True)

        meta_buffer = io.StringIO()
        json.dump(meta, meta_buffer)
        meta_buffer.seek(0)
        self._dataset = ShardListDataset(meta_buffer, cache_dir=str(cache_dir))

        # Pre-compute shard ranges for fast index mapping.
        self._shard_lengths = list(self._dataset.lengths)
        self._shard_ranges: List[Tuple[int, int]] = []
        cursor = 0
        for length in self._shard_lengths:
            self._shard_ranges.append((cursor, cursor + length))
            cursor += length

        # Handle multi-node/multi-process training: split shards across nodes/ranks.
        rank = int(os.environ.get("LOCAL_RANK", os.environ.get("RANK", 0)))
        world_size = int(os.environ.get("WORLD_SIZE", 1))
        node_rank = int(os.environ.get("NODE_RANK", 0))
        num_nodes = int(os.environ.get("NNODES", 1))

        def _partition(total: int, parts: int, idx: int) -> Tuple[int, int]:
            if parts <= 1 or total == 0:
                return 0, total
            parts = max(parts, 1)
            idx = max(0, min(idx, parts - 1))
            base = total // parts
            remainder = total % parts
            start = idx * base + min(idx, remainder)
            end = start + base + (1 if idx < remainder else 0)
            return start, min(end, total)

        total_shards = len(self._dataset.shards)
        node_start, node_end = _partition(total_shards, num_nodes, node_rank)
        node_shard_indices = list(range(node_start, node_end))
        print(
            f"[WebDataset] Node {node_rank}/{num_nodes}: assigned shard indices {node_start}-{max(node_end-1, node_start-1)} "
            f"({len(node_shard_indices)} shards)",
            flush=True,
        )

        if world_size > 1 and len(node_shard_indices) > 0:
            rank_start, rank_end = _partition(len(node_shard_indices), world_size, rank)
            assigned_indices = node_shard_indices[rank_start:rank_end]
            print(
                f"[WebDataset] Rank {rank}/{world_size}: shard slots {rank_start}-{max(rank_end-1, rank_start-1)} "
                f"-> {len(assigned_indices)} shards",
                flush=True,
            )
        else:
            assigned_indices = node_shard_indices

        self.assigned_shard_indices = assigned_indices
        self.shard_urls = [self._dataset.shards[i]["url"] for i in assigned_indices]

        if self.shard_urls:
            print(f"[WebDataset] First shard URL for this rank: {self.shard_urls[0]}", flush=True)
        else:
            print(f"[WebDataset] WARNING: No shards assigned to rank {rank}!", flush=True)

        # Build flat sample index for random access.
        self.samples = []
        self.shard_indices = {}
        total_samples = 0
        for local_idx, global_shard_idx in enumerate(assigned_indices):
            start, end = self._shard_ranges[global_shard_idx]
            shard_url = self._dataset.shards[global_shard_idx]["url"]
            shard_sample_count = end - start
            if rank == 0 and local_idx == 0:
                print(
                    f"[WebDataset] Processing shard {local_idx} (global {global_shard_idx}): {shard_url} "
                    f"with {shard_sample_count} samples",
                    flush=True,
                )
            shard_entries = []
            for sample_offset, dataset_index in enumerate(range(start, end)):
                global_idx = len(self.samples)
                self.samples.append(
                    {
                        "shard_idx": global_shard_idx,
                        "local_shard_idx": local_idx,
                        "shard_url": shard_url,
                        "dataset_index": dataset_index,
                        "sample_idx_in_shard": sample_offset,
                    }
                )
                shard_entries.append(global_idx)
            if shard_entries:
                self.shard_indices[global_shard_idx] = shard_entries
            total_samples += shard_sample_count

        print(
            f"[WebDataset] Indexed {total_samples} samples from {len(assigned_indices)} shards for this rank.",
            flush=True,
        )

        if len(self.samples) == 0:
            import warnings

            warnings.warn(
                "WebDataset initialized with 0 samples! "
                "Check if shard URLs are accessible and contain valid data."
            )
    
    def _setup_shard_path_resolution(self):
        """Prepare helper structures for shard-path remapping."""
        self._rewrite_log_budget = 5
        self._missing_log_budget = 5
        self._meta_rewrite_log_budget = 3
        self._user_name = self._extract_user_name()
        self._user_dataset_root = self._default_user_dataset_root(self._user_name)
        self._original_dataset_root = self._decode_original_dataset_root()
        self._mirror_dataset_root = self._infer_local_dataset_root(self._original_dataset_root)
        self._path_rewrite_rules: List[Tuple[Path, Path]] = []
        default_rules = self._default_shard_path_rules()
        if default_rules:
            self._path_rewrite_rules.extend(default_rules)
        env_rules = self._load_env_path_rules()
        if env_rules:
            self._path_rewrite_rules.extend(env_rules)
        if self._original_dataset_root is not None and self._mirror_dataset_root is not None:
            self._path_rewrite_rules.append(
                (self._original_dataset_root, self._mirror_dataset_root)
            )
            print(
                "[WebDataset] Auto-mapped shard root "
                f"{self._original_dataset_root} -> {self._mirror_dataset_root}",
                flush=True,
            )
        self._additional_search_roots = self._build_additional_search_roots()
    
    def _resolve_wdsmeta_path(self, candidate: Path) -> Path:
        """Try to locate metadata file even if prefix differs."""
        candidate = Path(candidate)
        if candidate.exists():
            return candidate
        for rewritten in self._iter_meta_path_candidates(candidate):
            if rewritten.exists():
                self._log_meta_rewrite(candidate, rewritten)
                return rewritten
        return candidate
    
    def _iter_meta_path_candidates(self, candidate: Path):
        """Yield possible metadata path rewrites."""
        path_str = candidate.as_posix()
        for src, dst in self._default_meta_path_rules():
            if path_str.startswith(src):
                yield Path(path_str.replace(src, dst, 1))
        for src, dst in self._load_meta_env_rules():
            if path_str.startswith(src):
                yield Path(path_str.replace(src, dst, 1))
    
    @staticmethod
    def _default_meta_path_rules() -> List[Tuple[str, str]]:
        return [
            (
                "/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users",
                "/lustre/fsw/portfolios/nvr/users",
            ),
        ]
    
    def _load_meta_env_rules(self) -> List[Tuple[str, str]]:
        env_spec = os.environ.get("LMFLOW_WDS_META_PATH_MAP")
        rules: List[Tuple[str, str]] = []
        if not env_spec:
            return rules
        for pair in env_spec.split(","):
            pair = pair.strip()
            if not pair or "=" not in pair:
                continue
            src, dst = pair.split("=", 1)
            src = src.strip()
            dst = dst.strip()
            if not src or not dst:
                continue
            rules.append((src, dst))
        if rules:
            print(
                "[WebDataset] Loaded metadata path rewrites from LMFLOW_WDS_META_PATH_MAP.",
                flush=True,
            )
        return rules
    
    def _log_meta_rewrite(self, original: Path, rewritten: Path):
        if getattr(self, "_meta_rewrite_log_budget", 0) > 0:
            print(
                f"[WebDataset] Metadata path rewrite: {original} -> {rewritten}",
                flush=True,
            )
            self._meta_rewrite_log_budget -= 1
    
    def _decode_original_dataset_root(self) -> Optional[Path]:
        """Infer the original dataset root path from the metadata filename."""
        meta_name = self.wdsmeta_path.name
        suffix = ".wdsmeta.json"
        if not (meta_name.startswith("--") and meta_name.endswith(suffix)):
            return None
        encoded = meta_name[: -len(suffix)]
        parts = [segment for segment in encoded.split("--") if segment]
        if not parts:
            return None
        decoded = Path("/" + "/".join(parts))
        return decoded
    
    def _infer_local_dataset_root(self, original_root: Optional[Path]) -> Optional[Path]:
        """Try to locate the dataset root on the current machine."""
        if original_root is None:
            return None
        dataset_name = original_root.name
        if not dataset_name:
            return None
        candidates: List[Path] = []
        env_mirror = os.environ.get("LMFLOW_WDS_MIRROR_ROOT")
        if env_mirror:
            candidates.append(Path(env_mirror))
        meta_parent = self.wdsmeta_path.parent
        candidates.extend([
            meta_parent,
            meta_parent.parent,
            meta_parent.parent.parent if meta_parent.parent is not None else None,
        ])
        if self._user_dataset_root is not None:
            candidates.append(self._user_dataset_root)
        unique_candidates = []
        for cand in candidates:
            if cand is None:
                continue
            cand = Path(cand)
            if cand not in unique_candidates:
                unique_candidates.append(cand)
        for cand in unique_candidates:
            if cand.name == dataset_name and cand.exists():
                return cand
            candidate_path = cand / dataset_name
            if candidate_path.exists():
                return candidate_path
        return None
    
    def _extract_user_name(self) -> Optional[str]:
        parts = self.wdsmeta_path.parts
        for idx, part in enumerate(parts):
            if part == "users" and idx + 1 < len(parts):
                return parts[idx + 1]
        return None
    
    def _default_user_dataset_root(self, user_name: Optional[str]) -> Optional[Path]:
        if not user_name:
            return None
        return Path("/lustre/fsw/portfolios/nvr/users") / user_name / "datasets"
    
    def _default_shard_path_rules(self) -> List[Tuple[Path, Path]]:
        rules: List[Tuple[Path, Path]] = []
        if self._original_dataset_root is not None and self._user_dataset_root is not None:
            rules.append((self._original_dataset_root, self._user_dataset_root))
        return rules
    
    def _load_env_path_rules(self) -> List[Tuple[Path, Path]]:
        """Load explicit rewrite rules from LMFLOW_WDS_PATH_MAP."""
        rules: List[Tuple[Path, Path]] = []
        env_spec = os.environ.get("LMFLOW_WDS_PATH_MAP")
        if not env_spec:
            return rules
        for pair in env_spec.split(","):
            pair = pair.strip()
            if not pair or "=" not in pair:
                continue
            src, dst = pair.split("=", 1)
            src = src.strip()
            dst = dst.strip()
            if not src or not dst:
                continue
            rules.append((Path(src), Path(dst)))
        if rules:
            print(
                "[WebDataset] Loaded shard path map from LMFLOW_WDS_PATH_MAP.",
                flush=True,
            )
        return rules
    
    def _build_additional_search_roots(self) -> List[Path]:
        """Prepare additional directories to probe for missing shards."""
        roots: List[Path] = []
        if self._user_dataset_root is not None:
            roots.append(self._user_dataset_root)
        meta_parent = self.wdsmeta_path.parent
        roots.append(meta_parent)
        roots.append(meta_parent.parent)
        env_extra = os.environ.get("LMFLOW_WDS_ADDITIONAL_ROOTS")
        if env_extra:
            for part in env_extra.split(":"):
                part = part.strip()
                if part:
                    roots.append(Path(part))
        deduped: List[Path] = []
        for root in roots:
            if root is None:
                continue
            root = Path(root)
            if root not in deduped:
                deduped.append(root)
        return deduped
    
    def _maybe_rewrite_shard_path(self, shard_url: str) -> str:
        """Rewrite shard paths when the recorded location is not accessible."""
        if not isinstance(shard_url, str):
            return shard_url
        if self._looks_like_remote_path(shard_url):
            return shard_url
        shard_path = Path(shard_url)
        if shard_path.exists():
            return shard_path.as_posix()
        rewritten = self._rewrite_user_home_path(shard_path)
        if rewritten is not None and rewritten.exists():
            self._log_rewrite(shard_path, rewritten)
            return rewritten.as_posix()
        # Apply explicit rewrite rules first.
        for src_root, dst_root in self._path_rewrite_rules:
            try:
                rel = shard_path.relative_to(src_root)
            except ValueError:
                continue
            candidate = dst_root / rel
            if candidate.exists():
                self._log_rewrite(shard_path, candidate)
                return candidate.as_posix()
        # Fallback: try additional search roots.
        dataset_dir = shard_path.parent.name
        for root in self._additional_search_roots:
            candidates = [root / shard_path.name]
            if dataset_dir:
                candidates.append(root / dataset_dir / shard_path.name)
            for candidate in candidates:
                if candidate.exists():
                    self._log_rewrite(shard_path, candidate)
                    return candidate.as_posix()
        self._log_missing_shard(shard_path)
        return shard_url
    
    def _rewrite_user_home_path(self, shard_path: Path) -> Optional[Path]:
        if not self._user_name or self._user_dataset_root is None:
            return None
        home_prefix = Path("/home") / self._user_name / "nvr_elm_llm" / "dataset"
        try:
            rel = shard_path.relative_to(home_prefix)
        except ValueError:
            return None
        return self._user_dataset_root / rel
    
    @staticmethod
    def _looks_like_remote_path(path_str: str) -> bool:
        remote_prefixes = ("s3://", "gs://", "http://", "https://", "pipe:")
        return path_str.startswith(remote_prefixes)
    
    def _log_rewrite(self, original: Path, rewritten: Path):
        if self._rewrite_log_budget > 0:
            print(
                f"[WebDataset] Remapped shard path: {original} -> {rewritten}",
                flush=True,
            )
            self._rewrite_log_budget -= 1
    
    def _log_missing_shard(self, shard_path: Path):
        if self._missing_log_budget > 0:
            print(
                f"[WebDataset] WARNING: Unable to locate shard file {shard_path}.",
                flush=True,
            )
            self._missing_log_budget -= 1
    
    def __len__(self):
        return len(self.samples)
    
    def register_tokenizer(self, tokenizer, image_processor=None):
        """Register tokenizer and image processor (for compatibility with CustomMultiModalDataset)."""
        self.tokenizer = tokenizer
        self.image_processor = getattr(
            tokenizer, "image_processor", image_processor)
    
    def __getitem__(self, idx):
        """Get a sample by index.
        
        Returns:
            dict: {"messages": [...]} format compatible with CustomMultiModalDataset
        """
        sample_info = self.samples[idx]
        dataset_index = sample_info["dataset_index"]
        
        # Get sample from global dataset index
        sample = self._dataset[dataset_index]
        if sample is None:
            raise ValueError(
                f"Sample at dataset index {dataset_index} not found "
                f"(shard {sample_info.get('shard_idx')} url {sample_info.get('shard_url')})"
            )
        
        # Extract image
        image_data = None
        for key in (".jpg", ".jpeg", ".png", ".webp", ".bmp"):
            if key in sample:
                image_data = sample[key]
                break
            alt_key = key.lstrip(".")
            if alt_key in sample:
                image_data = sample[alt_key]
                break
        if image_data is None:
            raise ValueError(f"Sample at index {idx} has no image")
        
        # Convert bytes to PIL Image
        if isinstance(image_data, bytes):
            image = Image.open(io.BytesIO(image_data)).convert('RGB')
        elif isinstance(image_data, Image.Image):
            image = image_data.convert('RGB')
        else:
            raise ValueError(f"Unexpected image type: {type(image_data)}")
        
        # Extract text caption
        txt_data = None
        for key in (".txt", "txt", ".text", "text"):
            if key in sample:
                txt_data = sample[key]
                break
        if txt_data is None:
            raise ValueError(f"Sample at index {idx} has no text")
        
        if isinstance(txt_data, bytes):
            caption = txt_data.decode("utf-8", errors="replace").strip()
        elif isinstance(txt_data, str):
            caption = txt_data.strip()
        else:
            caption = str(txt_data).strip()
        
        # Build conversations format
        # Simple format: user asks about the image, assistant responds with caption
        # This matches the format expected by CustomMultiModalDataset
        conversations = [
            {
                'from': 'human',
                'value': 'Describe this image.'
            },
            {
                'from': 'gpt',
                'value': caption
            }
        ]
        
        # Convert to messages format (same as CustomMultiModalDataset)
        messages = []
        image_added = False
        for message in conversations:
            if message['from'] == 'human':
                if not image_added:
                    messages.append({
                        "role": "user",
                        "content": [
                            {"type": "text", "text": message['value']},
                            {"type": "image", "image": image},  # PIL Image object
                        ],
                    })
                    image_added = True
                else:
                    messages.append({
                        "role": "user",
                        "content": [
                            {"type": "text", "text": message['value']},
                        ],
                    })
            else:
                messages.append({
                    "role": "assistant",
                    "content": [
                        {"type": "text", "text": message['value']},
                    ],
                })
        
        return {"messages": messages}
