# Copyright 2022 The Feast 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
#
#     https://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.
from abc import ABC, abstractmethod
from datetime import datetime, timezone
from enum import Enum
from typing import Dict, List, Union

import pyarrow

from feast.value_type import ValueType

PRIMITIVE_FEAST_TYPES_TO_VALUE_TYPES = {
    "INVALID": "UNKNOWN",
    "BYTES": "BYTES",
    "PDF_BYTES": "PDF_BYTES",
    "IMAGE_BYTES": "IMAGE_BYTES",
    "STRING": "STRING",
    "INT32": "INT32",
    "INT64": "INT64",
    "FLOAT64": "DOUBLE",
    "FLOAT32": "FLOAT",
    "BOOL": "BOOL",
    "UNIX_TIMESTAMP": "UNIX_TIMESTAMP",
    "MAP": "MAP",
}


def _utc_now() -> datetime:
    return datetime.now(tz=timezone.utc)


class ComplexFeastType(ABC):
    """
    A ComplexFeastType represents a structured type that is recognized by Feast.
    """

    def __init__(self):
        """Creates a ComplexFeastType object."""
        pass

    @abstractmethod
    def to_value_type(self) -> ValueType:
        """
        Converts a ComplexFeastType object to the corresponding ValueType enum.
        """
        raise NotImplementedError

    def __hash__(self):
        return hash(self.to_value_type().value)

    def __eq__(self, other):
        if isinstance(other, ComplexFeastType):
            return self.to_value_type() == other.to_value_type()
        else:
            return False


class PrimitiveFeastType(Enum):
    """
    A PrimitiveFeastType represents a primitive type in Feast.

    Note that these values must match the values in /feast/protos/types/Value.proto.
    """

    INVALID = 0
    BYTES = 1
    STRING = 2
    INT32 = 3
    INT64 = 4
    FLOAT64 = 5
    FLOAT32 = 6
    BOOL = 7
    UNIX_TIMESTAMP = 8
    PDF_BYTES = 9
    IMAGE_BYTES = 10
    MAP = 11

    def to_value_type(self) -> ValueType:
        """
        Converts a PrimitiveFeastType object to the corresponding ValueType enum.
        """
        value_type_name = PRIMITIVE_FEAST_TYPES_TO_VALUE_TYPES[self.name]
        return ValueType[value_type_name]

    def __str__(self):
        return PRIMITIVE_FEAST_TYPES_TO_STRING[self.name]

    def __eq__(self, other):
        if isinstance(other, PrimitiveFeastType):
            return self.value == other.value
        else:
            return False

    def __hash__(self):
        return hash((PRIMITIVE_FEAST_TYPES_TO_STRING[self.name]))


Invalid = PrimitiveFeastType.INVALID
Bytes = PrimitiveFeastType.BYTES
PdfBytes = PrimitiveFeastType.PDF_BYTES
ImageBytes = PrimitiveFeastType.IMAGE_BYTES
String = PrimitiveFeastType.STRING
Bool = PrimitiveFeastType.BOOL
Int32 = PrimitiveFeastType.INT32
Int64 = PrimitiveFeastType.INT64
Float32 = PrimitiveFeastType.FLOAT32
Float64 = PrimitiveFeastType.FLOAT64
UnixTimestamp = PrimitiveFeastType.UNIX_TIMESTAMP
Map = PrimitiveFeastType.MAP

SUPPORTED_BASE_TYPES = [
    Invalid,
    String,
    Bytes,
    PdfBytes,
    ImageBytes,
    Bool,
    Int32,
    Int64,
    Float32,
    Float64,
    UnixTimestamp,
    Map,
]

PRIMITIVE_FEAST_TYPES_TO_STRING = {
    "INVALID": "Invalid",
    "STRING": "String",
    "BYTES": "Bytes",
    "PDF_BYTES": "PdfBytes",
    "IMAGE_BYTES": "ImageBytes",
    "BOOL": "Bool",
    "INT32": "Int32",
    "INT64": "Int64",
    "FLOAT32": "Float32",
    "FLOAT64": "Float64",
    "UNIX_TIMESTAMP": "UnixTimestamp",
    "MAP": "Map",
}


class Array(ComplexFeastType):
    """
    An Array represents a list of types.

    Attributes:
        base_type: The base type of the array.
    """

    base_type: Union[PrimitiveFeastType, ComplexFeastType]

    def __init__(self, base_type: Union[PrimitiveFeastType, ComplexFeastType]):
        if base_type not in SUPPORTED_BASE_TYPES:
            raise ValueError(
                f"Type {type(base_type)} is currently not supported as a base type for Array."
            )

        self.base_type = base_type

    def to_value_type(self) -> ValueType:
        assert isinstance(self.base_type, PrimitiveFeastType)
        value_type_name = PRIMITIVE_FEAST_TYPES_TO_VALUE_TYPES[self.base_type.name]
        value_type_list_name = value_type_name + "_LIST"
        return ValueType[value_type_list_name]

    def __str__(self):
        return f"Array({self.base_type})"


FeastType = Union[ComplexFeastType, PrimitiveFeastType]

VALUE_TYPES_TO_FEAST_TYPES: Dict["ValueType", FeastType] = {
    ValueType.UNKNOWN: Invalid,
    ValueType.BYTES: Bytes,
    ValueType.PDF_BYTES: PdfBytes,
    ValueType.IMAGE_BYTES: ImageBytes,
    ValueType.STRING: String,
    ValueType.INT32: Int32,
    ValueType.INT64: Int64,
    ValueType.DOUBLE: Float64,
    ValueType.FLOAT: Float32,
    ValueType.BOOL: Bool,
    ValueType.UNIX_TIMESTAMP: UnixTimestamp,
    ValueType.BYTES_LIST: Array(Bytes),
    ValueType.STRING_LIST: Array(String),
    ValueType.INT32_LIST: Array(Int32),
    ValueType.INT64_LIST: Array(Int64),
    ValueType.DOUBLE_LIST: Array(Float64),
    ValueType.FLOAT_LIST: Array(Float32),
    ValueType.BOOL_LIST: Array(Bool),
    ValueType.UNIX_TIMESTAMP_LIST: Array(UnixTimestamp),
    ValueType.MAP: Map,
    ValueType.MAP_LIST: Array(Map),
}

FEAST_TYPES_TO_PYARROW_TYPES = {
    String: pyarrow.string(),
    Bool: pyarrow.bool_(),
    Int32: pyarrow.int32(),
    Int64: pyarrow.int64(),
    Float32: pyarrow.float32(),
    Float64: pyarrow.float64(),
    # Note: datetime only supports microseconds https://github.com/python/cpython/blob/3.8/Lib/datetime.py#L1559
    UnixTimestamp: pyarrow.timestamp("us", tz=_utc_now().tzname()),
}

FEAST_VECTOR_TYPES: List[Union[ValueType, PrimitiveFeastType, ComplexFeastType]] = [
    ValueType.BYTES_LIST,
    ValueType.INT32_LIST,
    ValueType.INT64_LIST,
    ValueType.FLOAT_LIST,
    ValueType.BOOL_LIST,
    ValueType.MAP_LIST,
]
for k in VALUE_TYPES_TO_FEAST_TYPES:
    if k in FEAST_VECTOR_TYPES:
        FEAST_VECTOR_TYPES.append(VALUE_TYPES_TO_FEAST_TYPES[k])


def from_feast_to_pyarrow_type(feast_type: FeastType) -> pyarrow.DataType:
    """
    Converts a Feast type to a PyArrow type.

    Args:
        feast_type: The Feast type to be converted.

    Raises:
        ValueError: The conversion could not be performed.
    """
    assert isinstance(feast_type, (ComplexFeastType, PrimitiveFeastType)), (
        f"Expected FeastType, got {type(feast_type)}"
    )
    if isinstance(feast_type, PrimitiveFeastType):
        if feast_type in FEAST_TYPES_TO_PYARROW_TYPES:
            return FEAST_TYPES_TO_PYARROW_TYPES[feast_type]
    elif isinstance(feast_type, ComplexFeastType):
        # Handle the case when feast_type is an instance of ComplexFeastType
        pass

    raise ValueError(f"Could not convert Feast type {feast_type} to PyArrow type.")


def from_value_type(
    value_type: ValueType,
) -> FeastType:
    """
    Converts a ValueType enum to a Feast type.

    Args:
        value_type: The ValueType to be converted.

    Raises:
        ValueError: The conversion could not be performed.
    """
    if value_type in VALUE_TYPES_TO_FEAST_TYPES:
        return VALUE_TYPES_TO_FEAST_TYPES[value_type]

    raise ValueError(f"Could not convert value type {value_type} to FeastType.")


def from_feast_type(
    feast_type: FeastType,
) -> ValueType:
    """
    Converts a Feast type to a ValueType enum.

    Args:
        feast_type: The Feast type to be converted.

    Returns:
        The corresponding ValueType enum.

    Raises:
        ValueError: The conversion could not be performed.
    """
    if feast_type in VALUE_TYPES_TO_FEAST_TYPES.values():
        return list(VALUE_TYPES_TO_FEAST_TYPES.keys())[
            list(VALUE_TYPES_TO_FEAST_TYPES.values()).index(feast_type)
        ]

    raise ValueError(f"Could not convert feast type {feast_type} to ValueType.")