from datetime import datetime as _datetime

import pytz as _pytz
import six as _six
from flyteidl.core import literals_pb2 as _literals_pb2
from google.protobuf.struct_pb2 import Struct

from flytekit.common.exceptions import user as _user_exceptions
from flytekit.models import common as _common
from flytekit.models.core import types as _core_types
from flytekit.models.types import OutputReference as _OutputReference
from flytekit.models.types import SchemaType as _SchemaType


class RetryStrategy(_common.FlyteIdlEntity):
    def __init__(self, retries):
        """
        :param int retries: Number of retries to attempt on recoverable failures.  If retries is 0, then
            only one attempt will be made.
        """
        self._retries = retries

    @property
    def retries(self):
        """
        Number of retries to attempt on recoverable failures.  If retries is 0, then only one attempt will be made.
        :rtype: int
        """
        return self._retries

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.RetryStrategy
        """
        return _literals_pb2.RetryStrategy(retries=self.retries)

    @classmethod
    def from_flyte_idl(cls, pb2_object):
        """
        :param flyteidl.core.literals_pb2.RetryStrategy pb2_object:
        :rtype: RetryStrategy
        """
        return cls(retries=pb2_object.retries)


class Primitive(_common.FlyteIdlEntity):
    def __init__(
        self,
        integer=None,
        float_value=None,
        string_value=None,
        boolean=None,
        datetime=None,
        duration=None,
    ):
        """
        This object proxies the primitives supported by the Flyte IDL system.  Only one value can be set.
        :param int integer: [Optional]
        :param float float_value: [Optional]
        :param Text string_value: [Optional]
        :param bool boolean: [Optional]
        :param datetime.timestamp datetime: [Optional]
        :param datetime.timedelta duration: [Optional]
        """
        self._integer = integer
        self._float_value = float_value
        self._string_value = string_value
        self._boolean = boolean
        if datetime is None:
            self._datetime = None
        elif isinstance(datetime, _datetime):
            self._datetime = datetime
        else:  # TODO Check for timestamp type?
            self._datetime = _datetime.utcfromtimestamp(datetime.seconds)
        self._duration = duration

    @property
    def integer(self):
        """
        :rtype: int
        """
        return self._integer

    @property
    def float_value(self):
        """
        :rtype: float
        """
        return self._float_value

    @property
    def string_value(self):
        """
        :rtype: Text
        """
        return self._string_value

    @property
    def boolean(self):
        """
        :rtype: bool
        """
        return self._boolean

    @property
    def datetime(self):
        """
        :rtype: datetime.datetime
        """
        if self._datetime is None or self._datetime.tzinfo is not None:
            return self._datetime
        return self._datetime.replace(tzinfo=_pytz.UTC)

    @property
    def duration(self):
        """
        :rtype: datetime.timedelta
        """
        return self._duration

    @property
    def value(self):
        """
        This returns whichever field is set.
        :rtype: T
        """
        for value in [
            self.integer,
            self.float_value,
            self.string_value,
            self.boolean,
            self.datetime,
            self.duration,
        ]:
            if value is not None:
                return value

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.Primitive
        """
        primitive = _literals_pb2.Primitive(
            integer=self.integer,
            float_value=self.float_value,
            string_value=self.string_value,
            boolean=self.boolean,
        )
        if self.datetime is not None:
            # Convert to UTC and remove timezone so protobuf behaves.
            primitive.datetime.FromDatetime(self.datetime.astimezone(_pytz.UTC).replace(tzinfo=None))
        if self.duration is not None:
            primitive.duration.FromTimedelta(self.duration)
        return primitive

    @classmethod
    def from_flyte_idl(cls, proto):
        """
        :param flyteidl.core.literals_pb2.Primitive proto:
        :rtype: Primitive
        """
        return cls(
            integer=proto.integer if proto.HasField("integer") else None,
            float_value=proto.float_value if proto.HasField("float_value") else None,
            string_value=proto.string_value if proto.HasField("string_value") else None,
            boolean=proto.boolean if proto.HasField("boolean") else None,
            datetime=proto.datetime.ToDatetime().replace(tzinfo=_pytz.UTC) if proto.HasField("datetime") else None,
            duration=proto.duration.ToTimedelta() if proto.HasField("duration") else None,
        )


class Binary(_common.FlyteIdlEntity):
    def __init__(self, value, tag):
        """
        :param bytes value:
        :param Text tag:
        """
        self._value = value
        self._tag = tag

    @property
    def value(self):
        """
        :rtype: bytes
        """
        return self._value

    @property
    def tag(self):
        """
        :rtype: Text
        """
        return self._tag

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.Binary
        """
        return _literals_pb2.Binary(value=self.value, tag=self.tag)

    @classmethod
    def from_flyte_idl(cls, pb2_object):
        """
        :param flyteidl.core.literals_pb2.Binary pb2_object:
        :rtype: Binary
        """
        return cls(value=pb2_object.value, tag=pb2_object.tag)


class BlobMetadata(_common.FlyteIdlEntity):
    def __init__(self, type):
        """
        :param flytekit.models.core.types.BlobType type: The type of the underlying blob
        """
        self._type = type

    @property
    def type(self):
        """
        :rtype: flytekit.models.core.types.BlobType
        """
        return self._type

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.BlobMetadata
        """
        return _literals_pb2.BlobMetadata(type=self.type.to_flyte_idl())

    @classmethod
    def from_flyte_idl(cls, proto):
        """
        :param flyteidl.core.literals_pb2.BlobMetadata proto:
        :rtype: BlobMetadata
        """
        return cls(type=_core_types.BlobType.from_flyte_idl(proto.type))


class Blob(_common.FlyteIdlEntity):
    def __init__(self, metadata, uri):
        """
        :param BlobMetadata metadata:
        :param Text uri: The location of this blob
        """
        self._metadata = metadata
        self._uri = uri

    @property
    def uri(self):
        """
        :rtype: Text
        """
        return self._uri

    @property
    def metadata(self):
        """
        :rtype: BlobMetadata
        """
        return self._metadata

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.Blob
        """
        return _literals_pb2.Blob(metadata=self.metadata.to_flyte_idl(), uri=self.uri)

    @classmethod
    def from_flyte_idl(cls, proto):
        """
        :param flyteidl.core.literals_pb2.Blob proto:
        :rtype: Blob
        """
        return cls(metadata=BlobMetadata.from_flyte_idl(proto.metadata), uri=proto.uri)


class Void(_common.FlyteIdlEntity):
    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.Void
        """
        return _literals_pb2.Void()

    @classmethod
    def from_flyte_idl(cls, proto):
        """
        :param flyteidl.core.literals_pb2.Void proto:
        :rtype: Void
        """
        return cls()


class BindingDataMap(_common.FlyteIdlEntity):
    def __init__(self, bindings):
        """
        A map of BindingData items.  Can be a recursive structure

        :param dict[string, BindingData] bindings: Map of strings to Bindings
        """
        self._bindings = bindings

    @property
    def bindings(self):
        """
        Map of strings to Bindings
        :rtype: dict[string, BindingData]
        """
        return self._bindings

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.BindingDataMap
        """
        return _literals_pb2.BindingDataMap(bindings={k: v.to_flyte_idl() for (k, v) in _six.iteritems(self.bindings)})

    @classmethod
    def from_flyte_idl(cls, pb2_object):
        """
        :param flyteidl.core.literals_pb2.BindingDataMap pb2_object:
        :rtype: flytekit.models.literals.BindingDataMap
        """

        return cls({k: BindingData.from_flyte_idl(v) for (k, v) in _six.iteritems(pb2_object.bindings)})


class BindingDataCollection(_common.FlyteIdlEntity):
    def __init__(self, bindings):
        """
        A list of BindingData items.

        :param list[BindingData] bindings:
        """
        self._bindings = bindings

    @property
    def bindings(self):
        """
        :rtype: list[BindingData]
        """
        return self._bindings

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.BindingDataCollection
        """
        return _literals_pb2.BindingDataCollection(bindings=[b.to_flyte_idl() for b in self.bindings])

    @classmethod
    def from_flyte_idl(cls, pb2_object):
        """
        :param flyteidl.core.literals_pb2.BindingDataCollection pb2_object:
        :rtype: flytekit.models.literals.BindingDataCollection
        """
        return cls([BindingData.from_flyte_idl(b) for b in pb2_object.bindings])


class BindingData(_common.FlyteIdlEntity):
    def __init__(self, scalar=None, collection=None, promise=None, map=None):
        """
        Specifies either a simple value or a reference to another output. Only one of the input arguments may be
        specified.

        :param Scalar scalar: [Optional] A simple scalar value.
        :param BindingDataCollection collection: [Optional] A collection of binding data. This allows nesting of
            binding data to any number of levels.
        :param flytekit.models.types.OutputReference promise: [Optional] References an output promised by another node.
        :param BindingDataMap map: [Optional] A map of bindings. The key is always a string.
        """
        self._scalar = scalar
        self._collection = collection
        self._promise = promise
        self._map = map

    @property
    def scalar(self):
        """
        A simple scalar value.
        :rtype: Scalar
        """
        return self._scalar

    @property
    def collection(self):
        """
        [Optional] A collection of binding data. This allows nesting of binding data to any number of levels.
        :rtype: BindingDataCollection
        """
        return self._collection

    @property
    def promise(self):
        """
        [Optional] References an output promised by another node.
        :rtype: flytekit.models.types.OutputReference
        """
        return self._promise

    @property
    def map(self):
        """
        [Optional] A map of bindings. The key is always a string.
        :rtype: BindingDataMap
        """
        return self._map

    @property
    def value(self):
        """
        Returns whichever value is set
        :rtype: T
        """
        return self.scalar or self.collection or self.promise or self.map

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.BindingData
        """
        return _literals_pb2.BindingData(
            scalar=self.scalar.to_flyte_idl() if self.scalar is not None else None,
            collection=self.collection.to_flyte_idl() if self.collection is not None else None,
            promise=self.promise.to_flyte_idl() if self.promise is not None else None,
            map=self.map.to_flyte_idl() if self.map is not None else None,
        )

    @classmethod
    def from_flyte_idl(cls, pb2_object):
        """
        :param flyteidl.core.literals_pb2.BindingData pb2_object:
        :return: BindingData
        """
        return cls(
            scalar=Scalar.from_flyte_idl(pb2_object.scalar) if pb2_object.HasField("scalar") else None,
            collection=BindingDataCollection.from_flyte_idl(pb2_object.collection)
            if pb2_object.HasField("collection")
            else None,
            promise=_OutputReference.from_flyte_idl(pb2_object.promise) if pb2_object.HasField("promise") else None,
            map=BindingDataMap.from_flyte_idl(pb2_object.map) if pb2_object.HasField("map") else None,
        )

    def to_literal_model(self):
        """
        Converts current binding data into a Literal asserting that there are no promises in the bindings.
        :rtype: Literal
        """
        if self.promise:
            raise _user_exceptions.FlyteValueException(
                self.promise,
                "Cannot convert BindingData to a Literal because " "it has a promise.",
            )
        elif self.scalar:
            return Literal(scalar=self.scalar)
        elif self.collection:
            return Literal(
                collection=LiteralCollection(
                    literals=[binding.to_literal_model() for binding in self.collection.bindings]
                )
            )
        elif self.map:
            return Literal(map=LiteralMap(literals={k: binding.to_literal_model() for k, binding in self.map.bindings}))


class Binding(_common.FlyteIdlEntity):
    def __init__(self, var, binding):
        """
        An input/output binding of a variable to either static value or a node output.

        :param Text var: A variable name, must match an input or output variable of the node.
        :param BindingData binding: Data to use to bind this variable.
        """
        self._var = var
        self._binding = binding

    @property
    def var(self):
        """
        A variable name, must match an input or output variable of the node.
        :rtype: Text
        """
        return self._var

    @property
    def binding(self):
        """
        Data to use to bind this variable.
        :rtype: BindingData
        """
        return self._binding

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.Binding
        """
        return _literals_pb2.Binding(var=self.var, binding=self.binding.to_flyte_idl())

    @classmethod
    def from_flyte_idl(cls, pb2_object):
        """
        :param flyteidl.core.literals_pb2.Binding pb2_object:
        :rtype: flytekit.core.models.literals.Binding
        """
        return cls(pb2_object.var, BindingData.from_flyte_idl(pb2_object.binding))


class Schema(_common.FlyteIdlEntity):
    def __init__(self, uri, type):
        """
        A strongly typed schema that defines the interface of data retrieved from the underlying storage medium.

        :param Text uri:
        :param flytekit.models.types.SchemaType type:
        """
        self._uri = uri
        self._type = type

    @property
    def uri(self):
        """
        :rtype: Text
        """
        return self._uri

    @property
    def type(self):
        """
        :rtype: flytekit.models.types.SchemaType
        """
        return self._type

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.Schema
        """
        return _literals_pb2.Schema(uri=self.uri, type=self.type.to_flyte_idl())

    @classmethod
    def from_flyte_idl(cls, pb2_object):
        """
        :param flyteidl.core.literals_pb2.Schema pb2_object:
        :rtype: Schema
        """
        return cls(uri=pb2_object.uri, type=_SchemaType.from_flyte_idl(pb2_object.type))


class LiteralCollection(_common.FlyteIdlEntity):
    def __init__(self, literals):
        """
        :param list[Literal] literals: underlying list of literals in this collection.
        """
        self._literals = literals

    @property
    def literals(self):
        """
        :rtype: list[Literal]
        """
        return self._literals

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.LiteralCollection
        """
        return _literals_pb2.LiteralCollection(literals=[l.to_flyte_idl() for l in self.literals])

    @classmethod
    def from_flyte_idl(cls, pb2_object):
        """
        :param flyteidl.core.literals_pb2.LiteralCollection pb2_object:
        :rtype: LiteralCollection
        """
        return cls([Literal.from_flyte_idl(l) for l in pb2_object.literals])


class LiteralMap(_common.FlyteIdlEntity):
    def __init__(self, literals):
        """
        :param dict[Text, Literal] literals: A dictionary mapping Text key names to Literal objects.
        """
        self._literals = literals

    @property
    def literals(self):
        """
        A dictionary mapping Text key names to Literal objects.
        :rtype: dict[Text, Literal]
        """
        return self._literals

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.LiteralMap
        """
        return _literals_pb2.LiteralMap(literals={k: v.to_flyte_idl() for k, v in _six.iteritems(self.literals)})

    @classmethod
    def from_flyte_idl(cls, pb2_object):
        """
        :param flyteidl.core.literals_pb2.LiteralMap pb2_object:
        :rtype: LiteralMap
        """
        return cls({k: Literal.from_flyte_idl(v) for k, v in _six.iteritems(pb2_object.literals)})


class Scalar(_common.FlyteIdlEntity):
    def __init__(
        self,
        primitive: Primitive = None,
        blob: Blob = None,
        binary: Binary = None,
        schema: Schema = None,
        none_type: Void = None,
        error=None,
        generic: Struct = None,
    ):
        """
        Scalar wrapper around Flyte types.  Only one can be specified.

        :param Primitive primitive:
        :param Blob blob:
        :param Binary binary:
        :param Schema schema:
        :param Void none_type:
        :param error:
        :param google.protobuf.struct_pb2.Struct generic:
        """

        self._primitive = primitive
        self._blob = blob
        self._binary = binary
        self._schema = schema
        self._none_type = none_type
        self._error = error
        self._generic = generic

    @property
    def primitive(self):
        """
        :rtype: Primitive
        """
        return self._primitive

    @property
    def blob(self):
        """
        :rtype: Blob
        """
        return self._blob

    @property
    def binary(self):
        """
        :rtype: Binary
        """
        return self._binary

    @property
    def schema(self):
        """
        :rtype: Schema
        """
        return self._schema

    @property
    def none_type(self):
        """
        :rtype: Void
        """
        return self._none_type

    @property
    def error(self):
        """
        :rtype: TODO
        """
        return self._error

    @property
    def generic(self):
        """
        :rtype: google.protobuf.struct_pb2.Struct
        """
        return self._generic

    @property
    def value(self):
        """
        Returns whichever value is set
        :rtype: T
        """
        return self.primitive or self.blob or self.binary or self.schema or self.none_type or self.error

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.Scalar
        """
        return _literals_pb2.Scalar(
            primitive=self.primitive.to_flyte_idl() if self.primitive is not None else None,
            blob=self.blob.to_flyte_idl() if self.blob is not None else None,
            binary=self.binary.to_flyte_idl() if self.binary is not None else None,
            schema=self.schema.to_flyte_idl() if self.schema is not None else None,
            none_type=self.none_type.to_flyte_idl() if self.none_type is not None else None,
            error=self.error if self.error is not None else None,
            generic=self.generic,
        )

    @classmethod
    def from_flyte_idl(cls, pb2_object):
        """
        :param flyteidl.core.literals_pb2.Scalar pb2_object:
        :rtype: flytekit.models.literals.Scalar
        """
        # todo finish
        return cls(
            primitive=Primitive.from_flyte_idl(pb2_object.primitive) if pb2_object.HasField("primitive") else None,
            blob=Blob.from_flyte_idl(pb2_object.blob) if pb2_object.HasField("blob") else None,
            binary=Binary.from_flyte_idl(pb2_object.binary) if pb2_object.HasField("binary") else None,
            schema=Schema.from_flyte_idl(pb2_object.schema) if pb2_object.HasField("schema") else None,
            none_type=Void.from_flyte_idl(pb2_object.none_type) if pb2_object.HasField("none_type") else None,
            error=pb2_object.error if pb2_object.HasField("error") else None,
            generic=pb2_object.generic if pb2_object.HasField("generic") else None,
        )


class Literal(_common.FlyteIdlEntity):
    def __init__(self, scalar: Scalar = None, collection: LiteralCollection = None, map: LiteralMap = None):
        """
        :param Scalar scalar:
        :param LiteralCollection collection:
        :param LiteralMap map:
        """
        self._scalar = scalar
        self._collection = collection
        self._map = map

    @property
    def scalar(self):
        """
        If not None, this value holds a scalar value which can be further unpacked.
        :rtype: Scalar
        """
        return self._scalar

    @property
    def collection(self):
        """
        If not None, this value holds a collection of Literal values which can be further unpacked.
        :rtype: LiteralCollection
        """
        return self._collection

    @property
    def map(self):
        """
        If not None, this value holds a map of Literal values which can be further unpacked.
        :rtype: LiteralMap
        """
        return self._map

    @property
    def value(self):
        """
        Returns one of the scalar, collection, or map properties based on which one is set.
        :rtype: T
        """
        return self.scalar or self.collection or self.map

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.core.literals_pb2.Literal
        """
        return _literals_pb2.Literal(
            scalar=self.scalar.to_flyte_idl() if self.scalar is not None else None,
            collection=self.collection.to_flyte_idl() if self.collection is not None else None,
            map=self.map.to_flyte_idl() if self.map is not None else None,
        )

    @classmethod
    def from_flyte_idl(cls, pb2_object):
        """
        :param flyteidl.core.literals_pb2.Literal pb2_object:
        :rtype: Literal
        """
        collection = None
        if pb2_object.HasField("collection"):
            collection = LiteralCollection.from_flyte_idl(pb2_object.collection)

        return cls(
            scalar=Scalar.from_flyte_idl(pb2_object.scalar) if pb2_object.HasField("scalar") else None,
            collection=collection,
            map=LiteralMap.from_flyte_idl(pb2_object.map) if pb2_object.HasField("map") else None,
        )
