# coding=utf-8
# Copyright 2020 The HuggingFace Datasets Authors and the DataLab Datasets Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from argparse import ArgumentParser, Namespace
import os
import re
import shutil

from datalabs.commands import BaseDatasetsCLICommand
from datalabs.utils.logging import get_logger

HIGHLIGHT_MESSAGE_PRE = (
    """<<<<<<< This should probably be modified because it mentions: """
)

HIGHLIGHT_MESSAGE_POST = """=======
>>>>>>>
"""

TO_HIGHLIGHT = [
    "TextEncoderConfig",
    "ByteTextEncoder",
    "SubwordTextEncoder",
    "encoder_config",
    "maybe_build_from_corpus",
    "manual_dir",
]

TO_CONVERT = [
    # (pattern, replacement)
    # Order is important here for some replacements
    (r"tfds\.core", r"datalab"),
    (r"tf\.io\.gfile\.GFile", r"open"),
    (r"tf\.([\w\d]+)", r"datalab.Value('\1')"),
    (r"tfds\.features\.Text\(\)", r"datalab.Value('string')"),
    (r"tfds\.features\.Text\(", r"datalab.Value('string'),"),
    (r"features\s*=\s*tfds.features.FeaturesDict\(", r"features=datalab.Features("),
    (r"tfds\.features\.FeaturesDict\(", r"dict("),
    (
        r"The TensorFlow Datasets Authors",
        r"The TensorFlow Datasets Authors and the HuggingFace Datasets Authors",
    ),
    (r"tfds\.", r"datalab."),
    (r"dl_manager\.manual_dir", r"self.config.data_dir"),
    (r"self\.builder_config", r"self.config"),
]


def convert_command_factory(args: Namespace):
    """
    Factory function used to convert a model TF 1.0 checkpoint in a PyTorch checkpoint.

    Returns: ConvertCommand
    """
    return ConvertCommand(args.tfds_path, args.datasets_directory)


class ConvertCommand(BaseDatasetsCLICommand):
    @staticmethod
    def register_subcommand(parser: ArgumentParser):
        """
        Register this command to argparse so it's available for the datalab-cli

        Args:
            parser: Root parser to register command-specific arguments
        """
        train_parser = parser.add_parser(
            "convert",
            help="Convert a TensorFlow Datasets dataset to a HuggingFace"
            " Datasets dataset.",
        )
        train_parser.add_argument(
            "--tfds_path",
            type=str,
            required=True,
            help="Path to a TensorFlow Datasets folder to convert or a single"
            " tfds file to convert.",
        )
        train_parser.add_argument(
            "--datasets_directory",
            type=str,
            required=True,
            help="Path to the HuggingFace Datasets folder.",
        )
        train_parser.set_defaults(func=convert_command_factory)

    def __init__(self, tfds_path: str, datasets_directory: str, *args):
        self._logger = get_logger("datalab-cli/converting")

        self._tfds_path = tfds_path
        self._datasets_directory = datasets_directory

    def run(self):
        if os.path.isdir(self._tfds_path):
            abs_tfds_path = os.path.abspath(self._tfds_path)
        elif os.path.isfile(self._tfds_path):
            abs_tfds_path = os.path.dirname(self._tfds_path)
        else:
            raise ValueError(
                "--tfds_path is neither a directory nor a file. Please check path."
            )

        abs_datasets_path = os.path.abspath(self._datasets_directory)

        self._logger.info(
            f"Converting datalab from {abs_tfds_path} to {abs_datasets_path}"
        )

        utils_files = []
        with_manual_update = []
        imports_to_builder_map = {}

        if os.path.isdir(self._tfds_path):
            file_names = os.listdir(abs_tfds_path)
        else:
            file_names = [os.path.basename(self._tfds_path)]

        for f_name in file_names:
            self._logger.info(f"Looking at file {f_name}")
            input_file = os.path.join(abs_tfds_path, f_name)
            output_file = os.path.join(abs_datasets_path, f_name)

            if (
                not os.path.isfile(input_file)
                or "__init__" in f_name
                or "_test" in f_name
                or ".py" not in f_name
            ):
                self._logger.info("Skipping file")
                continue

            with open(input_file, "r", encoding="utf-8") as f:
                lines = f.readlines()

            out_lines = []
            is_builder = False
            needs_manual_update = False
            tfds_imports = []
            for line in lines:
                out_line = line

                # Convert imports
                if "import tensorflow.compat.v2 as tf" in out_line:
                    continue
                elif "@tfds.core" in out_line:
                    continue
                elif "builder=self" in out_line:
                    continue
                elif "import tensorflow_datasets.public_api as tfds" in out_line:
                    out_line = "import datalab\n"
                elif "import tensorflow" in out_line:
                    # order is important here
                    out_line = ""
                    continue
                elif "from absl import logging" in out_line:
                    out_line = "from datalab import logging\n"
                elif "getLogger" in out_line:
                    out_line = out_line.replace("getLogger", "get_logger")
                elif any(expression in out_line for expression in TO_HIGHLIGHT):
                    needs_manual_update = True
                    to_remove = list(filter(lambda e: e in out_line, TO_HIGHLIGHT))
                    out_lines.append(HIGHLIGHT_MESSAGE_PRE + str(to_remove) + "\n")
                    out_lines.append(out_line)
                    out_lines.append(HIGHLIGHT_MESSAGE_POST)
                    continue
                else:
                    for pattern, replacement in TO_CONVERT:
                        out_line = re.sub(pattern, replacement, out_line)

                # Take care of saving utilities (to later move them together
                # with main script)
                if "tensorflow_datasets" in out_line:
                    match = re.match(
                        r"from\stensorflow_datasets.*import\s([^\.\r\n]+)", out_line
                    )
                    tfds_imports.extend(
                        imp.strip() for imp in match.group(1).split(",")
                    )
                    out_line = "from . import " + match.group(1)

                # Check we have not forget anything
                assert (
                    "tf." not in out_line
                    and "tfds." not in out_line
                    and "tensorflow_datasets" not in out_line
                ), f"Error converting {out_line.strip()}"

                if (
                    "GeneratorBasedBuilder" in out_line
                    or "BeamBasedBuilder" in out_line
                ):
                    is_builder = True
                out_lines.append(out_line)

            if is_builder or "wmt" in f_name:
                # We create a new directory for each dataset
                dir_name = f_name.replace(".py", "")
                output_dir = os.path.join(abs_datasets_path, dir_name)
                output_file = os.path.join(output_dir, f_name)
                os.makedirs(output_dir, exist_ok=True)
                self._logger.info(f"Adding directory {output_dir}")
                imports_to_builder_map.update({imp: output_dir for imp in tfds_imports})
            else:
                # Utilities will be moved at the end
                utils_files.append(output_file)

            if needs_manual_update:
                with_manual_update.append(output_file)

            with open(output_file, "w", encoding="utf-8") as f:
                f.writelines(out_lines)
            self._logger.info(f"Converted in {output_file}")

        for utils_file in utils_files:
            try:
                f_name = os.path.basename(utils_file)
                dest_folder = imports_to_builder_map[f_name.replace(".py", "")]
                self._logger.info(f"Moving {dest_folder} to {utils_file}")
                shutil.copy(utils_file, dest_folder)
            except KeyError:
                self._logger.error(
                    f"Cannot find destination folder for {utils_file}. Please"
                    f" copy manually."
                )

        if with_manual_update:
            for file_path in with_manual_update:
                self._logger.warning(
                    f"You need to manually update file {file_path} to remove"
                    f" configurations using 'TextEncoderConfig'."
                )
