import builtins
from copy import deepcopy
from datetime import datetime
import importlib
import types
from types import GenericAlias
from typing import (
    Any,
    ClassVar,
    Dict,
    Generic,
    List,
    Optional,
    Set,
    Tuple,
    Type,
    TypeVar,
    Union,
    cast,
    get_args,
    get_origin,
    get_type_hints,
)

from vellum import ArrayVellumValue, ArrayVellumValueRequest, ChatMessagePromptBlock
from vellum.workflows.constants import undefined
from vellum.workflows.descriptors.base import BaseDescriptor
from vellum.workflows.types.core import Json, SpecialGenericAlias, UnderGenericAlias, UnionGenericAlias

_T = TypeVar("_T")
_LHS = TypeVar("_LHS")
_RHS = TypeVar("_RHS")

LOCAL_NS = {
    "Json": Json,
    "ArrayVellumValueRequest": ArrayVellumValueRequest,
    "ArrayVellumValue": ArrayVellumValue,
    "ChatMessagePromptBlock": ChatMessagePromptBlock,
}


def resolve_types(value: Union[BaseDescriptor[_T], _T]) -> Tuple[Type[_T], ...]:
    if isinstance(value, BaseDescriptor):
        return value.types

    return (value.__class__,)


def resolve_combined_types(
    lhs: Union[BaseDescriptor[_LHS], _LHS], rhs: Union[BaseDescriptor[_RHS], _RHS]
) -> Tuple[Union[Type[_LHS], Type[_RHS]], ...]:
    lhs_types = resolve_types(lhs)
    rhs_types = resolve_types(rhs)

    result_types: List[Union[Type[_LHS], Type[_RHS]]] = list(lhs_types)
    for rhs_type in rhs_types:
        if rhs_type not in result_types:
            result_types.append(rhs_type)

    return tuple(result_types)


def infer_types(object_: Type, attr_name: str, localns: Optional[Dict[str, Any]] = None) -> Tuple[Type, ...]:
    try:
        class_ = object_
        type_var_mapping = {}
        if isinstance(object_, UnderGenericAlias):
            origin = get_origin(object_)
            if origin and Generic in origin.__bases__:
                class_ = origin
                args = get_args(object_)
                type_var_mapping = {t: a for t, a in zip(origin.__parameters__, args)}

        if hasattr(object_, "__annotations__") and attr_name in object_.__annotations__:
            annotation_str = object_.__annotations__[attr_name]
            if isinstance(annotation_str, str) and "|" in annotation_str:
                parts = [part.strip() for part in annotation_str.split("|")]
                types_list: List[Type] = []
                for part in parts:
                    if part == "None":
                        types_list.append(type(None))
                    else:
                        try:
                            module = importlib.import_module(object_.__module__)
                            resolved_type = getattr(module, part, None)
                            if resolved_type is None:
                                resolved_type = getattr(builtins, part, None)
                            if resolved_type is not None and isinstance(resolved_type, type):
                                types_list.append(resolved_type)
                        except (ImportError, AttributeError):
                            pass
                if len(types_list) == len(parts):
                    return tuple(types_list)

        try:
            type_hints = get_type_hints(class_, localns=LOCAL_NS if localns is None else {**LOCAL_NS, **localns})
        except AttributeError:
            type_hints = {}
        if attr_name in type_hints:
            type_hint = type_hints[attr_name]
            if get_origin(type_hint) is ClassVar:
                return get_args(type_hint)
            if isinstance(type_hint, type):
                return (type_hint,)
            if isinstance(type_hint, UnionGenericAlias):
                if get_origin(type_hint) is Union:
                    return get_args(type_hint)
            if isinstance(type_hint, UnderGenericAlias):
                return (type_hint,)
            if isinstance(type_hint, SpecialGenericAlias):
                return (type_hint,)
            if isinstance(type_hint, GenericAlias):
                # In future versions of python, list[str] will be a `GenericAlias`
                return (cast(Type, type_hint),)
            if isinstance(type_hint, TypeVar):
                if type_hint in type_var_mapping:
                    return (type_var_mapping[type_hint],)
                return type_hint.__constraints__
            if type_hint is Any:
                return cast(Tuple[Type[Any], ...], (Any,))

        for base in reversed(class_.__mro__):
            class_attributes = vars(base)
            if attr_name in class_attributes:
                class_attribute = class_attributes[attr_name]
                return resolve_types(class_attribute)

        raise AttributeError(f"'{object_.__name__}' has no attribute '{attr_name}'")
    except TypeError:
        # Python 3.13+: object class doesn't have __annotations__ by default
        # Use getattr with default to safely access annotations
        annotations = getattr(object_, "__annotations__", {})
        annotation_value = annotations.get(attr_name, undefined)
        if annotation_value is undefined:
            raise AttributeError(f"'{object_.__name__}' has no attribute '{attr_name}'")
        raise AttributeError(
            f"Found 3.9+ typing syntax for field '{attr_name}' on class '{object_.__name__}' – {annotation_value}. Type annotations must be compatible with python version 3.8. "  # noqa: E501
        )


def get_class_attr_names(cls: Type) -> List[str]:
    # make sure we don't duplicate attributes
    collected_attributes: Set[str] = set()

    # we want to preserve the order of attributes on each class
    ordered_attr_names: List[str] = []

    for base in cls.__mro__:
        # gets attributes declared `foo = 1`
        for class_attribute in vars(base).keys():
            if class_attribute in collected_attributes:
                continue

            if class_attribute.startswith("_"):
                continue

            collected_attributes.add(class_attribute)
            ordered_attr_names.append(class_attribute)

        # gets type-annotated attributes `foo: int`
        ann = base.__dict__.get("__annotations__", {})
        for attr_name in ann.keys():
            if not isinstance(attr_name, str):
                continue

            if attr_name in collected_attributes:
                continue

            if attr_name.startswith("_"):
                continue

            collected_attributes.add(attr_name)
            ordered_attr_names.append(attr_name)

    # combine and filter out private attributes
    return ordered_attr_names


def deepcopy_with_exclusions(
    obj: _T,
    memo: Any,
    exclusions: Optional[Dict[str, Any]] = None,
) -> _T:
    cls = obj.__class__
    new_instance = cls.__new__(cls)
    new_instance.__dict__.update(obj.__dict__)

    exclusions = exclusions or {}

    for key, value in obj.__dict__.items():
        if key in exclusions:
            continue
        if isinstance(value, types.GeneratorType):
            new_instance.__dict__[key] = value
        else:
            new_instance.__dict__[key] = deepcopy(value, memo)

    for key, value in exclusions.items():
        new_instance.__dict__[key] = value

    return new_instance


def get_class_by_qualname(qualname: str) -> Type:
    module_name, class_name = qualname.rsplit(".", 1)
    module = importlib.import_module(module_name)
    imported_class = getattr(module, class_name)
    if not isinstance(imported_class, type):
        raise ValueError(f"Class {qualname} is not a valid type")

    return imported_class


def datetime_now() -> datetime:
    """
    There's a race condition between freezegun and pydantic that causes `PydanticSchemaGenerationError`
    because freezegun monkey-patches `datetime.now` to return a frozen datetime. This helper provides
    an alternative way to facilitate testing that doesn't rely on freezegun.
    """

    return datetime.now()


def get_original_base(cls: Type) -> Type:
    if not hasattr(cls, "__orig_bases__"):
        return type(None)

    # in Python 3.12, there is `from types import get_original_bases`, making this future proof
    # https://docs.python.org/3/library/types.html#types.get_original_bases
    return cls.__orig_bases__[0]  # type: ignore[attr-defined]


def coerce_to_declared_type(value: Any, declared_type: Type, field_name: str) -> Any:
    """Coerce a value to the declared type if needed.

    This handles cases where the API returns a float for an int field
    (since NumberVellumValue.value is typed as float).

    Args:
        value: The value to coerce
        declared_type: The expected type from the field annotation
        field_name: The name of the field (for error messages)

    Returns:
        The coerced value

    Raises:
        ValueError: If the value cannot be safely coerced to the declared type
    """
    if declared_type is int and isinstance(value, float):
        if value.is_integer():
            return int(value)
        raise ValueError(f"Expected integer for input '{field_name}', but received non-integer float: {value}")

    return value
