#!/usr/bin/env python
# coding=utf-8
"""
VILA Dataset Adapter for LMFlow

Wraps VILA datasets to convert their output format to LMFlow's expected format.
VILA datasets return tokenized data, but LMFlow needs raw messages for custom collators.
"""

from torch.utils.data import Dataset
from typing import Any, Dict
import os
import io

        
# Get the VILA dataset item
# For ConcatDataset, we need to find which sub-dataset this index belongs to
from torch.utils.data import ConcatDataset

class VILADatasetAdapter(Dataset):
    """
    Adapter that wraps VILA datasets to return LMFlow-compatible format.
    
    VILA datasets return: {"input_ids": ..., "labels": ..., "image": ...}
    LMFlow expects: {"messages": [...]}
    
    This adapter extracts the raw conversation data from VILA datasets
    before tokenization happens.
    """
    
    def __init__(self, vila_dataset: Dataset):
        self.vila_dataset = vila_dataset
        
    def __len__(self):
        return len(self.vila_dataset)
    
    def __getitem__(self, idx) -> Dict[str, Any]:
        """
        Get item from VILA dataset and convert to LMFlow format.
        
        Strategy:
        1. Access the underlying VILA dataset's raw data (before tokenization)
        2. Extract conversations and images
        3. Convert to LMFlow's messages format
        """
        # Debug logging (only for first few samples and every 1000th)
        if idx < 5 or idx % 1000 == 0:
            rank = os.environ.get("RANK", "0")
            print(f"[Rank {rank}] VILAAdapter.__getitem__({idx})", flush=True)

        
        if isinstance(self.vila_dataset, ConcatDataset):
            # Find which sub-dataset this index belongs to
            dataset_idx = 0
            cumulative_sizes = self.vila_dataset.cumulative_sizes
            for i, size in enumerate(cumulative_sizes):
                if idx < size:
                    dataset_idx = i
                    break
            
            # Get the actual sub-dataset
            sub_dataset = self.vila_dataset.datasets[dataset_idx]
            # Calculate index within sub-dataset
            if dataset_idx == 0:
                sub_idx = idx
            else:
                sub_idx = idx - cumulative_sizes[dataset_idx - 1]
            
            # Access raw data from sub-dataset
            raw_data = self._get_raw_data(sub_dataset, sub_idx)
        else:
            raw_data = self._get_raw_data(self.vila_dataset, idx)
        
        return raw_data
    
    def _get_raw_data(self, dataset, idx):
        """
        Extract raw conversation data from VILA dataset.
        
        VILA datasets have a list_data_dict or instances attribute with raw data.
        """
        try:
            # Try to get raw data before tokenization
            if hasattr(dataset, "instances"):
                # BaseDataset format
                instance = dataset.instances[idx]
                return self._convert_vila_instance_to_messages(instance, dataset)
            elif hasattr(dataset, "list_data_dict"):
                # LazySupervisedDataset format
                instance = dataset.list_data_dict[idx]
                return self._convert_vila_instance_to_messages(instance, dataset)
            else:
                # Fallback: call __getitem__ and hope it returns usable data
                # This will be tokenized data, which is not ideal
                print(f"Warning: Dataset {type(dataset).__name__} has no instances/list_data_dict", flush=True)
                return self._get_fallback_messages()
        except Exception as e:
            print(f"Warning: Failed to get raw data at idx {idx}: {e}", flush=True)
            return self._get_fallback_messages()
    
    def _get_fallback_messages(self):
        """Return a minimal valid message structure as fallback."""
        return {"messages": [
            {"role": "user", "content": [{"type": "text", "text": "Hello."}]},
            {"role": "assistant", "content": [{"type": "text", "text": "Hello!"}]}
        ]}

    def _extract_images(self, instance, dataset):
        """Extract and normalize image objects from heterogeneous instance formats."""
        from PIL import Image

        image_folder = getattr(dataset, "media_dir", None) or getattr(dataset, "image_folder", None)
        raw_images = []
        for key in ("image", "images"):
            if key not in instance:
                continue
            value = instance.get(key)
            if isinstance(value, (list, tuple)):
                raw_images.extend(value)
            elif value is not None:
                raw_images.append(value)

        images = []
        for image_value in raw_images:
            try:
                if isinstance(image_value, bytes):
                    image_obj = Image.open(io.BytesIO(image_value)).convert("RGB")
                elif isinstance(image_value, str):
                    full_path = os.path.join(image_folder, image_value) if image_folder else image_value
                    image_obj = Image.open(full_path).convert("RGB")
                elif isinstance(image_value, Image.Image):
                    image_obj = image_value.convert("RGB")
                else:
                    print(f"Warning: Unsupported image type {type(image_value)} in VILADatasetAdapter", flush=True)
                    continue
                images.append(image_obj)
            except Exception as e:
                # Keep training progressing even when a single sample has broken media.
                print(f"Warning: Failed to load image ({type(image_value)}): {e}, using dummy image", flush=True)
                images.append(Image.new("RGB", (384, 384), (128, 128, 128)))
        return images

    def _extract_text(self, value):
        """Extract plain text from str/list/dict content payloads."""
        if value is None:
            return ""
        if isinstance(value, str):
            return value
        if isinstance(value, list):
            chunks = []
            for item in value:
                if isinstance(item, dict):
                    if item.get("type") == "text":
                        chunks.append(str(item.get("text", "")))
                elif isinstance(item, str):
                    chunks.append(item)
            return "\n".join([x for x in chunks if x]).strip()
        return str(value)

    def _build_turns(self, instance):
        """
        Convert heterogeneous raw schemas to simple (role, text) turns.
        Priority:
        1) HF datasets with `texts=[{user, assistant}]`
        2) LLaVA-style `conversations`
        3) Generic `messages`
        4) Fallback `question/instruction + answer/code`
        """
        turns = []

        texts = instance.get("texts")
        if isinstance(texts, list) and texts:
            for row in texts:
                if not isinstance(row, dict):
                    continue
                user_text = self._extract_text(row.get("user"))
                assistant_text = self._extract_text(row.get("assistant"))
                if user_text:
                    turns.append(("user", user_text))
                if assistant_text:
                    turns.append(("assistant", assistant_text))
            if turns:
                return turns

        conversations = instance.get("conversations")
        if isinstance(conversations, list) and conversations:
            for conv in conversations:
                if not isinstance(conv, dict):
                    continue
                from_field = conv.get("from", conv.get("role", "human"))
                role = "user" if from_field in ("human", "user") else "assistant"
                value = conv.get("value", conv.get("content", ""))
                turns.append((role, self._extract_text(value)))
            turns = [(r, t) for r, t in turns if t]
            if turns:
                return turns

        messages = instance.get("messages")
        if isinstance(messages, list) and messages:
            for msg in messages:
                if not isinstance(msg, dict):
                    continue
                role_raw = msg.get("role", msg.get("from", "user"))
                role = "user" if role_raw in ("human", "user") else "assistant"
                value = msg.get("content", msg.get("value", ""))
                turns.append((role, self._extract_text(value)))
            turns = [(r, t) for r, t in turns if t]
            if turns:
                return turns

        question = instance.get("question")
        if question is None:
            question = instance.get("instruction")
            response_key = "code"
        else:
            response_key = "qwen3vl_235b_thinking_response"
        if question is not None:
            response = instance.get(response_key, "")
            response = str(response or "")
            response = response.replace("</think>", "").replace("<think>", "")
            response = response.replace("<answer>", "").replace("</answer>", "").strip()
            turns.append(("user", str(question).replace("<image>", "").strip()))
            turns.append(("assistant", response))
            turns = [(r, t) for r, t in turns if t]
        return turns
    
    def _convert_vila_instance_to_messages(self, instance, dataset):
        """
        Convert VILA instance format to LMFlow messages format.
        
        VILA format: {"id": ..., "image": "path/to/img.jpg", "conversations": [{"from": "human", "value": "..."}, ...]}
        LMFlow format: {"messages": [{"role": "user", "content": [{"type": "text", "text": "..."}, {"type": "image", "image": PIL.Image}]}]}
        """
        turns = self._build_turns(instance)
        images = self._extract_images(instance, dataset)

        # Handle empty/invalid examples robustly
        if not turns:
            print("Warning: Empty conversations/texts in instance, using placeholder", flush=True)
            if images:
                return {"messages": [
                    {"role": "user", "content": [
                        {"type": "text", "text": "Describe this image."},
                        *[{"type": "image", "image": img} for img in images],
                    ]},
                    {"role": "assistant", "content": [
                        {"type": "text", "text": "This is an image."},
                    ]},
                ]}
            return self._get_fallback_messages()

        messages = []
        for role, text in turns:
            messages.append({
                "role": role,
                "content": [{"type": "text", "text": text}],
            })

        # Attach all available images to the first user turn.
        if images:
            user_idx = next((i for i, msg in enumerate(messages) if msg.get("role") == "user"), None)
            image_payload = [{"type": "image", "image": img} for img in images]
            if user_idx is None:
                messages.insert(0, {
                    "role": "user",
                    "content": [{"type": "text", "text": "Describe this image."}, *image_payload],
                })
            else:
                messages[user_idx]["content"].extend(image_payload)

        return {"messages": messages}
    
    def register_tokenizer(self, tokenizer, image_processor=None):
        """Pass-through for compatibility."""
        # VILA datasets don't need this, but LMFlow calls it
        pass

