import dataclasses
import gc
import os
import random
from collections.abc import Callable
from typing import Any

import numpy as np
import pyarrow as pa
import torch
from datasets import Dataset, Video, load_dataset

from fastvideo.configs.configs import (DatasetType, PreprocessConfig, VideoLoaderType)
from fastvideo.dataset.dataloader.parquet_io import (ParquetDatasetWriter, records_to_table)
from fastvideo.distributed.parallel_state import get_world_rank, get_world_size
from fastvideo.logger import init_logger
from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch

logger = init_logger(__name__)


class PreprocessingDataValidator:

    def __init__(self,
                 max_height: int = 1024,
                 max_width: int = 1024,
                 num_frames: int = 16,
                 train_fps: int = 24,
                 speed_factor: float = 1.0,
                 video_length_tolerance_range: float = 5.0,
                 drop_short_ratio: float = 0.0):
        self.max_height = max_height
        self.max_width = max_width
        self.num_frames = num_frames
        self.train_fps = train_fps
        self.speed_factor = speed_factor
        self.video_length_tolerance_range = video_length_tolerance_range
        self.drop_short_ratio = drop_short_ratio
        self.validators: dict[str, Callable[[dict[str, Any]], bool]] = {}
        self.filter_counts: dict[str, int] = {}

        self.num_items_before_filtering = 0
        self.num_items_after_filtering = 0

        self.register_validators()

    def register_validators(self) -> None:
        self.add_validator("data_type_validator", self._validate_data_type)
        self.add_validator("resolution_validator", self._validate_resolution)
        self.add_validator("frame_sampling_validator", self._validate_frame_sampling)

    def add_validator(self, name: str, validator: Callable[[dict[str, Any]], bool]) -> None:
        self.validators[name] = validator
        self.filter_counts[name] = 0

    def __call__(self, batch: dict[str, Any]) -> bool:
        """
        Validate whether the preprocessing data batch is valid.
        """
        self.num_items_before_filtering += 1

        for name, validator in self.validators.items():
            if not validator(batch):
                self.filter_counts[name] += 1
                return False

        self.num_items_after_filtering += 1
        return True

    def _validate_data_type(self, batch: dict[str, Any]) -> bool:
        """Validate basic validity of data items"""
        return not (batch["caption"] is None or batch["caption"] == "" or batch["fps"] is None or batch["fps"] <= 0
                    or batch["num_frames"] is None or batch["num_frames"] <= 0)

    def _validate_resolution(self, batch: dict[str, Any]) -> bool:
        """Validate resolution constraints"""

        if batch["resolution"] is not None:
            height = batch["resolution"].get("height", None)
            width = batch["resolution"].get("width", None)

        return not (height is None or width is None)

    def _validate_frame_sampling(self, batch: dict[str, Any]) -> bool:
        """Validate frame sampling constraints"""

        if (batch["num_frames"] / batch["fps"]
                > self.video_length_tolerance_range * (self.num_frames / self.train_fps * self.speed_factor)):
            return False

        frame_interval = batch["fps"] / self.train_fps
        start_frame_idx = 0
        frame_indices = np.arange(start_frame_idx, batch["num_frames"], frame_interval).astype(int)
        return not (len(frame_indices) < self.num_frames and random.random() < self.drop_short_ratio)

    def log_validation_stats(self):
        info = ""
        for name, count in self.filter_counts.items():
            info += f"failed in {name}: {count}, "
        info += f"number of items before filtering: {self.num_items_before_filtering}, "
        info += f"number of items after filtering: {self.num_items_after_filtering}"

        logger.info(info)


class VideoForwardBatchBuilder:

    def __init__(self, seed: int):
        self.seed = seed

    def __call__(self, batch: list) -> PreprocessBatch:
        forward_batch = PreprocessBatch(
            video_loader=[item["video"] for item in batch],
            video_file_name=[item["name"] for item in batch],
            height=[item["resolution"]["height"] for item in batch],
            width=[item["resolution"]["width"] for item in batch],
            fps=[item["fps"] for item in batch],
            num_frames=[item["num_frames"] for item in batch],
            prompt=[item["caption"] for item in batch],
            prompt_attention_mask=[],
            data_type="video",
            generator=torch.Generator("cpu").manual_seed(self.seed),
        )
        return forward_batch


class ParquetDatasetSaver:
    """Component for saving and writing Parquet datasets using shared parquet_io."""

    def __init__(self, flush_frequency: int, samples_per_file: int, schema: pa.Schema,
                 record_creator: Callable[..., list[dict[str, Any]]]):
        self.flush_frequency = flush_frequency
        self.samples_per_file = samples_per_file
        self.schema = schema
        self.create_records_from_batch = record_creator
        self.num_processed_samples: int = 0
        self._writer: ParquetDatasetWriter | None = None

    def save_and_write_parquet_batch(self,
                                     batch: PreprocessBatch,
                                     output_dir: str,
                                     extra_features: dict[str, Any] | None = None) -> None:
        """
        Save and write Parquet dataset batch
        
        Args:
            batch: PreprocessBatch containing video and metadata information
            output_dir: Output directory
            extra_features: Extra features
            
        Returns:
            Number of processed samples
        """
        assert isinstance(batch.latents, torch.Tensor)
        assert isinstance(batch.prompt_embeds, list)
        assert isinstance(batch.prompt_attention_mask, list)

        # Process non-padded embeddings (if needed)
        if batch.prompt_attention_mask is not None:
            batch.prompt_embeds = self._process_non_padded_embeddings(batch.prompt_embeds[0],
                                                                      batch.prompt_attention_mask[0])
        else:
            raise ValueError("prompt_attention_mask is None")

        # Prepare batch data for Parquet dataset
        batch_data: list[dict[str, Any]] = []

        for key in dataclasses.fields(batch):
            value = getattr(batch, key.name)
            if isinstance(value, list):
                for idx in range(len(value)):
                    if isinstance(value[idx], torch.Tensor):
                        value[idx] = value[idx].cpu().numpy()
            elif isinstance(value, torch.Tensor):
                value = value.cpu().numpy()
                setattr(batch, key.name, value)

        # Create record for Parquet dataset
        records = self.create_records_from_batch(batch)
        batch_data.extend(records)

        if batch_data:
            self.num_processed_samples += len(batch_data)
            table = records_to_table(batch_data, self.schema)
            if self._writer is None:
                os.makedirs(output_dir, exist_ok=True)
                self._writer = ParquetDatasetWriter(out_dir=output_dir, samples_per_file=self.samples_per_file)
            self._writer.append_table(table)
            logger.debug("Collected batch with %s samples", len(table))

        # If flush is needed
        if self.num_processed_samples >= self.flush_frequency:
            self.flush_tables()

    def _process_non_padded_embeddings(self, prompt_embeds: torch.Tensor,
                                       prompt_attention_mask: torch.Tensor) -> list[torch.Tensor]:
        """Process non-padded embeddings"""
        assert isinstance(prompt_embeds, torch.Tensor)
        assert isinstance(prompt_attention_mask, torch.Tensor)
        assert prompt_embeds.shape[0] == prompt_attention_mask.shape[0]

        # Get sequence lengths from attention masks (number of 1s)
        seq_lens = prompt_attention_mask.sum(dim=1)

        non_padded_embeds = []

        # Process each item in the batch
        for i in range(prompt_embeds.size(0)):
            seq_len = seq_lens[i].item()
            # Slice the embeddings and masks to keep only non-padding parts
            non_padded_embeds.append(prompt_embeds[i, :seq_len])

        return non_padded_embeds

    def flush_tables(self, write_remainder: bool = False):
        """Flush buffered records to disk.

        Args:
            output_dir: Directory where parquet files are written. Kept for API
                symmetry (writer already configured with this path).
            write_remainder: If True, also write any leftover rows smaller than
                ``samples_per_file`` as a final small file. Useful for the last flush.
        """
        if self._writer is None:
            return
        _ = self._writer.flush(write_remainder=write_remainder)
        # Reset processed sample count modulo samples_per_file
        remainder = self.num_processed_samples % self.samples_per_file
        self.num_processed_samples = 0 if write_remainder else remainder

    def clean_up(self) -> None:
        """Clean up all tables"""
        self.flush_tables(write_remainder=True)
        self._writer = None
        self.num_processed_samples = 0
        gc.collect()

    def __del__(self):
        self.clean_up()


def build_dataset(preprocess_config: PreprocessConfig, split: str, validator: Callable[[dict[str, Any]],
                                                                                       bool]) -> Dataset:
    if preprocess_config.dataset_type == DatasetType.HF:
        dataset = load_dataset(preprocess_config.dataset_path, split=split)
        dataset = dataset.filter(validator)
        dataset = dataset.shard(num_shards=get_world_size(), index=get_world_rank())
    elif preprocess_config.dataset_type == DatasetType.MERGED:
        metadata_json_path = os.path.join(preprocess_config.dataset_path, "videos2caption.json")
        video_folder = os.path.join(preprocess_config.dataset_path, "videos")
        dataset = load_dataset("json", data_files=metadata_json_path, split=split)
        column_names = dataset.column_names
        # rename columns to match the schema
        if "cap" in column_names:
            dataset = dataset.rename_column("cap", "caption")
        if "path" in column_names:
            dataset = dataset.rename_column("path", "name")

        dataset = dataset.filter(validator)
        dataset = dataset.shard(num_shards=get_world_size(), index=get_world_rank())

        # add video column
        def add_video_column(item: dict[str, Any]) -> dict[str, Any]:
            item["video"] = os.path.join(video_folder, item["name"])
            return item

        dataset = dataset.map(add_video_column)
        if preprocess_config.video_loader_type == VideoLoaderType.TORCHCODEC:
            dataset = dataset.cast_column("video", Video())
    else:
        raise ValueError(f"Invalid dataset type: {preprocess_config.dataset_type}")

    return dataset