from pathlib import Path

import datasets
import pandas as pd

from seacrowd.utils import schemas
from seacrowd.utils.configs import SEACrowdConfig
from seacrowd.utils.constants import Licenses, Tasks


_CITATION = ""

_DATASETNAME = "thai_gpteacher"

_DESCRIPTION = """This is a Thai-instructed dataset translated using Google Cloud Translation from GPTeacher, a
collection of modular datasets generated by GPT-4, General-Instruct & Roleplay-Instruct and is comprised of around
20,000 examples with deduplication. The dataset was asked to include reasoning and thought steps in the example
responses where appropriate.
"""

_HOMEPAGE = "https://huggingface.co/datasets/Thaweewat/gpteacher-20k-th"

_LANGUAGES = ["tha"]

_LICENSE = Licenses.CC_BY_SA_3_0.value

_LOCAL = False

_URLS = {
    _DATASETNAME: {
        "train": {
            "qa": "https://huggingface.co/datasets/Thaweewat/gpteacher-20k-th/resolve/main/gpteacher-gpt4-instruct-qa-18k-th.parquet",
            "role_play": "https://huggingface.co/datasets/Thaweewat/gpteacher-20k-th/resolve/main/gpteacher-gpt4-instruct-roleplay-2k-th.parquet",
        }
    },
}

_SUPPORTED_TASKS = [Tasks.INSTRUCTION_TUNING]

_SOURCE_VERSION = "1.0.0"

_SEACROWD_VERSION = "1.0.0"


class ThaiGPTeacherDataset(datasets.GeneratorBasedBuilder):
    """Thai-instructed dataset translated using Google Cloud Translation from GPTeacher."""

    SOURCE_VERSION = datasets.Version(_SOURCE_VERSION)
    SEACROWD_VERSION = datasets.Version(_SEACROWD_VERSION)

    BUILDER_CONFIGS = [
        SEACrowdConfig(
            name=f"{_DATASETNAME}_source",
            version=SOURCE_VERSION,
            description=f"{_DATASETNAME} source schema",
            schema="source",
            subset_id=_DATASETNAME,
        ),
        SEACrowdConfig(
            name=f"{_DATASETNAME}_seacrowd_t2t",
            version=SEACROWD_VERSION,
            description=f"{_DATASETNAME} SEACrowd schema",
            schema="seacrowd_t2t",
            subset_id=_DATASETNAME,
        ),
    ]

    DEFAULT_CONFIG_NAME = f"{_DATASETNAME}_source"

    def _info(self) -> datasets.DatasetInfo:
        if self.config.schema == "source":
            features = datasets.Features(
                {
                    "instruction": datasets.Value("string"),
                    "input": datasets.Value("string"),
                    "output": datasets.Value("string"),
                }
            )

        elif self.config.schema == "seacrowd_t2t":
            features = schemas.text2text_features

        return datasets.DatasetInfo(
            description=_DESCRIPTION,
            features=features,
            homepage=_HOMEPAGE,
            license=_LICENSE,
            citation=_CITATION,
        )

    def _split_generators(self, dl_manager: datasets.DownloadManager) -> list[datasets.SplitGenerator]:
        urls = _URLS[_DATASETNAME]
        data_dir = dl_manager.download_and_extract(urls)

        return [
            datasets.SplitGenerator(
                name=datasets.Split.TRAIN,
                gen_kwargs={"filepath": data_dir, "split": "train"},
            ),
        ]

    def _generate_examples(self, filepath: Path, split: str) -> tuple[int, dict]:
        df1 = pd.read_parquet(filepath[split]["qa"])
        df2 = pd.read_parquet(filepath[split]["role_play"])
        df = pd.concat([df1, df2], ignore_index=True)
        if self.config.schema == "source":
            for i, row in df.iterrows():
                yield i, {"instruction": row["instruction"], "input": row["input"], "output": row["output"]}

        elif self.config.schema == "seacrowd_t2t":
            for i, row in df.iterrows():
                yield i, {
                    "id": str(i),
                    "text_1": row["instruction"] + "\n" + row["input"],
                    "text_2": row["output"],
                    "text_1_name": "instruction + input",
                    "text_2_name": "output",
                }
