#
# Copyright (c) 2021, Neptune Labs Sp. z o.o.
#
# 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.
#

"""Remember to set environment values:
        # self.exp[SYSTEM_TAGS_ATTRIBUTE_PATH].pop('tag2_to_remove')
        # self.exp[SYSTEM_TAGS_ATTRIBUTE_PATH].pop('tag4_remove_non_existing')
* NEPTUNE_API_TOKEN
* NEPTUNE_PROJECT
"""
import sys
from datetime import datetime

from PIL import Image

import neptune.new as neptune
from alpha_integration_dev.common_client_code import ClientFeatures
from neptune.new.attributes.constants import (
    ARTIFACT_ATTRIBUTE_SPACE,
    LOG_ATTRIBUTE_SPACE,
    PROPERTIES_ATTRIBUTE_SPACE,
    SOURCE_CODE_FILES_ATTRIBUTE_PATH,
    SYSTEM_TAGS_ATTRIBUTE_PATH,
)
from neptune.new.exceptions import MissingFieldException
from neptune.new.types import File


class NewClientFeatures(ClientFeatures):
    def __init__(self):
        super().__init__()
        self.exp = neptune.init(
            source_files="alpha_integration_dev/*.py",
        )

        # download sources
        self.exp.sync()
        with self.with_check_if_file_appears("files.zip"):
            self.exp[SOURCE_CODE_FILES_ATTRIBUTE_PATH].download()

    def modify_tags(self):
        self.exp[SYSTEM_TAGS_ATTRIBUTE_PATH].add("tag1")
        self.exp[SYSTEM_TAGS_ATTRIBUTE_PATH].add(["tag2_to_remove", "tag3"])
        self.exp[SYSTEM_TAGS_ATTRIBUTE_PATH].remove("tag2_to_remove")
        self.exp[SYSTEM_TAGS_ATTRIBUTE_PATH].remove("tag4_remove_non_existing")
        # del self.exp[SYSTEM_TAGS_ATTRIBUTE_PATH]  # TODO: NPT-9222

        self.exp.sync()
        assert self.exp[SYSTEM_TAGS_ATTRIBUTE_PATH].fetch() == {"tag1", "tag3"}

    def modify_properties(self):
        self.exp[PROPERTIES_ATTRIBUTE_SPACE]["prop"] = "some text"
        self.exp[PROPERTIES_ATTRIBUTE_SPACE]["prop_number"] = 42
        self.exp[PROPERTIES_ATTRIBUTE_SPACE]["nested/prop"] = 42
        self.exp[PROPERTIES_ATTRIBUTE_SPACE]["prop_to_del"] = 42
        self.exp[PROPERTIES_ATTRIBUTE_SPACE]["prop_list"] = [1, 2, 3]
        with open(self.text_file_path, mode="r") as f:
            self.exp[PROPERTIES_ATTRIBUTE_SPACE]["prop_IO"] = f
        self.exp[PROPERTIES_ATTRIBUTE_SPACE]["prop_datetime"] = datetime.now()
        self.exp.sync()
        del self.exp[PROPERTIES_ATTRIBUTE_SPACE]["prop_to_del"]

        assert self.exp[PROPERTIES_ATTRIBUTE_SPACE]["prop"].fetch() == "some text"
        assert self.exp[PROPERTIES_ATTRIBUTE_SPACE]["prop_number"].fetch() == 42
        assert self.exp[PROPERTIES_ATTRIBUTE_SPACE]["nested/prop"].fetch() == 42
        prop_to_del_absent = False
        try:
            self.exp[PROPERTIES_ATTRIBUTE_SPACE]["prop_to_del"].fetch()
        except AttributeError:
            prop_to_del_absent = True
        assert prop_to_del_absent

    def log_std(self):
        print("stdout text1")
        print("stdout text2")
        print("stderr text1", file=sys.stderr)
        print("stderr text2", file=sys.stderr)

    def log_series(self):
        # floats
        self.exp[LOG_ATTRIBUTE_SPACE]["m1"].log(1)
        self.exp[LOG_ATTRIBUTE_SPACE]["m1"].log(2)
        self.exp[LOG_ATTRIBUTE_SPACE]["m1"].log(3)
        self.exp[LOG_ATTRIBUTE_SPACE]["m1"].log(2)
        self.exp[LOG_ATTRIBUTE_SPACE]["nested"]["m1"].log(1)

        # texts
        self.exp[LOG_ATTRIBUTE_SPACE]["m2"].log("a")
        self.exp[LOG_ATTRIBUTE_SPACE]["m2"].log("b")
        self.exp[LOG_ATTRIBUTE_SPACE]["m2"].log("c")

        # images
        im_frame = Image.open(self.img_path)
        g_img = File.as_image(im_frame)
        self.exp[LOG_ATTRIBUTE_SPACE]["g_img"].log(g_img)

    def handle_files_and_images(self):
        # image
        im_frame = Image.open(self.img_path)
        g_img = File.as_image(im_frame)
        self.exp[ARTIFACT_ATTRIBUTE_SPACE]["assigned image"] = g_img
        self.exp.wait()
        with self.with_check_if_file_appears("assigned image.png"):
            self.exp[ARTIFACT_ATTRIBUTE_SPACE]["assigned image"].download()
        with self.with_check_if_file_appears("custom_dest.png"):
            self.exp[ARTIFACT_ATTRIBUTE_SPACE]["assigned image"].download("custom_dest.png")

        self.exp[ARTIFACT_ATTRIBUTE_SPACE]["logged image"].log(g_img)

        with open(self.img_path, mode="r") as f:
            # self.exp[ARTIFACT_ATTRIBUTE_SPACE]['assigned image stream'] = f
            self.exp[ARTIFACT_ATTRIBUTE_SPACE]["logged image stream"].log(f)

        # artifact
        text_file = neptune.types.File(self.text_file_path)
        self.exp[ARTIFACT_ATTRIBUTE_SPACE]["assigned file"] = text_file
        # self.exp[ARTIFACT_ATTRIBUTE_SPACE]['logged file'].log(text_file)  # wrong type
        with open(self.text_file_path, mode="r") as f:
            self.exp[ARTIFACT_ATTRIBUTE_SPACE]["assigned file stream"] = f
            self.exp[ARTIFACT_ATTRIBUTE_SPACE]["logged file stream"].log(f)

    def handle_directories(self):
        pass

    def finalize(self):
        return


if __name__ == "__main__":
    NewClientFeatures().run()
