import asyncio
import inspect
import logging
import os
import re
import typing
from collections.abc import Callable, Iterator
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from importlib import import_module
from types import ModuleType
from typing import (
    Annotated,
    Any,
    Literal,
    Optional,
    Union,
    cast,
    get_args,
    get_origin,
    get_type_hints,
)

from pydantic import BaseModel, ConfigDict, Field, create_model, model_serializer
from pydantic.fields import FieldInfo
from pydantic_core import PydanticUndefined

from arcade_core.annotations import Inferrable
from arcade_core.auth import OAuth2, ToolAuthorization
from arcade_core.errors import (
    ToolDefinitionError,
    ToolInputSchemaError,
    ToolkitLoadError,
    ToolOutputSchemaError,
)
from arcade_core.metadata import ToolMetadata
from arcade_core.schema import (
    TOOL_NAME_SEPARATOR,
    FullyQualifiedName,
    InputParameter,
    OAuth2Requirement,
    ToolAuthRequirement,
    ToolContext,
    ToolDefinition,
    ToolInput,
    ToolkitDefinition,
    ToolMetadataKey,
    ToolMetadataRequirement,
    ToolOutput,
    ToolRequirements,
    ToolSecretRequirement,
    ValueSchema,
)
from arcade_core.toolkit import Toolkit
from arcade_core.utils import (
    does_function_return_value,
    first_or_none,
    is_strict_optional,
    is_string_literal,
    is_union,
    snake_to_pascal_case,
    space_to_snake_case,
)

logger = logging.getLogger(__name__)

InnerWireType = Literal["string", "integer", "number", "boolean", "json"]
WireType = Union[InnerWireType, Literal["array"]]


def is_typeddict(tp: type) -> bool:
    """
    Check if a type is a TypedDict.
    Works with both typing.TypedDict and typing_extensions.TypedDict.
    """
    try:
        # TypedDict creates classes that inherit from dict
        if not isinstance(tp, type) or not issubclass(tp, dict):
            return False

        # Check for TypedDict-specific attributes
        return (
            hasattr(tp, "__annotations__")
            and hasattr(tp, "__total__")
            and hasattr(tp, "__required_keys__")
            and hasattr(tp, "__optional_keys__")
        )
    except TypeError:
        # Some special forms raise TypeError when checking issubclass
        return False


@dataclass
class WireTypeInfo:
    """
    Represents the wire type information for a value, including its inner type if it's a list.
    """

    wire_type: WireType
    inner_wire_type: InnerWireType | None = None
    enum_values: list[str] | None = None
    properties: dict[str, "WireTypeInfo"] | None = None
    inner_properties: dict[str, "WireTypeInfo"] | None = None
    description: str | None = None
    required_keys: list[str] | None = None
    inner_required_keys: list[str] | None = None
    nullable: bool | None = None


class ToolMeta(BaseModel):
    """
    Metadata for a tool once it's been materialized.
    """

    module: str
    toolkit: str | None = None
    package: str | None = None
    path: str | None = None
    date_added: datetime = Field(default_factory=datetime.now)
    date_updated: datetime = Field(default_factory=datetime.now)


class MaterializedTool(BaseModel):
    """
    Data structure that holds tool information while stored in the Catalog
    """

    tool: Callable
    definition: ToolDefinition
    meta: ToolMeta

    # Thought (Sam): Should generate create these from ToolDefinition?
    input_model: type[BaseModel]
    output_model: type[BaseModel]

    @property
    def name(self) -> str:
        return self.definition.name

    @property
    def version(self) -> str | None:
        return self.definition.toolkit.version

    @property
    def description(self) -> str:
        return self.definition.description

    @property
    def requires_auth(self) -> bool:
        return self.definition.requirements.authorization is not None


class ToolCatalog(BaseModel):
    """Singleton class that holds all tools for a given worker"""

    _tools: dict[FullyQualifiedName, MaterializedTool] = {}

    _disabled_tools: set[str] = set()
    _disabled_toolkits: set[str] = set()

    def __init__(self, **data) -> None:  # type: ignore[no-untyped-def]
        super().__init__(**data)
        self._load_disabled_tools()
        self._load_disabled_toolkits()

    def _load_disabled_tools(self) -> None:
        """Load disabled tools from the environment variable.

        The ARCADE_DISABLED_TOOLS environment variable should contain a
        comma-separated list of tools that are to be excluded from the
        catalog.

        The expected format for each disabled tool is:
        - [CamelCaseToolkitName][TOOL_NAME_SEPARATOR][CamelCaseToolName]
        """
        disabled_tools = os.getenv("ARCADE_DISABLED_TOOLS", "").strip().split(",")
        if not disabled_tools:
            return

        pattern = re.compile(rf"^[a-zA-Z]+{re.escape(TOOL_NAME_SEPARATOR)}[a-zA-Z]+$")

        for tool in disabled_tools:
            if not pattern.match(tool):
                continue

            self._disabled_tools.add(tool.lower())

    def _load_disabled_toolkits(self) -> None:
        """Load disabled toolkits from the environment variable.

        The ARCADE_DISABLED_TOOLKITS environment variable should contain a
        comma-separated list of toolkits that are to be excluded from the
        catalog.

        The expected format for each disabled toolkit is:
        - [CamelCaseToolkitName]
        """
        disabled_toolkits = os.getenv("ARCADE_DISABLED_TOOLKITS", "").strip().split(",")
        if not disabled_toolkits:
            return

        for toolkit in disabled_toolkits:
            self._disabled_toolkits.add(toolkit.lower())

    def add_tool(
        self,
        tool_func: Callable,
        toolkit_or_name: str | Toolkit,
        module: ModuleType | None = None,
        toolkit_version: str | None = None,
        toolkit_description: str | None = None,
    ) -> None:
        """
        Add a function to the catalog as a tool.

        Args:
            tool_func: The function to add to the catalog.
            toolkit_or_name: The toolkit or name of the toolkit.
            module: The module to add the tool from.
            toolkit_version: The version of the toolkit. Overrides the version of the toolkit if provided.
            toolkit_description: The description of the toolkit. Overrides the description of the toolkit if provided.
        """

        input_model, output_model = create_func_models(tool_func)

        if isinstance(toolkit_or_name, Toolkit):
            toolkit = toolkit_or_name
            toolkit_name = toolkit.name
            toolkit_version = toolkit_version or toolkit.version
            toolkit_description = toolkit_description or toolkit.description
        elif isinstance(toolkit_or_name, str):
            toolkit = None
            toolkit_name = toolkit_or_name

        if not toolkit_name:
            raise ValueError("A server name or server must be provided.")

        definition = ToolCatalog.create_tool_definition(
            tool_func,
            toolkit_name,
            toolkit_version,
            toolkit_description,
        )

        fully_qualified_name = definition.get_fully_qualified_name()

        if fully_qualified_name in self._tools:
            raise ToolkitLoadError(
                f"Tool '{definition.name}' in server '{toolkit_name}' already exists in the catalog."
            )

        if str(fully_qualified_name).lower() in self._disabled_tools:
            logger.info(f"Tool '{fully_qualified_name!s}' is disabled and will not be cataloged.")
            return

        if str(toolkit_name).lower() in self._disabled_toolkits:
            logger.info(f"Server '{toolkit_name!s}' is disabled and will not be cataloged.")
            return

        self._tools[fully_qualified_name] = MaterializedTool(
            definition=definition,
            tool=tool_func,
            meta=ToolMeta(
                module=module.__name__ if module else tool_func.__module__,
                toolkit=toolkit_name,
                package=toolkit.package_name if toolkit else None,
                path=module.__file__ if module else None,
            ),
            input_model=input_model,
            output_model=output_model,
        )

    def add_module(
        self,
        module: ModuleType,
        name: str | None = None,
        version: str | None = None,
        description: str | None = None,
    ) -> None:
        """
        Add all the tools in a module to the catalog.

        Args:
            module: The module to add.
            name: Optionally override the name of the toolkit with this parameter
            version: Optionally override the version of the toolkit with this parameter
            description: Optionally override the description of the toolkit with this parameter
        """
        toolkit = Toolkit.from_module(module)
        if name:
            toolkit.name = name
        self.add_toolkit(toolkit, version=version, description=description)

    def add_toolkit(
        self, toolkit: Toolkit, version: str | None = None, description: str | None = None
    ) -> None:
        """
        Add the tools from a loaded toolkit to the catalog.

        Args:
            toolkit: The toolkit to add.
            version: Optionally override the version of the toolkit with this parameter
            description: Optionally override the description of the toolkit with this parameter
        """

        if str(toolkit).lower() in self._disabled_toolkits:
            logger.info(f"Server '{toolkit.name!s}' is disabled and will not be cataloged.")
            return

        for module_name, tool_names in toolkit.tools.items():
            for tool_name in tool_names:
                try:
                    module = import_module(module_name)
                    tool_func = getattr(module, tool_name)
                    self.add_tool(
                        tool_func,
                        toolkit,
                        module,
                        toolkit_version=version,
                        toolkit_description=description,
                    )
                except ToolDefinitionError as e:
                    raise e.with_context(tool_name) from e
                except ToolkitLoadError as e:
                    raise e.with_context(toolkit.name) from e
                except ImportError as e:
                    raise ToolkitLoadError(
                        f"Could not import module {module_name}. Reason: {e}"
                    ).with_context(tool_name)
                except AttributeError as e:
                    raise ToolDefinitionError(
                        f"Could not import tool {tool_name} in module {module_name}. Reason: {e}"
                    ).with_context(tool_name)
                except TypeError as e:
                    raise ToolDefinitionError(
                        f"Type error encountered while adding tool {tool_name} from {module_name}. Reason: {e}"
                    ).with_context(tool_name)
                except Exception as e:
                    raise ToolDefinitionError(
                        f"Error encountered while adding tool {tool_name} from {module_name}. Reason: {e}"
                    ).with_context(tool_name)

    def __getitem__(self, name: FullyQualifiedName) -> MaterializedTool:
        return self.get_tool(name)

    def __contains__(self, name: FullyQualifiedName) -> bool:
        return name in self._tools

    def __iter__(self) -> Iterator[MaterializedTool]:  # type: ignore[override]
        yield from self._tools.values()

    def __len__(self) -> int:
        return len(self._tools)

    def is_empty(self) -> bool:
        return len(self._tools) == 0

    def get_tool_names(self) -> list[FullyQualifiedName]:
        return [tool.definition.get_fully_qualified_name() for tool in self._tools.values()]

    def find_tool_by_func(self, func: Callable) -> ToolDefinition:
        """
        Find a tool by its function.
        """
        for _, tool in self._tools.items():
            if tool.tool == func:
                return tool.definition
        raise ValueError(f"Tool {func} not found in the catalog.")

    def get_tool_by_name(
        self,
        name: str,
        version: str | None = None,
        separator: str = TOOL_NAME_SEPARATOR,
    ) -> MaterializedTool:
        """Get a tool from the catalog by name.

        Args:
            name: The name of the tool, potentially including the toolkit name separated by the `separator`.
            version: The version of the toolkit. Defaults to None.
            separator: The separator between toolkit and tool names. Defaults to `TOOL_NAME_SEPARATOR`.

        Returns:
            MaterializedTool: The matching tool from the catalog.

        Raises:
            ValueError: If the tool is not found in the catalog.
        """
        if separator in name:
            toolkit_name, tool_name = name.split(separator, 1)
            fq_name = FullyQualifiedName(
                name=tool_name, toolkit_name=toolkit_name, toolkit_version=version
            )
            return self.get_tool(fq_name)
        else:
            # No toolkit name provided, search tools with matching tool name
            matching_tools = [
                tool
                for fq_name, tool in self._tools.items()
                if fq_name.name.lower() == name.lower()
                and (
                    version is None
                    or (fq_name.toolkit_version or "").lower() == (version or "").lower()
                )
            ]
            if matching_tools:
                return matching_tools[0]

        raise ValueError(f"Tool {name} not found in the catalog.")

    def get_tool(self, name: FullyQualifiedName) -> MaterializedTool:
        """
        Get a tool from the catalog by fully-qualified name and version.
        If the version is not specified, the any version is returned.
        """
        if name.toolkit_version:
            try:
                return self._tools[name]
            except KeyError:
                raise ValueError(f"Tool {name}@{name.toolkit_version} not found in the catalog.")

        for key, tool in self._tools.items():
            if key.equals_ignoring_version(name):
                return tool

        raise ValueError(f"Tool {name} not found.")

    def get_tool_count(self) -> int:
        """
        Get the number of tools in the catalog.
        """
        return len(self._tools)

    @staticmethod
    def create_tool_definition(
        tool: Callable,
        toolkit_name: str,
        toolkit_version: str | None = None,
        toolkit_desc: str | None = None,
    ) -> ToolDefinition:
        """
        Given a tool function, create a ToolDefinition
        """

        raw_tool_name = getattr(tool, "__tool_name__", tool.__name__)

        # Hard requirement: tools must have descriptions
        tool_description = getattr(tool, "__tool_description__", None)
        if not tool_description:
            raise ToolDefinitionError(
                f"Tool '{raw_tool_name}' is missing a description. Tool descriptions are specified as docstrings for the tool function."
            )

        # If the function returns a value, it must have a type annotation
        if does_function_return_value(tool) and tool.__annotations__.get("return") is None:
            raise ToolOutputSchemaError(f"Tool '{raw_tool_name}' must have a return type")

        auth_requirement = create_auth_requirement(tool)
        secrets_requirement = create_secrets_requirement(tool)
        metadata_requirement = create_metadata_requirement(tool, auth_requirement)

        toolkit_definition = ToolkitDefinition(
            name=snake_to_pascal_case(space_to_snake_case(toolkit_name)),
            description=toolkit_desc,
            version=toolkit_version,
        )

        tool_name = snake_to_pascal_case(raw_tool_name)
        fully_qualified_name = FullyQualifiedName.from_toolkit(tool_name, toolkit_definition)
        deprecation_message = getattr(tool, "__tool_deprecation_message__", None)
        tool_metadata = getattr(tool, "__tool_metadata__", None)

        if tool_metadata is not None:
            if not isinstance(tool_metadata, ToolMetadata):
                raise ToolDefinitionError(
                    f"Expected a ToolMetadata instance for 'metadata', "
                    f"but got {type(tool_metadata).__name__}. "
                )
            tool_metadata.validate_for_tool()

        return ToolDefinition(
            name=tool_name,
            fully_qualified_name=str(fully_qualified_name),
            description=tool_description,
            toolkit=toolkit_definition,
            input=create_input_definition(tool),
            output=create_output_definition(tool),
            requirements=ToolRequirements(
                authorization=auth_requirement,
                secrets=secrets_requirement,
                metadata=metadata_requirement,
            ),
            deprecation_message=deprecation_message,
            metadata=tool_metadata,
        )


def create_input_definition(func: Callable) -> ToolInput:
    """
    Create an input model for a function based on its parameters.
    """
    input_parameters = []
    tool_context_param_name: str | None = None

    for _, param in inspect.signature(func, follow_wrapped=True).parameters.items():
        ann = param.annotation
        if isinstance(ann, type) and issubclass(ann, ToolContext):
            # Soft guidance for developers using legacy ToolContext
            if tool_context_param_name is not None:
                raise ToolInputSchemaError(
                    f"Only one ToolContext parameter is supported, but tool {func.__name__} has multiple."
                )

            tool_context_param_name = param.name
            continue  # No further processing of this param (don't add it to the list of inputs)

        tool_field_info = extract_field_info(param)

        # If the field has an explicit default in the signature, it is not required.
        # If the annotation is Optional[...], it is not required.
        # (A bare `param: str = None` has has_explicit_default=True but
        # is_optional=False, and must still be treated as not-required.)
        is_required = not tool_field_info.is_optional and not tool_field_info.has_explicit_default

        input_parameters.append(
            InputParameter(
                name=tool_field_info.name,
                description=tool_field_info.description,
                required=is_required,
                inferrable=tool_field_info.is_inferrable,
                value_schema=wire_type_info_to_value_schema(tool_field_info.wire_type_info),
            )
        )

    return ToolInput(
        parameters=input_parameters,
        tool_context_parameter_name=tool_context_param_name,
    )


def create_output_definition(func: Callable) -> ToolOutput:
    """
    Create an output model for a function based on its return annotation.
    """
    return_type = inspect.signature(func, follow_wrapped=True).return_annotation
    description = "No description provided."

    if return_type is inspect.Signature.empty:
        return ToolOutput(
            value_schema=None,
            description="No description provided.",
            available_modes=["null"],
        )

    if hasattr(return_type, "__metadata__"):
        description = (
            return_type.__metadata__[0]
            if return_type.__metadata__
            else "No description provided for return type."
        )
        return_type = return_type.__origin__

    # Unwrap Optional types
    # Both Optional[T] and T | None are supported
    is_optional = is_strict_optional(return_type)
    if is_optional:
        return_type = next(arg for arg in get_args(return_type) if arg is not type(None))

    wire_type_info = get_wire_type_info(return_type)

    available_modes = ["value", "error"]

    if is_optional:
        available_modes.append("null")

    return ToolOutput(
        description=description,
        available_modes=available_modes,
        value_schema=wire_type_info_to_value_schema(wire_type_info),
    )


def create_auth_requirement(tool: Callable) -> ToolAuthRequirement | None:
    """
    Create an auth requirement for a tool.
    """
    auth_requirement = getattr(tool, "__tool_requires_auth__", None)
    if isinstance(auth_requirement, ToolAuthorization):
        new_auth_requirement = ToolAuthRequirement(
            provider_id=auth_requirement.provider_id,
            provider_type=auth_requirement.provider_type,
            id=auth_requirement.id,
        )
        if isinstance(auth_requirement, OAuth2):
            new_auth_requirement.oauth2 = OAuth2Requirement(**auth_requirement.model_dump())
        auth_requirement = new_auth_requirement

    return auth_requirement


def create_secrets_requirement(tool: Callable) -> list[ToolSecretRequirement] | None:
    """
    Create a secrets requirement for a tool.
    """
    raw_tool_name = getattr(tool, "__tool_name__", tool.__name__)
    secrets_requirement = getattr(tool, "__tool_requires_secrets__", None)
    if isinstance(secrets_requirement, list):
        if any(not isinstance(secret, str) for secret in secrets_requirement):
            raise ToolDefinitionError(
                f"Secret keys must be strings (error in tool {raw_tool_name})."
            )

        secrets_requirement = to_tool_secret_requirements(secrets_requirement)
        if any(secret.key is None or secret.key.strip() == "" for secret in secrets_requirement):
            raise ToolDefinitionError(
                f"Secrets must have a non-empty key (error in tool {raw_tool_name})."
            )

    return secrets_requirement


def create_metadata_requirement(
    tool: Callable, auth_requirement: ToolAuthRequirement | None
) -> list[ToolMetadataRequirement] | None:
    """
    Create a metadata requirement for a tool.
    """
    raw_tool_name = getattr(tool, "__tool_name__", tool.__name__)
    metadata_requirement = getattr(tool, "__tool_requires_metadata__", None)
    if isinstance(metadata_requirement, list):
        for metadata in metadata_requirement:
            if not isinstance(metadata, str):
                raise ToolDefinitionError(
                    f"Metadata must be strings (error in tool {raw_tool_name})."
                )
            if ToolMetadataKey.requires_auth(metadata) and auth_requirement is None:
                raise ToolDefinitionError(
                    f"Tool {raw_tool_name} declares metadata key '{metadata}', "
                    "which requires that the tool has an auth requirement, "
                    "but no auth requirement was provided. Please specify an auth requirement."
                )

        metadata_requirement = to_tool_metadata_requirements(metadata_requirement)
        if any(
            metadata.key is None or metadata.key.strip() == "" for metadata in metadata_requirement
        ):
            raise ToolDefinitionError(
                f"Metadata must have a non-empty key (error in tool {raw_tool_name})."
            )

    return metadata_requirement


@dataclass
class ParamInfo:
    """
    Information about a function parameter found through inspection.
    """

    name: str
    default: Any
    original_type: type
    field_type: type
    description: str | None = None
    is_optional: bool = True
    # True iff the Python signature (or pydantic Field) supplies a default.
    # Distinct from `default is not None` — `def f(x: str = None)` has an
    # explicit default (None) and must not be treated as required.
    has_explicit_default: bool = False


@dataclass
class ToolParamInfo:
    """
    Information about a tool parameter, including computed values.
    """

    name: str
    default: Any
    original_type: type
    field_type: type
    wire_type_info: WireTypeInfo
    description: str | None = None
    is_optional: bool = True
    is_inferrable: bool = True
    has_explicit_default: bool = False

    @classmethod
    def from_param_info(
        cls,
        param_info: ParamInfo,
        wire_type_info: WireTypeInfo,
        is_inferrable: bool = True,
    ) -> "ToolParamInfo":
        return cls(
            name=param_info.name,
            default=param_info.default,
            original_type=param_info.original_type,
            field_type=param_info.field_type,
            description=param_info.description,
            is_optional=param_info.is_optional,
            wire_type_info=wire_type_info,
            is_inferrable=is_inferrable,
            has_explicit_default=param_info.has_explicit_default,
        )


def extract_field_info(param: inspect.Parameter) -> ToolParamInfo:
    """
    Extract type and field parameters from a function parameter.
    """
    annotation = param.annotation
    if annotation == inspect.Parameter.empty:
        raise ToolInputSchemaError(f"Parameter {param} has no type annotation.")

    # Get the majority of the param info from either the Pydantic Field() or regular inspection
    if isinstance(param.default, FieldInfo):
        param_info = extract_pydantic_param_info(param)
    else:
        param_info = extract_python_param_info(param)

    metadata = getattr(annotation, "__metadata__", [])
    str_annotations = [m for m in metadata if isinstance(m, str)]

    # Get the description from annotations, if present
    if len(str_annotations) == 0:
        pass
    elif len(str_annotations) == 1:
        param_info.description = str_annotations[0]
    elif len(str_annotations) == 2:
        new_name = str_annotations[0]
        if not new_name.isidentifier():
            raise ToolInputSchemaError(
                f"Invalid parameter name: '{new_name}' is not a valid identifier. "
                "Identifiers must start with a letter or underscore, "
                "and can only contain letters, digits, or underscores."
            )
        param_info.name = new_name
        param_info.description = str_annotations[1]
    else:
        raise ToolInputSchemaError(
            f"Parameter {param} has too many string annotations. Expected 0, 1, or 2, got {len(str_annotations)}."
        )

    # Get the Inferrable annotation, if it exists
    inferrable_annotation = first_or_none(Inferrable, get_args(annotation))

    # Params are inferrable by default
    is_inferrable = inferrable_annotation.value if inferrable_annotation else True

    # Get the wire (serialization) type information for the type
    wire_type_info = get_wire_type_info(param_info.field_type)

    # Final reality check
    if param_info.description is None:
        raise ToolInputSchemaError(
            f"Parameter '{param_info.name}' is missing a description. Parameter descriptions are specified as string annotations using the typing.Annotated class."
        )

    if wire_type_info.wire_type is None:
        raise ToolInputSchemaError(f"Unknown parameter type: {param_info.field_type}")

    return ToolParamInfo.from_param_info(param_info, wire_type_info, is_inferrable)


def get_wire_type_info(_type: type) -> WireTypeInfo:
    """
    Get the wire type information for a given type.
    """

    # Is this a list type?
    # If so, get the inner (enclosed) type
    is_list = get_origin(_type) is list
    inner_properties = None
    inner_required_keys = None

    if is_list:
        inner_type = get_args(_type)[0]

        # Recursively get wire type info for inner type
        inner_info = get_wire_type_info(inner_type)
        inner_wire_type = cast(InnerWireType, inner_info.wire_type)

        # If inner type has a known object shape (possibly empty), propagate it. A known-empty
        # shape ({}) is distinct from None (unknown shape) so converters can close it.
        if inner_info.properties is not None:
            inner_properties = inner_info.properties
            inner_required_keys = inner_info.required_keys
        # If inner type is array (nested arrays), propagate inner_properties
        elif inner_info.inner_properties is not None:
            inner_properties = inner_info.inner_properties
            inner_required_keys = inner_info.inner_required_keys
    else:
        inner_wire_type = None

    # Get the outer wire type
    wire_type = get_wire_type(str) if is_string_literal(_type) else get_wire_type(_type)

    # Handle enums (known/fixed lists of values)
    is_enum = False
    enum_values: list[str] = []

    type_to_check = inner_type if is_list else _type

    # Strip generic parameters if type_to_check is a parameterized generic
    actual_type = get_origin(type_to_check) or type_to_check

    # Special case: Literal["string1", "string2"] can be enumerated on the wire
    if is_string_literal(type_to_check):
        is_enum = True
        enum_values = [str(e) for e in get_args(type_to_check)]

    # Special case: Enum can be enumerated on the wire
    elif isinstance(actual_type, type) and issubclass(actual_type, Enum):
        is_enum = True
        enum_values = [e.value for e in actual_type]

    # Extract properties for complex types
    properties = None
    required_keys = None
    if wire_type == "json" and not is_list:
        properties, required_keys = extract_properties(type_to_check)

    return WireTypeInfo(
        wire_type,
        inner_wire_type,
        enum_values if is_enum else None,
        properties,
        inner_properties,
        required_keys=required_keys,
        inner_required_keys=inner_required_keys,
    )


def _extract_typeddict_field_descriptions(typeddict_class: type) -> dict[str, str]:
    """
    Extract field descriptions from TypedDict docstrings.

    TypedDict classes typically have field descriptions as docstrings after each field.
    This function attempts to parse the source code to extract these descriptions.
    """
    descriptions = {}

    try:
        source = inspect.getsource(typeddict_class)
        # Simple regex to match field: type pattern followed by a docstring
        # This is a simplified approach - a full AST parser would be more robust
        import re

        # Pattern to match field definition followed by docstring
        pattern = r'(\w+):\s*[^"\n]+\n\s*"""([^"]+)"""'
        matches = re.findall(pattern, source)

        for field_name, description in matches:
            descriptions[field_name] = description.strip()

    except (OSError, TypeError):
        # If we can't get the source, return empty descriptions
        pass

    return descriptions


def extract_properties(
    type_to_check: type,
) -> tuple[dict[str, WireTypeInfo] | None, list[str] | None]:
    """
    Extract properties from TypedDict, Pydantic models, or other structured types.

    Returns (properties, required_keys). For a known structured type (TypedDict or
    Pydantic model) properties is a dict of its fields, empty when the type declares no
    fields, and required_keys is a sorted list of required property names, empty when every
    field is optional. A known-empty shape (no fields) therefore yields ``({}, [])``, which
    converters can distinguish from a type with no known shape (e.g. a freeform ``dict``),
    which yields ``(None, None)`` to signal that both shape and requiredness are unknown.
    """
    properties = {}

    # Handle Pydantic BaseModel
    if isinstance(type_to_check, type) and issubclass(type_to_check, BaseModel):
        pydantic_required_keys: list[str] = []
        for field_name, field_info in type_to_check.model_fields.items():
            # Get the field type
            field_type = field_info.annotation
            if field_type is None:
                continue

            # Handle Optional types (Union[T, None])
            is_nullable = is_strict_optional(field_type)
            if is_nullable:
                # Extract the non-None type from Optional
                field_type = next(arg for arg in get_args(field_type) if arg is not type(None))

            # Get wire type info recursively
            # field_type cannot be None here due to the check above
            wire_info = get_wire_type_info(field_type)
            if is_nullable:
                wire_info.nullable = True
            properties[field_name] = wire_info

            # A Pydantic field is required when it has no default (default or factory).
            if field_info.is_required():
                pydantic_required_keys.append(field_name)

        return (properties, sorted(pydantic_required_keys))

    # Handle TypedDict
    elif is_typeddict(type_to_check):
        # Get type hints for the TypedDict
        type_hints = get_type_hints(type_to_check, include_extras=True)

        # Try to extract field descriptions from the class source
        field_descriptions = _extract_typeddict_field_descriptions(type_to_check)

        for field_name, field_type in type_hints.items():
            # Handle Optional types (Union[T, None])
            is_nullable = is_strict_optional(field_type)
            if is_nullable:
                # Extract the non-None type from Optional
                field_type = next(arg for arg in get_args(field_type) if arg is not type(None))
            wire_info = get_wire_type_info(field_type)
            if is_nullable:
                wire_info.nullable = True

            # Add description if available
            if field_name in field_descriptions:
                wire_info.description = field_descriptions[field_name]

            properties[field_name] = wire_info

        required_keys = sorted(getattr(type_to_check, "__required_keys__", frozenset()))
        return (properties, required_keys)

    # Handle regular dict with type annotations (e.g., dict[str, Any])
    elif get_origin(type_to_check) is dict:
        # For generic dicts, we can't extract specific properties
        return (None, None)

    return (properties or None, None)


def wire_type_info_to_value_schema(wire_info: WireTypeInfo) -> ValueSchema:
    """
    Convert WireTypeInfo to ValueSchema, including nested properties.
    """
    # Convert nested properties if they exist. A known-empty object shape ({}) is preserved
    # as {} rather than collapsed to None so converters can emit it as a closed object.
    properties = None
    if wire_info.properties is not None:
        properties = {
            name: wire_type_info_to_value_schema(nested_info)
            for name, nested_info in wire_info.properties.items()
        }

    # Convert inner properties for array items
    inner_properties = None
    if wire_info.inner_properties is not None:
        inner_properties = {
            name: wire_type_info_to_value_schema(nested_info)
            for name, nested_info in wire_info.inner_properties.items()
        }

    return ValueSchema(
        val_type=wire_info.wire_type,
        inner_val_type=wire_info.inner_wire_type,
        enum=wire_info.enum_values,
        properties=properties,
        inner_properties=inner_properties,
        description=wire_info.description,
        required_keys=wire_info.required_keys,
        inner_required_keys=wire_info.inner_required_keys,
        nullable=wire_info.nullable,
    )


def extract_python_param_info(param: inspect.Parameter) -> ParamInfo:
    # If the param is Annotated[], unwrap the annotation to get the "real" type
    # Otherwise, use the literal type
    annotation = param.annotation
    original_type = annotation.__args__[0] if get_origin(annotation) is Annotated else annotation
    field_type = original_type

    # Handle optional types
    # Both Optional[T] and T | None are supported
    is_optional = is_strict_optional(field_type)
    if is_optional:
        field_type = next(arg for arg in get_args(field_type) if arg is not type(None))

    # Union types are not currently supported
    # (other than optional, which is handled above)
    if is_union(field_type):
        raise ToolInputSchemaError(
            f"Parameter {param.name} is a union type. Only optional types are supported."
        )

    has_explicit_default = param.default is not inspect.Parameter.empty
    return ParamInfo(
        name=param.name,
        default=param.default if has_explicit_default else None,
        is_optional=is_optional,
        original_type=original_type,
        field_type=field_type,
        has_explicit_default=has_explicit_default,
    )


def extract_pydantic_param_info(param: inspect.Parameter) -> ParamInfo:
    has_explicit_default = (
        param.default.default is not PydanticUndefined or param.default.default_factory is not None
    )
    default_value = None if param.default.default is PydanticUndefined else param.default.default

    if param.default.default_factory is not None:
        if callable(param.default.default_factory):
            default_value = param.default.default_factory()
        else:
            raise ToolInputSchemaError(f"Default factory for parameter {param} is not callable.")

    # If the param is Annotated[], unwrap the annotation to get the "real" type
    # Otherwise, use the literal type
    original_type = (
        param.annotation.__args__[0]
        if get_origin(param.annotation) is Annotated
        else param.annotation
    )
    field_type = original_type

    # Unwrap Optional types
    # Both Optional[T] and T | None are supported
    is_optional = is_strict_optional(field_type)
    if is_optional:
        field_type = next(arg for arg in get_args(field_type) if arg is not type(None))

    return ParamInfo(
        name=param.name,
        description=param.default.description,
        default=default_value,
        is_optional=is_optional,
        original_type=original_type,
        field_type=field_type,
        has_explicit_default=has_explicit_default,
    )


def get_wire_type(
    _type: type,
) -> WireType:
    """
    Mapping between Python types and HTTP/JSON types
    """
    # TODO ensure Any is not allowed
    type_mapping: dict[type, WireType] = {
        str: "string",
        bool: "boolean",
        int: "integer",
        float: "number",
        dict: "json",
    }
    outer_type_mapping: dict[type, WireType] = {
        list: "array",
        dict: "json",
    }
    wire_type = type_mapping.get(_type)
    if wire_type:
        return wire_type

    if hasattr(_type, "__origin__"):
        wire_type = outer_type_mapping.get(cast(type, get_origin(_type)))
        if wire_type:
            return wire_type

    if isinstance(_type, type) and issubclass(_type, Enum):
        return "string"

    if isinstance(_type, type) and issubclass(_type, BaseModel):
        return "json"

    if is_typeddict(_type):
        return "json"

    raise ToolDefinitionError(f"Unsupported parameter type: {_type}")


def _wrap_typeddicts_as_models(field_type: Any, model_name_prefix: str) -> Any:
    """Route TypedDict input-field types through a strict Pydantic model.

    The raw TypedDict → Pydantic validator path does NOT honor extra='forbid',
    so typos inside nested request dicts are silently dropped. Wrapping the
    TypedDict in a strict Pydantic model closes that hole. Handles:
      - TypedDict (direct)
      - list[TypedDict]
      - list[Optional[TypedDict]]
    Top-level Optional[TypedDict] is already unwrapped by
    extract_*_param_info before this helper runs; the caller in
    create_func_models re-wraps with Optional[...] when is_optional is true.
    """
    if is_typeddict(field_type):
        return create_model_from_typeddict(
            field_type, f"{model_name_prefix}_{field_type.__name__}", strict=True
        )
    # Unwrap Optional[TypedDict] inside a list so the TypedDict can be wrapped
    # and then re-Optionalized.
    if is_strict_optional(field_type):
        inner = next(arg for arg in get_args(field_type) if arg is not type(None))
        wrapped = _wrap_typeddicts_as_models(inner, model_name_prefix)
        if wrapped is not inner:
            return Optional[wrapped]
        return field_type
    if get_origin(field_type) is list:
        args = get_args(field_type)
        if not args:
            # Bare `list` (no type parameter) — nothing to wrap.
            return field_type
        inner = args[0]
        wrapped = _wrap_typeddicts_as_models(inner, model_name_prefix)
        if wrapped is not inner:
            return list[wrapped]  # type: ignore[valid-type]
    return field_type


def create_func_models(func: Callable) -> tuple[type[BaseModel], type[BaseModel]]:
    """
    Analyze a function to create corresponding Pydantic models for its input and output.
    """
    input_fields: dict[str, Any] = {}
    # TODO figure this out (Sam)
    if asyncio.iscoroutinefunction(func) and hasattr(func, "__wrapped__"):
        func = func.__wrapped__
    model_prefix = snake_to_pascal_case(func.__name__)
    for name, param in inspect.signature(func, follow_wrapped=True).parameters.items():
        # Skip ToolContext parameters (including subclasses like arcade_mcp_server.Context)
        ann = param.annotation
        if isinstance(ann, type) and issubclass(ann, ToolContext):
            continue

        tool_field_info = extract_field_info(param)

        # A param is required iff the annotation is not Optional[...] AND
        # there is no explicit default in the signature. Uses the
        # has_explicit_default flag (not `default is not None`) so that
        # `def f(x: str = None)` is correctly treated as optional.
        is_required = not tool_field_info.is_optional and not tool_field_info.has_explicit_default

        # extract_field_info raises ToolInputSchemaError when the Annotated
        # description is missing, so by this point it is guaranteed non-None.
        description = tool_field_info.description

        # Route TypedDict fields (and list[TypedDict]) through Pydantic models
        # so the extra='forbid' rule applies inside nested request dicts too.
        # Include the parameter name in the model-name prefix so two params
        # of the same TypedDict type don't collide in Pydantic's $defs.
        field_type = _wrap_typeddicts_as_models(
            tool_field_info.field_type, f"{model_prefix}_{name}"
        )

        # extract_*_param_info unwraps Optional[T] to T before this point, so
        # re-wrap when the original annotation permitted None — otherwise the
        # field would be annotated as a non-nullable Pydantic model and callers
        # omitting the field (which falls back to default=None) would fail
        # validation at access time. Also re-wrap when the signature has an
        # explicit `= None` default on a non-Optional annotation
        # (`def f(x: str = None)`), which is legal Python and should behave as
        # `Optional[str]`.
        if tool_field_info.is_optional or (
            tool_field_info.has_explicit_default and tool_field_info.default is None
        ):
            field_type = Optional[field_type]

        if is_required:
            # Omit `default=` entirely — Pydantic treats this as a required field
            # and raises ValidationError when the caller omits it.
            input_fields[name] = (field_type, Field(description=description))
        else:
            input_fields[name] = (
                field_type,
                Field(default=tool_field_info.default, description=description),
            )

    input_model = create_model(
        f"{model_prefix}Input",
        __config__=ConfigDict(extra="forbid"),
        **input_fields,
    )

    output_model = determine_output_model(func)
    return input_model, output_model


def determine_output_model(func: Callable) -> type[BaseModel]:
    """
    Determine the output model for a function based on its return annotation.
    """
    return_annotation = inspect.signature(func).return_annotation
    output_model_name = f"{snake_to_pascal_case(func.__name__)}Output"

    # If the return annotation is empty, create a model with no fields
    if return_annotation is inspect.Signature.empty:
        return create_model(output_model_name)

    # If the return annotation has an __origin__ attribute
    # and does not have a __metadata__ attribute.
    # This is the case for TypedDicts.
    elif hasattr(return_annotation, "__origin__"):
        if hasattr(return_annotation, "__metadata__"):
            field_type = return_annotation.__args__[0]
            description = (
                return_annotation.__metadata__[0] if return_annotation.__metadata__ else ""
            )

            # Check if the field type is a TypedDict
            if is_typeddict(field_type):
                # Create a Pydantic model from TypedDict
                typeddict_model = create_model_from_typeddict(
                    field_type, f"{output_model_name}TypedDict"
                )
                return create_model(
                    output_model_name,
                    result=(
                        typeddict_model,
                        Field(
                            description=str(description)
                            if description
                            else "No description provided."
                        ),
                    ),
                )

            # If the return annotation has a description, use it
            if description:
                try:
                    return create_model(
                        output_model_name,
                        result=(field_type, Field(description=str(description))),
                    )
                except Exception:
                    raise ToolOutputSchemaError(
                        f"Unsupported output type '{field_type}'. Only built-in Python types, TypedDicts, "
                        "Pydantic models, and standard collections are supported as tool output types."
                    )

        # If the return annotation is a Union type
        origin = return_annotation.__origin__
        if origin is typing.Union:
            # For union types, create a model with the first non-None argument
            # TODO handle multiple non-None arguments. Raise error?
            for arg in get_args(return_annotation):
                if arg is not type(None):
                    # Check if the arg is a TypedDict
                    if is_typeddict(arg):
                        typeddict_model = create_model_from_typeddict(
                            arg, f"{output_model_name}TypedDict"
                        )
                        return create_model(
                            output_model_name,
                            result=(
                                typeddict_model,
                                Field(description="No description provided."),
                            ),
                        )
                    return create_model(
                        output_model_name,
                        result=(
                            arg,
                            Field(description="No description provided."),
                        ),
                    )

        # If the return annotation has an __origin__ attribute
        # and does not have a __metadata__ attribute.
        # This is the case for TypedDicts.
        return create_model(
            output_model_name,
            result=(
                return_annotation,
                Field(description="No description provided."),
            ),
        )
    else:
        # If the return annotation is a TypedDict
        if is_typeddict(return_annotation):
            typeddict_model = create_model_from_typeddict(return_annotation, output_model_name)
            return create_model(
                output_model_name,
                result=(
                    typeddict_model,
                    Field(description="No description provided."),
                ),
            )

        # If the return annotation is a simple type (like str)
        return create_model(
            output_model_name,
            result=(
                return_annotation,
                Field(description="No description provided."),
            ),
        )


class _TypedDictBaseModel(BaseModel):
    """Base for Pydantic models derived from TypedDict (output side).

    Drops unset fields at serialization time so absent optional fields
    (total=False) don't appear as None. Uses @model_serializer instead
    of overriding model_dump() because Pydantic v2's outer serializer
    does not call Python-level model_dump() on nested models — the
    override is only respected when the Python method is invoked
    directly.

    NOTE: this base does NOT set extra='forbid' — it is used for output
    models, and tools that return TypedDicts may include extra keys from
    upstream APIs that shouldn't fail serialization. Input-side strictness
    is applied via the `strict` parameter in create_model_from_typeddict.
    """

    @model_serializer(mode="wrap")
    def _drop_unset_fields(self, handler: Any) -> dict[str, Any]:
        result: dict[str, Any] = handler(self)
        for name in type(self).model_fields.keys() - self.model_fields_set:
            result.pop(name, None)
        return result

    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
        kwargs.setdefault("exclude_unset", True)
        return super().model_dump(**kwargs)


class _StrictTypedDictBaseModel(_TypedDictBaseModel):
    """Strict input-side base: rejects unknown keys.

    Typos inside nested request dicts (e.g. {"txt": "x"} instead of
    {"text": "x"}) raise ValidationError instead of being silently dropped.
    """

    model_config = ConfigDict(extra="forbid")


def create_model_from_typeddict(
    typeddict_class: type, model_name: str, strict: bool = False
) -> type[BaseModel]:
    """
    Create a Pydantic model from a TypedDict class.

    When strict=True, the generated model (and all nested TypedDict models)
    forbid unknown keys. Used for input-side TypedDicts so typos surface as
    ValidationError. Output-side callers leave strict=False (default) to
    preserve the pass-through behavior for tools whose return dicts contain
    extra keys from upstream APIs.
    """
    # Get type hints for the TypedDict
    type_hints = get_type_hints(typeddict_class, include_extras=True)

    # Build field definitions for the Pydantic model
    field_definitions: dict[str, Any] = {}
    for field_name, field_type in type_hints.items():
        # Check if field is required
        is_required = field_name in getattr(typeddict_class, "__required_keys__", set())

        # Unwrap Optional[T] (i.e. T | None) so we can detect nested TypedDicts
        is_optional_type = is_strict_optional(field_type)
        inner_type = field_type
        if is_optional_type:
            inner_type = next(arg for arg in get_args(field_type) if arg is not type(None))

        # Handle nested TypedDict (works for both T and Optional[T] after unwrapping)
        if is_typeddict(inner_type):
            nested_model = create_model_from_typeddict(
                inner_type, f"{model_name}_{field_name}", strict=strict
            )
            if is_required:
                if is_optional_type:
                    field_definitions[field_name] = (Optional[nested_model], Field())
                else:
                    field_definitions[field_name] = (nested_model, Field())
            else:
                field_definitions[field_name] = (Optional[nested_model], Field(default=None))
        else:
            # In strict mode, propagate wrapping into list[TypedDict] fields
            # so typos inside list elements also raise.
            effective_type = field_type
            if strict:
                effective_type = _wrap_typeddicts_as_models(
                    field_type, f"{model_name}_{field_name}"
                )
            if is_required:
                field_definitions[field_name] = (effective_type, Field())
            else:
                field_definitions[field_name] = (effective_type, Field(default=None))

    base = _StrictTypedDictBaseModel if strict else _TypedDictBaseModel
    return create_model(model_name, __base__=base, **field_definitions)


def to_tool_secret_requirements(
    secrets_requirement: list[str],
) -> list[ToolSecretRequirement]:
    # De-dupe case-insensitively but preserve the original casing for env var lookup
    unique_map: dict[str, str] = {}
    for name in secrets_requirement:
        lowered = str(name).lower()
        if lowered not in unique_map:
            unique_map[lowered] = str(name)
    return [ToolSecretRequirement(key=orig_name) for orig_name in unique_map.values()]


def to_tool_metadata_requirements(
    metadata_requirement: list[str],
) -> list[ToolMetadataRequirement]:
    # Iterate through the list, de-dupe case-insensitively, and convert each string to a ToolMetadataRequirement
    unique_metadata = {name.lower(): name.lower() for name in metadata_requirement}.values()
    return [ToolMetadataRequirement(key=name) for name in unique_metadata]
