import inspect
from datetime import date, datetime, time, timedelta
from decimal import Decimal
from enum import Enum
from types import UnionType
from typing import Any, Literal, Union, get_args, get_origin

from yera.agents.dataclasses import AgentMetadata
from yera.agents.typing.utils import get_type_name
from yera.dsl.struct import Struct


def is_enum_type(expected_type: Any) -> bool:
    """Check if type is an Enum subclass."""
    try:
        return inspect.isclass(expected_type) and issubclass(expected_type, Enum)
    except TypeError:
        return False


def is_struct_type(expected_type: Any) -> bool:
    """Check if type is a Struct subclass."""
    try:
        return inspect.isclass(expected_type) and issubclass(expected_type, Struct)
    except TypeError:
        return False


def is_dataframe_type(type_hint) -> bool:
    """Check if type is pandas DataFrame without importing pandas."""
    try:
        module = getattr(type_hint, "__module__", None)
        name = getattr(type_hint, "__name__", None)
        module_match = module == "pandas" or module == "pandas.core.frame"
        return module_match and name == "DataFrame"
    except (AttributeError, TypeError):
        return False


SIMPLE_ALLOWED_TYPES = {
    type(None),
    None,
    int,
    float,
    str,
    bool,
    datetime,
    date,
    time,
    timedelta,
    Decimal,
}


def _validate_union(args):
    """All union members must be allowed."""
    return all(_is_allowed_type(arg) for arg in args)


def _validate_literal(args):
    """Literal values must be primitives."""
    return all(isinstance(arg, int | float | str | bool | type(None)) for arg in args)


def _validate_list(args):
    """Must have element type."""
    return bool(args) and _is_allowed_type(args[0])


def _validate_dict(args):
    """Must have key and value types."""
    if not args or len(args) != 2:
        return False
    return _is_allowed_type(args[0]) and _is_allowed_type(args[1])


def _validate_set(args):
    """Must have element type."""
    return bool(args) and _is_allowed_type(args[0])


GENERIC_VALIDATORS = {
    Union: _validate_union,
    UnionType: _validate_union,
    Literal: _validate_literal,
    list: _validate_list,
    dict: _validate_dict,
    set: _validate_set,
}


def _is_allowed_type(type_hint) -> bool:
    """Check if a type hint is allowed in agent signatures."""
    # None
    if type_hint is None or type_hint is type(None):
        return True

    # Simple types
    if type_hint in SIMPLE_ALLOWED_TYPES:
        return True

    # Struct/Enum types
    if (
        is_struct_type(type_hint)
        or is_enum_type(type_hint)
        or is_dataframe_type(type_hint)
    ):
        return True

    # Generic types with args validation
    origin = get_origin(type_hint)
    validator = GENERIC_VALIDATORS.get(origin)
    if validator:
        return validator(get_args(type_hint))

    return False


ERROR_HINTS = {
    list: "collections must specify element type, e.g., list[int]",
    dict: "collections must specify key and value types, e.g., dict[str, int]",
    set: "collections must specify element type, e.g., set[int]",
}


def _get_type_error_detail(type_hint) -> str:
    """Get helpful error message for invalid type."""
    type_name = get_type_name(type_hint)
    origin = get_origin(type_hint)

    hint = ERROR_HINTS.get(origin)
    if hint:
        return f"{type_name} ({hint})"
    return type_name


def _validate_parameters(signature, type_hints):
    """Validate parameter type hints, return (missing, invalid)."""
    missing = []
    invalid = []

    for param_name in signature.parameters:
        if param_name not in type_hints:
            missing.append(param_name)
        elif not _is_allowed_type(type_hints[param_name]):
            error_detail = _get_type_error_detail(type_hints[param_name])
            invalid.append((param_name, error_detail))

    return missing, invalid


def _validate_return_type(signature):
    """Validate return type hint, return error tuple or None."""
    return_annotation = signature.return_annotation
    if return_annotation is inspect.Parameter.empty:
        return_annotation = None  # void is OK

    if not _is_allowed_type(return_annotation):
        error_detail = _get_type_error_detail(return_annotation)
        return "function return", error_detail
    return None


def validate_signature(signature, type_hints, identifier):
    """Validate that all type hints in signature are allowed."""
    # Validate parameters
    missing_types, invalid_types = _validate_parameters(signature, type_hints)

    # Validate return type
    return_error = _validate_return_type(signature)
    if return_error:
        invalid_types.append(return_error)

    # Report errors
    if missing_types:
        raise TypeError(
            f"Parameters {missing_types} for agent '{identifier}' "
            f"are missing type hints. All parameters must be type-hinted."
        )

    if invalid_types:
        invalid_str = ",\n  ".join([f"'{n}': {t}" for n, t in invalid_types])
        raise TypeError(
            f"Agent '{identifier}' has invalid type hints: \n  {invalid_str}. "
            f"\nOnly supported types are allowed (primitives, Struct, Enum, datetime "
            f"types, Decimal, DataFrame, and typed collections like list[T], dict[K,V]"
            f", set[T])."
        )


def assert_value_is_return_type(
    return_value: Any, return_cls: type, meta: AgentMetadata
):
    if return_cls is None or return_cls is type(None):
        if return_value is not None:
            raise RuntimeError(
                f"Agent '{meta.identifier}' is marked as returning None, "
                f"but returned: {return_value!r}"
            )
        return

    if not isinstance(return_value, return_cls):
        raise TypeError(
            f"Agent '{meta.identifier}' expected to return "
            f"{meta.return_type}, but returned "
            f"{type(return_value).__name__}."
        )
