from typing import Any, List, Optional, Sequence, Union

import hydra
import lightning as pl
import torch
from lightning.pytorch.utilities.types import EVAL_DATALOADERS
from omegaconf import DictConfig
from torch.utils.data import DataLoader

from relik.common.log import get_logger
from relik.retriever.data.datasets import GoldenRetrieverDataset

logger = get_logger()


class GoldenRetrieverPLDataModule(pl.LightningDataModule):
    def __init__(
        self,
        train_dataset: Optional[GoldenRetrieverDataset] = None,
        val_datasets: Optional[Sequence[GoldenRetrieverDataset]] = None,
        test_datasets: Optional[Sequence[GoldenRetrieverDataset]] = None,
        num_workers: Optional[Union[DictConfig, int]] = None,
        datasets: Optional[DictConfig] = None,
        *args,
        **kwargs,
    ):
        super().__init__()
        self.datasets = datasets
        if num_workers is None:
            num_workers = 0
        if isinstance(num_workers, int):
            num_workers = DictConfig(
                {"train": num_workers, "val": num_workers, "test": num_workers}
            )
        self.num_workers = num_workers
        # data
        self.train_dataset: Optional[GoldenRetrieverDataset] = train_dataset
        self.val_datasets: Optional[Sequence[GoldenRetrieverDataset]] = val_datasets
        self.test_datasets: Optional[Sequence[GoldenRetrieverDataset]] = test_datasets

    def prepare_data(self, *args, **kwargs):
        """
        Method for preparing the data before the training. This method is called only once.
        It is used to download the data, tokenize the data, etc.
        """
        pass

    def setup(self, stage: Optional[str] = None):
        if stage == "fit" or stage is None:
            # usually there is only one dataset for train
            # if you need more train loader, you can follow
            # the same logic as val and test datasets
            if self.train_dataset is None:
                self.train_dataset = hydra.utils.instantiate(self.datasets.train)
                self.val_datasets = [
                    hydra.utils.instantiate(dataset_cfg)
                    for dataset_cfg in self.datasets.val
                ]
        if stage == "test":
            if self.test_datasets is None:
                self.test_datasets = [
                    hydra.utils.instantiate(dataset_cfg)
                    for dataset_cfg in self.datasets.test
                ]

    def train_dataloader(self, *args, **kwargs) -> DataLoader:
        torch_dataset = self.train_dataset.to_torch_dataset()
        return DataLoader(
            # self.train_dataset.to_torch_dataset(),
            torch_dataset,
            shuffle=False,
            batch_size=None,
            num_workers=self.num_workers.train,
            pin_memory=True,
            collate_fn=lambda x: x,
        )

    def val_dataloader(self, *args, **kwargs) -> Union[DataLoader, List[DataLoader]]:
        dataloaders = []
        for dataset in self.val_datasets:
            torch_dataset = dataset.to_torch_dataset()
            dataloaders.append(
                DataLoader(
                    torch_dataset,
                    shuffle=False,
                    batch_size=None,
                    num_workers=self.num_workers.val,
                    pin_memory=True,
                    collate_fn=lambda x: x,
                )
            )
        return dataloaders

    def test_dataloader(self, *args, **kwargs) -> Union[DataLoader, List[DataLoader]]:
        dataloaders = []
        for dataset in self.test_datasets:
            torch_dataset = dataset.to_torch_dataset()
            dataloaders.append(
                DataLoader(
                    torch_dataset,
                    shuffle=False,
                    batch_size=None,
                    num_workers=self.num_workers.test,
                    pin_memory=True,
                    collate_fn=lambda x: x,
                )
            )
        return dataloaders

    def predict_dataloader(self) -> EVAL_DATALOADERS:
        raise NotImplementedError

    def transfer_batch_to_device(
        self, batch: Any, device: torch.device, dataloader_idx: int
    ) -> Any:
        return super().transfer_batch_to_device(batch, device, dataloader_idx)

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__name__}(" f"{self.datasets=}, " f"{self.num_workers=}, "
        )