import os
import warnings
from pathlib import Path
from typing import List, Dict, Tuple, Optional, Union

import torchaudio
from torchaudio.datasets.utils import extract_archive, unicode_csv_reader, validate_file
from torch import Tensor
from torch.utils.data import Dataset

# Default TSV should be one of
# dev.tsv
# invalidated.tsv
# other.tsv
# test.tsv
# train.tsv
# validated.tsv

FOLDER_IN_ARCHIVE = "CommonVoice"
LANGUAGE = "english"
VERSION = "cv-corpus-5.1-2020-06-22"
TSV = "train.tsv"
_CHECKSUMS = {
    "cv-corpus-5.1-2020-06-22/tt.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/en.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/de.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/fr.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/cy.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/br.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/cv.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/tr.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/ky.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/ga-IE.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/kab.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/ca.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/zh-TW.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/sl.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/it.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/nl.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/cnh.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/eo.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/et.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/fa.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/eu.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/es.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/zh-CN.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/mn.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/sah.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/dv.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/rw.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/sv-SE.tar.gz": None,
    "cv-corpus-5.1-2020-06-22/ru.tar.gz": None,
}


def load_commonvoice_item(line: List[str],
                          header: List[str],
                          path: str,
                          folder_audio: str) -> Tuple[Tensor, int, Dict[str, str]]:
    # Each line as the following data:
    # client_id, path, sentence, up_votes, down_votes, age, gender, accent

    assert header[1] == "path"
    fileid = line[1]

    filename = os.path.join(path, folder_audio, fileid)

    waveform, sample_rate = torchaudio.load(filename)

    dic = dict(zip(header, line))

    return waveform, sample_rate, dic


class COMMONVOICE(Dataset):
    """Create a Dataset for `CommonVoice <https://commonvoice.mozilla.org/>`_.

    Args:
        root (str or Path): Path to the directory where the dataset is found or downloaded.
        tsv (str, optional): The name of the tsv file used to construct the metadata.
            (default: ``"train.tsv"``)
        url (str, optional): Deprecated.
        folder_in_archive (str, optional): The top-level directory of the dataset.
        version (str): Version string. (default: ``"cv-corpus-5.1-2020-06-22"``)
        language (str, optional): Language of the dataset. (default: None)
            The following values are mapped to their corresponding shortened version:
            ``"tatar"``, ``"english"``, ``"german"``,
            ``"french"``, ``"welsh"``, ``"breton"``, ``"chuvash"``, ``"turkish"``, ``"kyrgyz"``,
            ``"irish"``, ``"kabyle"``, ``"catalan"``, ``"taiwanese"``, ``"slovenian"``,
            ``"italian"``, ``"dutch"``, ``"hakha chin"``, ``"esperanto"``, ``"estonian"``,
            ``"persian"``, ``"portuguese"``, ``"basque"``, ``"spanish"``, ``"chinese"``,
            ``"mongolian"``, ``"sakha"``, ``"dhivehi"``, ``"kinyarwanda"``, ``"swedish"``,
            ``"russian"``, ``"indonesian"``, ``"arabic"``, ``"tamil"``, ``"interlingua"``,
            ``"latvian"``, ``"japanese"``, ``"votic"``, ``"abkhaz"``, ``"cantonese"`` and
            ``"romansh sursilvan"``.
            For the other allowed values, Please checkout https://commonvoice.mozilla.org/en/datasets.
        download (bool, optional): Deprecated.
    """

    _ext_txt = ".txt"
    _ext_audio = ".mp3"
    _folder_audio = "clips"

    def __init__(self,
                 root: Union[str, Path],
                 tsv: str = TSV,
                 url: Optional[str] = None,
                 folder_in_archive: str = FOLDER_IN_ARCHIVE,
                 version: str = VERSION,
                 language: str = LANGUAGE,
                 download: Optional[bool] = False) -> None:

        if download is True:
            raise RuntimeError(
                "The dataset is no longer publicly accessible. You need to "
                "download the archives externally and place them in the root "
                "directory."
            )
        elif download is False:
            warnings.warn(
                "The use of the download flag is deprecated, since the dataset "
                "is no longer directly accessible.", RuntimeWarning
            )

        if url is not None:
            warnings.warn(
                "The use of the url flag is deprecated, since the dataset "
                "is no longer publicly accessible. To specify the language of the dataset, "
                "please use the language parameter instead.", RuntimeWarning
            )

        languages = {
            "tatar": "tt",
            "english": "en",
            "german": "de",
            "french": "fr",
            "welsh": "cy",
            "breton": "br",
            "chuvash": "cv",
            "turkish": "tr",
            "kyrgyz": "ky",
            "irish": "ga-IE",
            "kabyle": "kab",
            "catalan": "ca",
            "taiwanese": "zh-TW",
            "slovenian": "sl",
            "italian": "it",
            "dutch": "nl",
            "hakha chin": "cnh",
            "esperanto": "eo",
            "estonian": "et",
            "persian": "fa",
            "portuguese": "pt",
            "basque": "eu",
            "spanish": "es",
            "chinese": "zh-CN",
            "mongolian": "mn",
            "sakha": "sah",
            "dhivehi": "dv",
            "kinyarwanda": "rw",
            "swedish": "sv-SE",
            "russian": "ru",
            "indonesian": "id",
            "arabic": "ar",
            "tamil": "ta",
            "interlingua": "ia",
            "latvian": "lv",
            "japanese": "ja",
            "votic": "vot",
            "abkhaz": "ab",
            "cantonese": "zh-HK",
            "romansh sursilvan": "rm-sursilv"
        }

        if language in languages:
            ext_archive = ".tar.gz"
            language = languages[language]
            url = os.path.join(version, language + ext_archive)
        else:
            raise ValueError(
                'Allowed language values are ``"tatar"``, ``"english"``, ``"german"``,'
                '``"french"``, ``"welsh"``, ``"breton"``, ``"chuvash"``, ``"turkish"``, ``"kyrgyz"``,'
                '``"irish"``, ``"kabyle"``, ``"catalan"``, ``"taiwanese"``, ``"slovenian"``,'
                '``"italian"``, ``"dutch"``, ``"hakha chin"``, ``"esperanto"``, ``"estonian"``,'
                '``"persian"``, ``"portuguese"``, ``"basque"``, ``"spanish"``, ``"chinese"``,'
                '``"mongolian"``, ``"sakha"``, ``"dhivehi"``, ``"kinyarwanda"``, ``"swedish"``,'
                '``"russian"``, ``"indonesian"``, ``"arabic"``, ``"tamil"``, ``"interlingua"``,'
                '``"latvian"``, ``"japanese"``, ``"votic"``, ``"abkhaz"``, ``"cantonese"`` and'
                '``"romansh sursilvan"``.'
            )

        # Get string representation of 'root' in case Path object is passed
        root = os.fspath(root)

        basename = os.path.basename(url)
        archive = os.path.join(root, basename)

        basename = basename.rsplit(".", 2)[0]
        folder_in_archive = os.path.join(folder_in_archive, version, basename)

        self._path = os.path.join(root, folder_in_archive)

        if not os.path.isdir(self._path):
            if os.path.isfile(archive):
                checksum = _CHECKSUMS.get(url, None)
                if checksum:
                    filepath = os.path.basename(url)
                    with open(filepath, "rb") as file_obj:
                        if not validate_file(file_obj, checksum, "sha256"):
                            raise RuntimeError(
                                f"The hash of {filepath} does not match. Delete the file manually and retry."
                            )
                extract_archive(archive)
            else:
                raise RuntimeError(
                    "The dataset is no longer publicly accessible. You need to "
                    "download the archives externally and place them in the root "
                    "directory."
                )

        self._tsv = os.path.join(root, folder_in_archive, tsv)

        with open(self._tsv, "r") as tsv:
            walker = unicode_csv_reader(tsv, delimiter="\t")
            self._header = next(walker)
            self._walker = list(walker)

    def __getitem__(self, n: int) -> Tuple[Tensor, int, Dict[str, str]]:
        """Load the n-th sample from the dataset.

        Args:
            n (int): The index of the sample to be loaded

        Returns:
            tuple: ``(waveform, sample_rate, dictionary)``,  where dictionary is built
            from the TSV file with the following keys: ``client_id``, ``path``, ``sentence``,
            ``up_votes``, ``down_votes``, ``age``, ``gender`` and ``accent``.
        """
        line = self._walker[n]
        return load_commonvoice_item(line, self._header, self._path, self._folder_audio)

    def __len__(self) -> int:
        return len(self._walker)
