# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
#
# 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
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 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.

"""Dataset upload modules"""
import tarfile
import os
import glob
import sys
from handlers.utilities import Code
from handlers.utilities import run_system_command, get_handler_root


# Simple helper class for ease of code migration
class SimpleHandler:
    """Helper class holding dataset information"""

    def __init__(self, handler_metadata):
        """Initialize the Handler helper class"""
        self.root = get_handler_root(handler_metadata.get("id"))
        self.type = handler_metadata.get("type")
        self.format = handler_metadata.get("format")


def _extract_images(tar_path, dest):
    """Function to extract images, other directories on same level as images to root of dataset"""
    # Infer how many components to strip to get images,labels to top of dataset directory
    # Assumes: images, other necessary directories are in the same level
    with tarfile.open(tar_path) as tar:
        strip_components = 0
        names = [tinfo.name for tinfo in tar.getmembers()]
        for name in names:
            if "/images/" in name:
                strip_components = name.split("/").index("images")
                break
    # Build shell command for untarring
    untar_command = f"tar -xf {tar_path} --strip-components={strip_components} -C {dest}/"
    # Run shell command
    print("Untarring data started", file=sys.stderr)
    run_system_command(untar_command)
    print("Untarring data complete", file=sys.stderr)

    # Remove .tar.gz file
    print("Removing data tar file", file=sys.stderr)
    os.remove(tar_path)
    print("Deleted data tar file", file=sys.stderr)


def write_dir_contents(directory, file):
    """Write contents of a directory to a file"""
    with open(file, "w", encoding='utf-8') as f:
        for dir_files in sorted(glob.glob(directory + "/*")):
            f.write(dir_files + "\n")


def _untar_and_delete(tar_path, dest):
    """Run untar command and delete the tar file"""
    # Build shell command for untarring
    untar_command = f"tar -xf {tar_path} -C {dest}/"
    # Run shell command
    run_system_command(untar_command)
    # Remove .tar.gz file
    os.remove(tar_path)


def object_detection(tar_path, handler_metadata):
    """
    OD Dataset structure
    Upload - uploads and untars
    - /images
    - /labels
    Convert
    - /tfrecords (generated by dataset convert)
    - /tfrecords/classes.json - A json file with a list
    Augment
    - Creates temp output folders and moves them to /images and /labels

    """
    handler = SimpleHandler(handler_metadata)

    try:
        _extract_images(tar_path, handler.root)
        # Validate images and labels paths exist
        assert os.path.exists(os.path.join(handler.root, "images"))
        if handler.format == "kitti":
            assert os.path.exists(os.path.join(handler.root, "labels"))
        elif handler.format == "coco":
            assert os.path.exists(os.path.join(handler.root, "annotations.json"))
        elif handler.format == "coco_raw":
            assert os.path.exists(os.path.join(handler.root, "label_map.txt"))
        msg = "Upload successful"
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


instance_segmentation = object_detection


def semantic_segmentation(tar_path, handler_metadata):
    """
    Upload - uploads and creates .txt files
    - /images/
    - /masks/

    No Actions
    """
    handler = SimpleHandler(handler_metadata)

    try:
        _extract_images(tar_path, handler.root)
        # Validate images and masks paths exist
        assert os.path.exists(os.path.join(handler.root, "images"))
        write_dir_contents(os.path.join(handler.root, "images"), os.path.join(handler.root, "images.txt"))
        if handler.format == "unet":
            assert os.path.exists(os.path.join(handler.root, "masks"))
            write_dir_contents(os.path.join(handler.root, "masks"), os.path.join(handler.root, "masks.txt"))
        elif handler.format == "coco":
            assert os.path.exists(os.path.join(handler.root, "annotations.json"))
        msg = "Upload successful"
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


def character_recognition(tar_path, handler_metadata):
    """
    LPRNET Dataset structure
    Upload - uploads and untars
    - /images
    - /labels
    - /characters.txt

    """
    handler = SimpleHandler(handler_metadata)

    try:
        _extract_images(tar_path, handler.root)
        # Validate images and labels paths exist
        assert os.path.exists(os.path.join(handler.root, "image"))
        if handler.format != "raw":
            assert os.path.exists(os.path.join(handler.root, "label"))
            assert os.path.exists(os.path.join(handler.root, "characters.txt"))
        msg = "Upload successful"
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


def ocrnet(tar_path, handler_metadata):
    """
    OCRNET Dataset structure
    Upload - uploads and untars
    - /train/gt_new.txt: train dataset
    - /test/gt_new.txt: val dataset

    """
    handler = SimpleHandler(handler_metadata)

    try:
        _extract_images(tar_path, handler.root)
        # Validate images and labels paths exist

        assert os.path.exists(os.path.join(handler.root, "character_list"))
        assert os.path.exists(os.path.join(handler.root, "train")) or os.path.exists(os.path.join(handler.root, "test"))
        if os.path.exists(os.path.join(handler.root, "train")):
            os.system(f"chmod -R 777 {os.path.join(handler.root, 'train')}")
        if os.path.exists(os.path.join(handler.root, "test")):
            os.system(f"chmod -R 777 {os.path.join(handler.root, 'test')}")
        assert os.path.exists(os.path.join(handler.root, "train/gt_new.txt")) or os.path.exists(os.path.join(handler.root, "test/gt_new.txt"))
        msg = "Upload successful"
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


def ocdnet(tar_path, handler_metadata):
    """
    OCDNET Dataset structure
    Upload - uploads and untars
    - /train/img: train images
    - /train/gt: train ground_truth
    - /test/img: val images
    - /test/gt: val ground_truth

    """
    handler = SimpleHandler(handler_metadata)

    try:
        _extract_images(tar_path, handler.root)
        # Validate images and labels paths exist

        assert os.path.exists(os.path.join(handler.root, "train/img")) or os.path.exists(os.path.join(handler.root, "test/img"))
        assert os.path.exists(os.path.join(handler.root, "train/gt")) or os.path.exists(os.path.join(handler.root, "test/gt"))
        msg = "Upload successful"
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


def optical_inspection(tar_path, handler_metadata):
    """
    Optical Inspection Dataset structure
    Upload - uploads and untars
    - images: images
    - dataset.csv: ground_truth

    """
    handler = SimpleHandler(handler_metadata)

    try:
        _extract_images(tar_path, handler.root)
        # Validate images and labels paths exist

        assert os.path.exists(os.path.join(handler.root, "images")) and os.path.exists(os.path.join(handler.root, "dataset.csv"))
        msg = "Upload successful"
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


def ml_recog(tar_path, handler_metadata):
    """
    Metric Learning Recognition Dataset structure
    Upload - uploads and untars
    - metric_learning_recognition/
      - retail-product-checkout-dataset/
      - retail-product-checkout-dataset_classification_demo/
        - known_classes
        - unknown_classes

    """
    handler = SimpleHandler(handler_metadata)

    try:
        _extract_images(tar_path, handler.root)
        # Validate images and labels paths exist

        assert os.path.exists(os.path.join(handler.root, "metric_learning_recognition"))
        assert os.path.exists(os.path.join(handler.root, "metric_learning_recognition", "retail-product-checkout-dataset_classification_demo"))
        msg = "Upload successful"
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


def image_classification(tar_path, handler_metadata):
    """
    Raw:
    images/

    Default:
    images/<class1>
    images/<class2>
    ...

    Custom:
    images/
    train.csv
    val.csv

    """
    handler = SimpleHandler(handler_metadata)
    try:
        print("Extracting images from data tarball file", file=sys.stderr)
        _extract_images(tar_path, handler.root)
        print("Extraction complete", file=sys.stderr)
        # Validate images and labels paths exist
        assert len(glob.glob(os.path.join(handler.root, "images*"))) == 1
        if handler.format == "custom":
            assert os.path.exists(os.path.join(handler.root, "val.csv"))
        if handler.format == "classification_pyt":
            assert os.path.exists(os.path.join(handler.root, "classes.txt"))
        msg = "Upload successful"
        print("Returning sucess code to the api call", file=sys.stderr)
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


def bpnet(tar_path, handler_metadata):
    """
    OD Dataset structure
    Upload - uploads and untars
    - /images
    - /labels
    Convert
    - /tfrecords (generated by dataset convert)
    - /tfrecords/classes.json - A json file with a list
    Augment
    - Creates temp output folders and moves them to /images and /labels

    """
    handler = SimpleHandler(handler_metadata)

    try:
        _extract_images(tar_path, handler.root)
        # Validate images and labels paths exist
        assert os.path.exists(os.path.join(handler.root, "annotations"))
        assert os.path.exists(os.path.join(handler.root, "annotations", "person_keypoints_train2017.json"))
        assert os.path.exists(os.path.join(handler.root, "annotations", "person_keypoints_val2017.json"))
        assert os.path.exists(os.path.join(handler.root, "train2017"))
        assert os.path.exists(os.path.join(handler.root, "val2017"))

        assert os.path.exists(os.path.join(handler.root, "coco_spec.json"))
        os.system(f"chmod -R 777 {os.path.join(handler.root, 'coco_spec.json')}")
        assert os.path.exists(os.path.join(handler.root, "bpnet_18joints.json"))
        assert os.path.exists(os.path.join(handler.root, "infer_spec.yaml"))
        msg = "Upload successful"
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


def fpenet(tar_path, handler_metadata):
    """
    Default:
    data/afw
    data/afw.json or data/afw_10.json
    ...

    """
    handler = SimpleHandler(handler_metadata)
    try:
        print("Extracting images from data tarball file", file=sys.stderr)
        _extract_images(tar_path, handler.root)
        print("Extraction complete", file=sys.stderr)
        # Validate images and labels paths exist
        assert os.path.exists(os.path.join(handler.root, "data", "afw"))
        assert (os.path.exists(os.path.join(handler.root, "data", "afw", "afw.json")) or os.path.exists(os.path.join(handler.root, "data", "afw_10", "afw_10.json")))
        assert os.path.exists(os.path.join(handler.root, "data.json"))
        os.system(f"chmod -R 777 {os.path.join(handler.root, 'data')}")
        os.system(f"chmod -R 777 {os.path.join(handler.root, 'data.json')}")
        msg = "Upload successful"
        print("Returning sucess code to the api call", file=sys.stderr)
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


def action_recognition(tar_path, handler_metadata):
    """
    Default:
    train
    test
    ...

    """
    handler = SimpleHandler(handler_metadata)
    try:
        print("Extracting images from data tarball file", file=sys.stderr)
        _extract_images(tar_path, handler.root)
        print("Extraction complete", file=sys.stderr)
        # Validate images and labels paths exist
        assert os.path.exists(os.path.join(handler.root, "train"))
        assert os.path.exists(os.path.join(handler.root, "test"))
        msg = "Upload successful"
        print("Returning sucess code to the api call", file=sys.stderr)
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


def pointpillars(tar_path, handler_metadata):
    """
    OD Dataset structure
    Upload - uploads and untars
    - /images
    - /labels
    - /velodyne
    - /calib

    """
    handler = SimpleHandler(handler_metadata)

    try:
        _extract_images(tar_path, handler.root)
        # Validate images and labels paths exist
        assert os.path.exists(os.path.join(handler.root, "train", "label"))
        assert os.path.exists(os.path.join(handler.root, "train", "lidar"))
        assert os.path.exists(os.path.join(handler.root, "val", "label"))
        assert os.path.exists(os.path.join(handler.root, "val", "lidar"))
        msg = "Upload successful"
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


def pose_classification(tar_path, handler_metadata):
    """
    Default:
    kinetics/nvidia : root_folder_path
    files:
    train_data.npy
    train_label.npy
    val_data.pkl
    val_label.pkl
    ...

    """
    handler = SimpleHandler(handler_metadata)
    try:
        print("Extracting images from data tarball file", file=sys.stderr)
        _extract_images(tar_path, handler.root)
        print("Extraction complete", file=sys.stderr)
        # Validate images and labels paths exist
        assert os.path.exists(os.path.join(handler.root, "kinetics")) or os.path.exists(os.path.join(handler.root, "nvidia"))

        model_type = ""
        if os.path.exists(os.path.join(handler.root, "kinetics")):
            model_type = "kinetics"
        elif os.path.exists(os.path.join(handler.root, "nvidia")):
            model_type = "nvidia"

        assert os.path.exists(os.path.join(handler.root, model_type, "train_data.npy")) and os.path.exists(os.path.join(handler.root, model_type, "train_label.pkl")) and os.path.exists(os.path.join(handler.root, model_type, "val_data.npy")) and os.path.exists(os.path.join(handler.root, model_type, "val_label.pkl"))

        msg = "Upload successful"
        print("Returning sucess code to the api call", file=sys.stderr)
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


def re_identification(tar_path, handler_metadata):
    """
    Default:
    sample_train
    sample_test
    sample_test
    ...

    """
    handler = SimpleHandler(handler_metadata)
    try:
        print("Extracting images from data tarball file", file=sys.stderr)
        _extract_images(tar_path, handler.root)
        print("Extraction complete", file=sys.stderr)
        # Validate images and labels paths exist
        assert os.path.exists(os.path.join(handler.root, "sample_train")) and os.path.exists(os.path.join(handler.root, "sample_test")) and os.path.exists(os.path.join(handler.root, "sample_query"))

        msg = "Upload successful"
        print("Returning sucess code to the api call", file=sys.stderr)
        return Code(201, {}, msg)

    except:
        msg = "Invalid tar file / tar file with no images and/or labels directory"
        return Code(400, {}, msg)


DS_UPLOAD_TO_FUNCTIONS = {"object_detection": object_detection,
                          "semantic_segmentation": semantic_segmentation,
                          "character_recognition": character_recognition,
                          "image_classification": image_classification,
                          "instance_segmentation": instance_segmentation,
                          "bpnet": bpnet,
                          "fpenet": fpenet,
                          "action_recognition": action_recognition,
                          "ml_recog": ml_recog,
                          "ocdnet": ocdnet,
                          "ocrnet": ocrnet,
                          "optical_inspection": optical_inspection,
                          "pointpillars": pointpillars,
                          "pose_classification": pose_classification,
                          "re_identification": re_identification}
