# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import builtins
import importlib
import inspect
import sys
import typing
from threading import Lock
from types import ModuleType, NoneType, UnionType
from typing import Any, Union, get_args

from haystack.core.errors import DeserializationError

_import_lock = Lock()


def _is_union_type(target: Any) -> bool:
    """
    Check if target is a Union type.

    This handles both `typing.Union[X, Y]` and `X | Y` syntax from PEP 604,
    including parameterized types like `Optional[str]`.
    """
    if target is Union or target is UnionType:
        return True
    origin = typing.get_origin(target)
    return origin is Union or origin is UnionType


def _build_pep604_union_type(types: list) -> Any:
    """Build a union type from a list of types using PEP 604 syntax (X | Y)."""
    result = types[0]
    for t in types[1:]:
        result = result | t
    return result


def serialize_type(target: Any) -> str:
    """
    Serializes a type or an instance to its string representation, including the module name.

    This function handles types, instances of types, and special typing objects.
    It assumes that non-typing objects will have a '__name__' attribute.

    :param target:
        The object to serialize, can be an instance or a type.
    :return:
        The string representation of the type.
    """
    if target is NoneType:
        return "None"

    args = get_args(target)

    if isinstance(target, UnionType):
        return " | ".join([serialize_type(a) for a in args])

    name = getattr(target, "__name__", str(target))
    if name.startswith("typing."):
        name = name[7:]
    if "[" in name:
        name = name.split("[")[0]

    # Get module name
    module = inspect.getmodule(target)
    module_name = ""
    # We omit the module name for builtins to not clutter the output
    if module and hasattr(module, "__name__") and module.__name__ != "builtins":
        module_name = f"{module.__name__}"

    if args:
        args_str = ", ".join([serialize_type(a) for a in args if a is not NoneType])
        return f"{module_name}.{name}[{args_str}]" if module_name else f"{name}[{args_str}]"

    return f"{module_name}.{name}" if module_name else f"{name}"


def _parse_generic_args(args_str):
    args = []
    bracket_count = 0
    current_arg = ""

    for char in args_str:
        if char == "[":
            bracket_count += 1
        elif char == "]":
            bracket_count -= 1

        if char == "," and bracket_count == 0:
            args.append(current_arg.strip())
            current_arg = ""
        else:
            current_arg += char

    if current_arg:
        args.append(current_arg.strip())

    return args


def _parse_pep604_union_args(union_str: str) -> list[str]:
    """
    Parse a PEP 604 union string (e.g., "str | int | None") into individual type strings.

    Handles nested generics properly, e.g., "list[str] | dict[str, int] | None".

    :param union_str: The union string to parse
    :returns: A list of individual type strings
    """
    args = []
    bracket_count = 0
    current_arg = ""

    for char in union_str:
        if char == "[":
            bracket_count += 1
        elif char == "]":
            bracket_count -= 1

        if char == "|" and bracket_count == 0:
            args.append(current_arg.strip())
            current_arg = ""
        else:
            current_arg += char

    if current_arg.strip():
        args.append(current_arg.strip())

    return args


def deserialize_type(type_str: str) -> Any:  # pylint: disable=too-many-return-statements
    """
    Deserializes a type given its full import path as a string, including nested generic types.

    This function will dynamically import the module if it's not already imported
    and then retrieve the type object from it. It also handles nested generic types like
    `list[dict[int, str]]`.

    :param type_str:
        The string representation of the type's full import path.
    :returns:
        The deserialized type object.
    :raises DeserializationError:
        If the type cannot be deserialized due to missing module or type.
    """
    # Handle PEP 604 union syntax at the top level (e.g., "str | int", "str | None")
    pep604_union_args = _parse_pep604_union_args(type_str)
    if len(pep604_union_args) > 1:
        deserialized_args = [deserialize_type(arg) for arg in pep604_union_args]
        return _build_pep604_union_type(deserialized_args)

    # Handle generics (including Union[X, Y])
    if "[" in type_str and type_str.endswith("]"):
        main_type_str, generics_str = type_str.split("[", 1)
        generics_str = generics_str[:-1]

        main_type = deserialize_type(main_type_str)
        generic_args = [deserialize_type(arg) for arg in _parse_generic_args(generics_str)]

        # Reconstruct
        try:
            return main_type[tuple(generic_args) if len(generic_args) > 1 else generic_args[0]]
        except (TypeError, AttributeError) as e:
            raise DeserializationError(f"Could not apply arguments {generic_args} to type {main_type}") from e

    # Handle non-generic types
    # First, check if there's a module prefix
    if "." in type_str:
        parts = type_str.split(".")
        module_name = ".".join(parts[:-1])
        type_name = parts[-1]

        module = sys.modules.get(module_name)
        if module is None:
            try:
                module = thread_safe_import(module_name)
            except ImportError as e:
                raise DeserializationError(f"Could not import the module: {module_name}") from e

        # Get the class from the module
        if hasattr(module, type_name):
            return getattr(module, type_name)

        raise DeserializationError(f"Could not locate the type: {type_name} in the module: {module_name}")

    # No module prefix, check builtins and typing
    # First check builtins
    if hasattr(builtins, type_str):
        return getattr(builtins, type_str)

    # Then check typing
    if hasattr(typing, type_str):
        return getattr(typing, type_str)

    # Special case for NoneType
    if type_str == "NoneType":
        return NoneType

    # Special case for None
    if type_str == "None":
        return None

    raise DeserializationError(f"Could not deserialize type: {type_str}")


def thread_safe_import(module_name: str) -> ModuleType:
    """
    Import a module in a thread-safe manner.

    Importing modules in a multi-threaded environment can lead to race conditions.
    This function ensures that the module is imported in a thread-safe manner without having impact
    on the performance of the import for single-threaded environments.

    :param module_name: the module to import
    """
    with _import_lock:
        return importlib.import_module(module_name)
