################################################################################
#  Licensed to the Apache Software Foundation (ASF) under one
#  or more contributor license agreements.  See the NOTICE file
#  distributed with this work for additional information
#  regarding copyright ownership.  The ASF licenses this file
#  to you under the Apache License, Version 2.0 (the
#  "License"); you may not use this file except in compliance
#  with the License.  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################
import typing
from inspect import signature
from typing import Any, Callable, Dict, Optional, Type, Union

from docstring_parser import parse
from mcp import types
from pydantic import BaseModel, create_model
from pydantic.fields import Field, FieldInfo


def create_schema_from_function(name: str, func: Callable) -> Type[BaseModel]:
    """Create a pydantic schema from a function's signature."""
    docstr = func.__doc__

    docstr = parse(docstr)
    doc_params = {}
    for param in docstr.params:
        doc_params[param.arg_name] = param

    fields = {}
    params = signature(func).parameters
    for param_name in params:
        param_type = params[param_name].annotation
        param_default = params[param_name].default
        description = doc_params.get(param_name)
        if description is not None:
            description = description.description
        else:
            description = f"Parameter: {param_name}"

        if typing.get_origin(param_type) is typing.Annotated:
            args = typing.get_args(param_type)
            param_type = args[0]
            if isinstance(args[1], str):
                description = args[1]
            elif isinstance(args[1], FieldInfo):
                description = args[1].description

        if param_type is params[param_name].empty:
            param_type = typing.Any

        if param_default is params[param_name].empty:
            # Required field
            fields[param_name] = (param_type, FieldInfo(description=description))
        elif isinstance(param_default, FieldInfo):
            # Field with pydantic.Field as default value
            fields[param_name] = (param_type, param_default)
        else:
            fields[param_name] = (
                param_type,
                FieldInfo(default=param_default, description=description),
            )

    return create_model(name, **fields)


TYPE_MAPPING: dict[str, type] = {
    "string": str,
    "integer": int,
    "number": float,
    "boolean": bool,
    "object": dict,
    "array": list,
    "null": type(None),
}

CONSTRAINT_MAPPING: dict[str, str] = {
    "minimum": "ge",
    "maximum": "le",
    "exclusiveMinimum": "gt",
    "exclusiveMaximum": "lt",
    "inclusiveMinimum": "ge",
    "inclusiveMaximum": "le",
    "minItems": "min_length",
    "maxItems": "max_length",
}


def __get_field_params_from_field_schema(field_schema: dict) -> dict:
    """Gets Pydantic field parameters from a JSON schema field."""
    field_params = {}
    for constraint, constraint_value in CONSTRAINT_MAPPING.items():
        if constraint in field_schema:
            field_params[constraint_value] = field_schema[constraint]
    if "description" in field_schema:
        field_params["description"] = field_schema["description"]
    if "default" in field_schema:
        field_params["default"] = field_schema["default"]
    return field_params


def create_model_from_schema(name: str, schema: dict) -> type[BaseModel]:
    """Create Pydantic model from a JSON schema generated by
    BaseModel.model_json_schema().
    """
    models: dict[str, type[BaseModel]] = {}

    def resolve_field_type(field_schema: dict) -> type[typing.Any]:
        """Resolves field type, including optional types and nullability."""
        if "$ref" in field_schema:
            model_reference = field_schema["$ref"].split("/")[-1]
            return models.get(model_reference, Any)  #

        if "anyOf" in field_schema:
            types = [
                TYPE_MAPPING.get(t["type"], typing.Any)
                for t in field_schema["anyOf"]
                if t.get("type")
            ]
            if type(None) in types:
                types.remove(type(None))
                if len(types) == 1:
                    return typing.Optional[types[0]] # noqa: UP007
                return Optional[tuple(types)]  # # noqa: UP007
            else:
                return Union[tuple(types)]  # noqa: UP007
        field_type = TYPE_MAPPING.get(field_schema.get("type"), typing.Any)  # type: ignore[arg-type]

        # Handle arrays (lists)
        if field_schema.get("type") == "array":
            items = field_schema.get("items", {})
            item_type = resolve_field_type(items)
            return list[item_type]  # type: ignore[valid-type]

        # Handle objects (dicts with specified value types)
        if field_schema.get("type") == "object":
            additional_props = field_schema.get("additionalProperties")
            value_type = (
                resolve_field_type(additional_props) if additional_props else typing.Any
            )
            return dict[str, value_type]  # type: ignore[valid-type]

        return field_type  # type: ignore[return-value]

    # First, create models for definitions
    definitions = schema.get("$defs", {})
    for model_name, model_schema in definitions.items():
        fields = {}
        for field_name, field_schema in model_schema.get("properties", {}).items():
            field_type = resolve_field_type(field_schema=field_schema)
            field_params = __get_field_params_from_field_schema(
                field_schema=field_schema
            )
            fields[field_name] = (field_type, Field(**field_params))

        models[model_name] = create_model(
            model_name, **fields, __doc__=model_schema.get("description", "")
        )  # type: ignore[call-overload]

    # Now, create the main model, resolving references
    main_fields = {}
    for field_name, field_schema in schema.get("properties", {}).items():
        if "$ref" in field_schema:
            model_reference = field_schema["$ref"].split("/")[-1]
            field_type = models.get(model_reference, Any)  # type: ignore[arg-type]
        else:
            field_type = resolve_field_type(field_schema=field_schema)

        field_params = __get_field_params_from_field_schema(field_schema=field_schema)
        main_fields[field_name] = (field_type, Field(**field_params))

    return create_model(name, **main_fields, __doc__=schema.get("description", ""))


def extract_mcp_content_item(content_item: Any) -> Dict[str, Any] | str:
    """Extract and normalize a single MCP content item.

    Args:
        content_item: A single MCP content item (TextContent, ImageContent, etc.)

    Returns:
        Dict representation of the content item

    Raises:
        ImportError: If MCP types are not available
    """
    if types is None:
        err_msg = "MCP types not available. Please install the mcp package."
        raise ImportError(err_msg)

    if isinstance(content_item, types.TextContent):
        return content_item.text
    elif isinstance(content_item, types.ImageContent):
        return {
            "type": "image",
            "data": content_item.data,
            "mimeType": content_item.mimeType
        }
    elif isinstance(content_item, types.EmbeddedResource):
        if isinstance(content_item.resource, types.TextResourceContents):
            return {
                "type": "resource",
                "uri": content_item.resource.uri,
                "text": content_item.resource.text
            }
        elif isinstance(content_item.resource, types.BlobResourceContents):
            return {
                "type": "resource",
                "uri": content_item.resource.uri,
                "blob": content_item.resource.blob
            }
    else:
        # Handle unknown content types as generic dict
        return content_item.model_dump() if hasattr(content_item, 'model_dump') else str(content_item)