from typing import TYPE_CHECKING, Optional

from tqdm import tqdm

from trainer.pipelines.pipeline_batch_info import PreprocessBatch
from trainer.workflow.preprocess.components import ParquetDatasetSaver
from trainer.workflow.preprocess.preprocess_workflow import PreprocessWorkflow

if TYPE_CHECKING:
    from torch.utils.data import DataLoader

    from trainer.pipelines.composed_pipeline_base import ComposedPipelineBase
    from trainer.workflow.preprocess.components import (
        VideoForwardBatchBuilder)


class PreprocessWorkflowT2V(PreprocessWorkflow):
    training_dataloader: "DataLoader"
    validation_dataloader: Optional["DataLoader"]
    preprocess_pipeline: "ComposedPipelineBase"
    processed_dataset_saver: "ParquetDatasetSaver"
    video_forward_batch_builder: "VideoForwardBatchBuilder"

    def run(self) -> None:
        # Training dataset preprocessing
        for batch in tqdm(self.training_dataloader,
                          desc="Preprocessing training dataset",
                          unit="batch"):
            forward_batch: PreprocessBatch = self.video_forward_batch_builder(
                batch)

            forward_batch = self.preprocess_pipeline.forward(
                forward_batch, self.trainer_args)

            self.processed_dataset_saver.save_and_write_parquet_batch(
                forward_batch, self.training_dataset_output_dir)

        self.processed_dataset_saver.flush_tables(
            self.training_dataset_output_dir)
        self.processed_dataset_saver.clean_up()

        # Validation dataset preprocessing
        if self.validation_dataloader is not None:
            for batch in tqdm(self.validation_dataloader,
                              desc="Preprocessing validation dataset",
                              unit="batch"):
                forward_batch = self.video_forward_batch_builder(batch)

                forward_batch = self.preprocess_pipeline.forward(
                    forward_batch, self.trainer_args)

                self.processed_dataset_saver.save_and_write_parquet_batch(
                    forward_batch, self.validation_dataset_output_dir)
            self.processed_dataset_saver.flush_tables(
                self.validation_dataset_output_dir)
            self.processed_dataset_saver.clean_up()