# coding: utf-8

import os
import pyarrow.parquet as pq
import torch

from data.parquet_utils import init_arrow_fs, read_parquet_rows
from data.datasets_factory.dataset_x2t_interleave import X2TInterleaveIterableDataset


class X2TInterleaveLocalIterableDataset(X2TInterleaveIterableDataset):
    def __init__(self, *args, **kwargs):
        # Debug/local smoke-test dataset: read local parquet rows with raw image/video bytes.
        kwargs.setdefault("raw_bytes_input", True)
        super().__init__(*args, **kwargs)
        # Debug-only option for quick smoke tests on tiny local datasets:
        # repeat the same parquet file list to avoid very short epochs / empty shards.
        self.debug_parquet_repeat = max(1, int(kwargs.get("debug_parquet_repeat", 1)))
        self.data_paths = (kwargs.get("all_data_paths") or self.data_dir_list) * self.debug_parquet_repeat
        self._validate_local_data_paths()

    def _validate_local_data_paths(self):
        for path in self.data_paths:
            if str(path).startswith(("hdfs://", "http://", "https://")):
                raise ValueError(
                    f"{self.__class__.__name__} only supports local parquet files with embedded bytes, got: {path}"
                )
            if not os.path.isfile(path):
                raise FileNotFoundError(f"Local parquet file not found: {path}")

    def _validate_local_row(self, row):
        if self.dataset_type == "x2t_general":
            if "image_bytes" not in row and "video_bytes" not in row:
                raise ValueError(
                    f"{self.__class__.__name__} expects x2t_general rows to contain image_bytes or video_bytes"
                )
        else:
            raise ValueError(
                f"{self.__class__.__name__} only supports dataset_type=x2t_general for local byte reading, got {self.dataset_type}"
            )

    def set_epoch(self, seed=42):
        if not self.data_paths:
            return

        data_paths = sorted(self.data_paths)
        self.rng.seed(seed)
        self.rng.shuffle(data_paths)
        self.data_paths_per_rank = data_paths
        self.num_files_per_rank = len(data_paths)

    def __iter__(self):
        if not hasattr(self, "data_paths_per_rank"):
            self.set_epoch(self.seed)

        info = torch.utils.data.get_worker_info()
        worker_id = 0 if info is None else info.id
        num_workers = 1 if info is None else info.num_workers
        global_worker_id = self.local_rank * num_workers + worker_id
        global_num_workers = max(1, self.world_size * num_workers)

        while True:
            for parquet_idx, parquet_file_path in enumerate(self.data_paths_per_rank):
                fs = init_arrow_fs(parquet_file_path)
                with fs.open_input_file(parquet_file_path) as f:
                    parquet_file = pq.ParquetFile(f)
                    for row_group_id in range(parquet_file.num_row_groups):
                        rows = read_parquet_rows(parquet_file, row_group_id)
                        for row_idx, row in enumerate(rows):
                            if row_idx % global_num_workers != global_worker_id:
                                continue
                            try:
                                self._validate_local_row(row)
                            except Exception as e:
                                self.logger.warning(
                                    f"Invalid local x2t row in rg#{row_group_id} row#{row_idx} {parquet_file_path}: {e}"
                                )
                                continue
                            sample = self._process_row(
                                row=row,
                                parquet_idx=parquet_idx,
                                row_group_id=row_group_id,
                                row_idx=row_idx,
                                worker_id=worker_id,
                                parquet_file_path=parquet_file_path,
                            )
                            if sample:
                                yield sample

            self.logger.info(f"{self.dataset_name} repeat in rank-{self.local_rank} worker-{worker_id}")