#
# Copyright (c) 2022, 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.
#
import abc
from dataclasses import dataclass
from datetime import datetime
from typing import (
    TYPE_CHECKING,
    Generic,
    List,
    Optional,
    Set,
    Type,
    TypeVar,
)

from neptune.core.components.operation_storage import OperationStorage
from neptune.exceptions import MalformedOperation
from neptune.internal.container_type import ContainerType

if TYPE_CHECKING:
    from neptune.attributes.attribute import Attribute
    from neptune.internal.backends.neptune_backend import NeptuneBackend
    from neptune.internal.operation_visitor import OperationVisitor

Ret = TypeVar("Ret")
T = TypeVar("T")


def all_subclasses(cls):
    return set(cls.__subclasses__()).union([s for c in cls.__subclasses__() for s in all_subclasses(c)])


@dataclass
class Operation(abc.ABC):

    path: List[str]

    @abc.abstractmethod
    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        pass

    def clean(self, operation_storage: OperationStorage):
        pass

    def to_dict(self) -> dict:
        return {"type": self.__class__.__name__, "path": self.path}

    @staticmethod
    def from_dict(data: dict) -> "Operation":
        if "type" not in data:
            raise ValueError("Malformed operation {} - type is missing".format(data))
        sub_classes = {cls.__name__: cls for cls in all_subclasses(Operation)}
        if not data["type"] in sub_classes:
            raise ValueError("Malformed operation {} - unknown type {}".format(data, data["type"]))
        return sub_classes[data["type"]].from_dict(data)


@dataclass
class AssignFloat(Operation):

    value: float

    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_assign_float(self)

    def to_dict(self) -> dict:
        ret = super().to_dict()
        ret["value"] = self.value
        return ret

    @staticmethod
    def from_dict(data: dict) -> "AssignFloat":
        return AssignFloat(data["path"], data["value"])


@dataclass
class AssignInt(Operation):

    value: int

    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_assign_int(self)

    def to_dict(self) -> dict:
        ret = super().to_dict()
        ret["value"] = self.value
        return ret

    @staticmethod
    def from_dict(data: dict) -> "AssignInt":
        return AssignInt(data["path"], data["value"])


@dataclass
class AssignBool(Operation):

    value: bool

    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_assign_bool(self)

    def to_dict(self) -> dict:
        ret = super().to_dict()
        ret["value"] = self.value
        return ret

    @staticmethod
    def from_dict(data: dict) -> "AssignBool":
        return AssignBool(data["path"], data["value"])


@dataclass
class AssignString(Operation):

    value: str

    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_assign_string(self)

    def to_dict(self) -> dict:
        ret = super().to_dict()
        ret["value"] = self.value
        return ret

    @staticmethod
    def from_dict(data: dict) -> "AssignString":
        return AssignString(data["path"], data["value"])


@dataclass
class AssignDatetime(Operation):

    value: datetime

    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_assign_datetime(self)

    def to_dict(self) -> dict:
        ret = super().to_dict()
        ret["value"] = int(1000 * self.value.timestamp())
        return ret

    @staticmethod
    def from_dict(data: dict) -> "AssignDatetime":
        return AssignDatetime(data["path"], datetime.fromtimestamp(data["value"] / 1000))


class LogOperation(Operation, abc.ABC):
    pass


@dataclass
class LogSeriesValue(Generic[T]):

    value: T
    step: Optional[float]
    ts: float

    def to_dict(self, value_serializer=lambda x: x) -> dict:
        return {"value": value_serializer(self.value), "step": self.step, "ts": self.ts}

    @staticmethod
    def from_dict(data: dict, value_deserializer=lambda x: x) -> "LogSeriesValue[T]":
        return LogSeriesValue[T](value_deserializer(data["value"]), data.get("step", None), data["ts"])


@dataclass
class LogFloats(LogOperation):

    ValueType = LogSeriesValue[float]

    values: List[ValueType]

    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_log_floats(self)

    def to_dict(self) -> dict:
        ret = super().to_dict()
        ret["values"] = [value.to_dict() for value in self.values]
        return ret

    @staticmethod
    def from_dict(data: dict) -> "LogFloats":
        return LogFloats(
            data["path"],
            [LogFloats.ValueType.from_dict(value) for value in data["values"]],
        )


@dataclass
class LogStrings(LogOperation):

    ValueType = LogSeriesValue[str]

    values: List[ValueType]

    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_log_strings(self)

    def to_dict(self) -> dict:
        ret = super().to_dict()
        ret["values"] = [value.to_dict() for value in self.values]
        return ret

    @staticmethod
    def from_dict(data: dict) -> "LogStrings":
        return LogStrings(
            data["path"],
            [LogStrings.ValueType.from_dict(value) for value in data["values"]],
        )


@dataclass
class ClearFloatLog(Operation):
    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_clear_float_log(self)

    @staticmethod
    def from_dict(data: dict) -> "ClearFloatLog":
        return ClearFloatLog(data["path"])


@dataclass
class ClearStringLog(Operation):
    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_clear_string_log(self)

    @staticmethod
    def from_dict(data: dict) -> "ClearStringLog":
        return ClearStringLog(data["path"])


@dataclass
class ConfigFloatSeries(Operation):

    min: Optional[float]
    max: Optional[float]
    unit: Optional[str]

    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_config_float_series(self)

    def to_dict(self) -> dict:
        ret = super().to_dict()
        ret["min"] = self.min
        ret["max"] = self.max
        ret["unit"] = self.unit
        return ret

    @staticmethod
    def from_dict(data: dict) -> "ConfigFloatSeries":
        return ConfigFloatSeries(data["path"], data["min"], data["max"], data["unit"])


@dataclass
class AddStrings(Operation):

    values: Set[str]

    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_add_strings(self)

    def to_dict(self) -> dict:
        ret = super().to_dict()
        ret["values"] = list(self.values)
        return ret

    @staticmethod
    def from_dict(data: dict) -> "AddStrings":
        return AddStrings(data["path"], set(data["values"]))


@dataclass
class RemoveStrings(Operation):

    values: Set[str]

    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_remove_strings(self)

    def to_dict(self) -> dict:
        ret = super().to_dict()
        ret["values"] = list(self.values)
        return ret

    @staticmethod
    def from_dict(data: dict) -> "RemoveStrings":
        return RemoveStrings(data["path"], set(data["values"]))


@dataclass
class ClearStringSet(Operation):
    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_clear_string_set(self)

    @staticmethod
    def from_dict(data: dict) -> "ClearStringSet":
        return ClearStringSet(data["path"])


@dataclass
class DeleteAttribute(Operation):
    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_delete_attribute(self)

    @staticmethod
    def from_dict(data: dict) -> "DeleteAttribute":
        return DeleteAttribute(data["path"])


@dataclass
class CopyAttribute(Operation):
    container_id: str
    container_type: ContainerType
    source_path: List[str]
    source_attr_cls: Type["Attribute"]

    def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
        return visitor.visit_copy_attribute(self)

    def to_dict(self) -> dict:
        ret = super().to_dict()
        ret["container_id"] = self.container_id
        ret["container_type"] = self.container_type.value
        ret["source_path"] = self.source_path
        ret["source_attr_name"] = self.source_attr_cls.__name__
        return ret

    @staticmethod
    def from_dict(data: dict) -> "CopyAttribute":
        from neptune.attributes.attribute import Attribute

        source_attr_cls = {cls.__name__: cls for cls in all_subclasses(Attribute) if cls.supports_copy}.get(
            data["source_attr_name"]
        )

        if source_attr_cls is None:
            raise MalformedOperation("Copy of non-copiable type found in queue!")

        return CopyAttribute(
            data["path"],
            data["container_id"],
            ContainerType(data["container_type"]),
            data["source_path"],
            source_attr_cls,
        )

    def resolve(self, backend: "NeptuneBackend") -> Operation:
        # repack CopyAttribute op into target attribute assignment
        getter = self.source_attr_cls.getter
        create_assignment_operation = self.source_attr_cls.create_assignment_operation
        value = getter(backend, self.container_id, self.container_type, self.source_path)
        return create_assignment_operation(self.path, value)
