from typing import Callable, Optional, Union
from tqdm import tqdm
import os
import torch
import torchaudio
import torchaudio.functional as F
from torch.utils.data import Dataset, DataLoader, IterableDataset, random_split
from pytorch_lightning import LightningDataModule
import webdataset


class VLSP2020Dataset(Dataset):
    def __init__(self, root: str, sample_rate: int = 16000):
        super().__init__()

        self.sample_rate = sample_rate
        self.memory = self._prepare_data(root)
        self._memory = tuple(
            (v["transcript"], v["audio"]) for v in self.memory.values()
        )

    @staticmethod
    def _prepare_data(root: str):
        memory = {}

        for f in os.scandir(root):
            file_name, file_ext = os.path.splitext(f.name)

            if file_ext == ".txt":
                if file_name not in memory:
                    memory[file_name] = {"transcript": f.path}
                elif "transcript" not in memory[file_name]:
                    memory[file_name]["transcript"] = f.path
                else:
                    raise ValueError(f"Duplicate transcript for {f.path}")
            else:
                if file_name not in memory:
                    memory[file_name] = {"audio": f.path}
                elif "audio" not in memory[file_name]:
                    memory[file_name]["audio"] = f.path
                else:
                    raise ValueError(f"Duplicate audio for {f.path}")

        for key, value in memory.items():
            if "audio" not in value:
                raise ValueError(f"Missing audio for {key}")
            elif "transcript" not in value:
                raise ValueError(f"Missing transcript for {key}")

        return memory

    def __len__(self):
        return len(self.memory)

    def __getitem__(self, index: int):
        transcript, audio = self._memory[index]

        with open(transcript, "r") as f:
            transcript = f.read()

        audio, sample_rate = torchaudio.load(audio)
        audio = F.resample(audio, sample_rate, self.sample_rate)

        return transcript, audio


class VLSP2020TarDataset:
    def __init__(self, outpath: str):
        self.outpath = outpath

    def convert(self, dataset: VLSP2020Dataset):
        writer = webdataset.TarWriter(self.outpath)

        for idx, (transcript, audio) in enumerate(tqdm(dataset, colour="green")):
            writer.write(
                {
                    "__key__": f"{idx:08d}",
                    "txt": transcript,
                    "pth": audio,
                }
            )

        writer.close()

    def load(self) -> webdataset.WebDataset:
        self.data = (
            webdataset.WebDataset(self.outpath)
            .decode(
                webdataset.handle_extension("txt", lambda x: x.decode("utf-8")),
                webdataset.torch_audio,
            )
            .to_tuple("txt", "pth")
        )

        return self.data


def get_dataloader(
    dataset: Union[VLSP2020Dataset, webdataset.WebDataset],
    return_transcript: bool = False,
    target_transform: Optional[Callable] = None,
    batch_size: int = 32,
    num_workers: int = 2,
):
    def collate_fn(batch):
        def get_audio(item):
            audio = item[1]

            assert (
                isinstance(audio, torch.Tensor)
                and audio.ndim == 2
                and audio.size(0) == 1
            )

            return audio.squeeze(0)

        audio = tuple(get_audio(item) for item in batch)

        if return_transcript:
            if target_transform is not None:
                transcript = tuple(target_transform(item[0]) for item in batch)
            else:
                transcript = tuple(item[0] for item in batch)

            return transcript, audio
        else:
            return audio

    return DataLoader(
        dataset, batch_size=batch_size, num_workers=num_workers, collate_fn=collate_fn
    )