# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.

"""
Core type definitions and conversions.

This module contains type-related utilities including VectorType and
pandas-to-featureform datatype mappings.
"""

from typing import Union

import numpy

from ..enums import ScalarType
from ..proto import metadata_pb2 as pb

# Mapping from pandas/numpy dtypes to Featureform scalar types
pd_to_ff_datatype = {
    numpy.dtype("float64"): ScalarType.FLOAT64,
    numpy.dtype("float32"): ScalarType.FLOAT32,
    numpy.dtype("int64"): ScalarType.INT64,
    numpy.dtype("int"): ScalarType.INT,
    numpy.dtype("int32"): ScalarType.INT32,
    numpy.dtype("O"): ScalarType.STRING,
    numpy.dtype("str"): ScalarType.STRING,
    numpy.dtype("bool"): ScalarType.BOOL,
}


class VectorType:
    """Represents a vector type with scalar element type and dimensions."""

    def __init__(
        self, scalarType: Union[ScalarType, str], dims: int, is_embedding: bool = False
    ):
        self.scalarType = (
            scalarType if isinstance(scalarType, ScalarType) else ScalarType(scalarType)
        )
        self.dims = dims
        self.is_embedding = is_embedding

    def to_proto(self):
        proto_vec = pb.VectorType(
            scalar=self.scalarType.to_proto_enum(),
            dimension=self.dims,
            is_embedding=self.is_embedding,
        )
        return pb.ValueType(vector=proto_vec)

    @classmethod
    def from_proto(cls, proto_val):
        proto_vec = proto_val.vector
        scalar = ScalarType.from_proto(proto_vec.scalar)
        dims = proto_vec.dimension
        is_embedding = proto_vec.is_embedding
        return VectorType(scalar, dims, is_embedding)

    def __eq__(self, other):
        if not isinstance(other, VectorType):
            return False
        return (
            self.scalarType == other.scalarType
            and self.dims == other.dims
            and self.is_embedding == other.is_embedding
        )

    def __repr__(self):
        return f"VectorType({self.scalarType!r}, {self.dims}, is_embedding={self.is_embedding})"


class ListType:
    """Represents a variable-length list type with a scalar element type.

    Unlike VectorType (which has a fixed dimension for embeddings), ListType
    supports variable-length collections suitable for features like tags,
    categories, or historical values.

    Example:
        >>> from featureform import List, String, Int64
        >>> tags_type = List(String)
        >>> scores_type = List(Int64)
    """

    def __init__(self, element_type: Union[ScalarType, str]):
        """Initialize a ListType with the given element type.

        Args:
            element_type: The scalar type of elements in the list.
                Can be a ScalarType enum or a string like "string", "int64", etc.
        """
        self.element_type = (
            element_type
            if isinstance(element_type, ScalarType)
            else ScalarType(element_type)
        )

    def to_proto(self):
        """Convert this ListType to a protobuf ValueType."""
        proto_list = pb.ListType(element_type=self.element_type.to_proto_enum())
        return pb.ValueType(list=proto_list)

    @classmethod
    def from_proto(cls, proto_val):
        """Create a ListType from a protobuf ValueType."""
        proto_list = proto_val.list
        element_type = ScalarType.from_proto(proto_list.element_type)
        return ListType(element_type)

    def __eq__(self, other):
        if not isinstance(other, ListType):
            return False
        return self.element_type == other.element_type

    def __repr__(self):
        return f"List({self.element_type!r})"

    def __str__(self):
        return f"List<{self.element_type.value}>"

    @property
    def name(self):
        """Return the type name for CLI display compatibility.

        This property enables ListType to be used interchangeably with
        ScalarType (enum) in contexts that expect a .name attribute,
        such as the CLI type display.
        """
        return str(self)


def List(element_type: Union[ScalarType, str]) -> ListType:
    """Factory function to create a ListType.

    This provides a convenient API for creating list types:
        ff.List(ff.String)
        ff.List(ff.Int64)
        ff.List("float32")

    Args:
        element_type: The scalar type of elements in the list.

    Returns:
        A ListType instance.
    """
    return ListType(element_type)


def type_from_proto(proto_val):
    """Convert a protobuf ValueType to a Python type."""
    value_type = proto_val.WhichOneof("Type")
    if value_type == "scalar":
        return ScalarType.from_proto(proto_val.scalar)
    elif value_type == "vector":
        return VectorType.from_proto(proto_val)
    elif value_type == "list":
        return ListType.from_proto(proto_val)
    else:
        return ScalarType.NIL


# Union type for all supported value types in column schemas.
# This centralizes the type alias so it can be used consistently
# throughout the codebase.
ValueType = Union[ScalarType, ListType, VectorType]
