import dataclasses
import inspect
from typing import TYPE_CHECKING, Annotated, Any, Callable, List, Literal, Optional, Type, Union, get_args, get_origin

from pydantic import BaseModel
from pydantic_core import PydanticUndefined
from pydash import snake_case

from vellum import Vellum
from vellum.client.types.function_definition import FunctionDefinition
from vellum.workflows.integrations.composio_service import ComposioService
from vellum.workflows.integrations.mcp_service import MCPService
from vellum.workflows.integrations.vellum_integration_service import VellumIntegrationService
from vellum.workflows.types.definition import (
    ComposioToolDefinition,
    DeploymentDefinition,
    MCPServer,
    MCPToolDefinition,
    VellumIntegrationToolDefinition,
)
from vellum.workflows.utils.vellum_variables import vellum_variable_type_to_openapi_type

if TYPE_CHECKING:
    from vellum.workflows.workflows.base import BaseWorkflow

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

for k, v in list(type_map.items()):
    if isinstance(k, type):
        type_map[k.__name__] = v


def compile_annotation(annotation: Optional[Any], defs: dict[str, Any]) -> dict:
    if annotation is None:
        return {"type": "null"}

    if get_origin(annotation) is Union:
        return {"anyOf": [compile_annotation(a, defs) for a in get_args(annotation)]}

    if get_origin(annotation) is Literal:
        values = list(get_args(annotation))
        types = {type(value) for value in values}
        if len(types) == 1:
            value_type = types.pop()
            if value_type in type_map:
                return {"type": type_map[value_type], "enum": values}
            else:
                return {"enum": values}
        else:
            return {"enum": values}

    if get_origin(annotation) is dict:
        _, value_type = get_args(annotation)
        return {"type": "object", "additionalProperties": compile_annotation(value_type, defs)}

    if get_origin(annotation) is list:
        item_type = get_args(annotation)[0]
        return {"type": "array", "items": compile_annotation(item_type, defs)}

    if get_origin(annotation) is tuple:
        args = get_args(annotation)
        if len(args) == 2 and args[1] is Ellipsis:
            # Tuple[int, ...] with homogeneous items
            return {"type": "array", "items": compile_annotation(args[0], defs)}
        else:
            # Tuple[int, str] with fixed length items
            result = {
                "type": "array",
                "prefixItems": [compile_annotation(arg, defs) for arg in args],
                "minItems": len(args),
                "maxItems": len(args),
            }
            return result

    if dataclasses.is_dataclass(annotation) and isinstance(annotation, type):
        if annotation.__name__ not in defs:
            properties = {}
            required = []
            for field in dataclasses.fields(annotation):
                properties[field.name] = compile_annotation(field.type, defs)
                if field.default is dataclasses.MISSING:
                    required.append(field.name)
                else:
                    properties[field.name]["default"] = _compile_default_value(field.default)
            defs[annotation.__name__] = {"type": "object", "properties": properties, "required": required}
        return {"$ref": f"#/$defs/{annotation.__name__}"}

    if inspect.isclass(annotation) and issubclass(annotation, BaseModel):
        if annotation.__name__ not in defs:
            properties = {}
            required = []
            for field_name, field_info in annotation.model_fields.items():
                # field_info is a FieldInfo object which has an annotation attribute
                properties[field_name] = compile_annotation(field_info.annotation, defs)

                if field_info.description is not None:
                    properties[field_name]["description"] = field_info.description

                if field_info.default is PydanticUndefined:
                    required.append(field_name)
                else:
                    properties[field_name]["default"] = _compile_default_value(field_info.default)
            defs[annotation.__name__] = {"type": "object", "properties": properties, "required": required}

        return {"$ref": f"#/$defs/{annotation.__name__}"}

    return {"type": type_map[annotation]}


def _compile_default_value(default: Any) -> Any:
    if dataclasses.is_dataclass(default):
        return {
            field.name: _compile_default_value(getattr(default, field.name)) for field in dataclasses.fields(default)
        }

    if isinstance(default, BaseModel):
        return {
            field_name: _compile_default_value(getattr(default, field_name))
            for field_name in default.__class__.model_fields.keys()
        }

    return default


def _compile_workflow_deployment_input(input_var: Any) -> dict[str, Any]:
    """
    Converts a deployment workflow input variable to a JSON schema type definition.
    """
    primitive_type = vellum_variable_type_to_openapi_type(input_var.type)
    input_schema = {"type": primitive_type}

    if input_var.default is not None:
        input_schema["default"] = input_var.default.value

    return input_schema


def compile_function_definition(function: Callable) -> FunctionDefinition:
    """
    Converts a Python function into our Vellum-native FunctionDefinition type.

    Args:
        function: The Python function to compile
    """

    try:
        signature = inspect.signature(function)
    except ValueError as e:
        raise ValueError(f"Failed to get signature for function {function.__name__}: {str(e)}")

    # Get inputs from the decorator if present
    inputs = getattr(function, "__vellum_inputs__", {})
    exclude_params = set(inputs.keys())

    properties = {}
    required = []
    defs: dict[str, Any] = {}
    for param in signature.parameters.values():
        # Skip parameters that are in the exclude_params set
        if exclude_params and param.name in exclude_params:
            continue

        # Check if parameter uses Annotated type hint
        if get_origin(param.annotation) is Annotated:
            args = get_args(param.annotation)
            actual_type = args[0]
            # Extract description from metadata
            description = args[1] if len(args) > 1 and isinstance(args[1], str) else None

            properties[param.name] = compile_annotation(actual_type, defs)
            if description:
                properties[param.name]["description"] = description
        else:
            properties[param.name] = compile_annotation(param.annotation, defs)

        if param.default is inspect.Parameter.empty:
            required.append(param.name)
        else:
            properties[param.name]["default"] = _compile_default_value(param.default)

    parameters = {"type": "object", "properties": properties, "required": required}
    if defs:
        parameters["$defs"] = defs

    return FunctionDefinition(
        name=function.__name__,
        description=function.__doc__,
        parameters=parameters,
    )


def compile_inline_workflow_function_definition(workflow_class: Type["BaseWorkflow"]) -> FunctionDefinition:
    """
    Converts a base workflow class into our Vellum-native FunctionDefinition type.
    """
    # Get the inputs class for the workflow
    inputs_class = workflow_class.get_inputs_class()
    vars_inputs_class = vars(inputs_class)

    properties = {}
    required = []
    defs: dict[str, Any] = {}

    for name, field_type in inputs_class.__annotations__.items():
        if name.startswith("__"):
            continue

        properties[name] = compile_annotation(field_type, defs)

        # Check if the field has a default value
        if name not in vars_inputs_class:
            required.append(name)
        else:
            # Field has a default value
            properties[name]["default"] = vars_inputs_class[name]

    parameters = {"type": "object", "properties": properties, "required": required}
    if defs:
        parameters["$defs"] = defs

    return FunctionDefinition(
        name=snake_case(workflow_class.__name__),
        description=workflow_class.__doc__,
        parameters=parameters,
    )


def compile_workflow_deployment_function_definition(
    deployment_definition: DeploymentDefinition,
    vellum_client: Vellum,
) -> FunctionDefinition:
    """
    Converts a deployment workflow config into our Vellum-native FunctionDefinition type.

    Args:
        deployment_definition: DeploymentDefinition instance
        vellum_client: Vellum client instance
    """
    release_info = deployment_definition.get_release_info(vellum_client)

    name = release_info["name"]
    description = release_info["description"]
    input_variables = release_info["input_variables"]

    properties = {}
    required = []

    for input_var in input_variables:
        properties[input_var.key] = _compile_workflow_deployment_input(input_var)

        if input_var.required and input_var.default is None:
            required.append(input_var.key)

    parameters = {"type": "object", "properties": properties, "required": required}

    return FunctionDefinition(
        name=name.replace("-", ""),
        description=description,
        parameters=parameters,
    )


def get_mcp_tool_name(tool_def: MCPToolDefinition) -> str:
    """Generate a unique name for an MCP tool by combining server and tool names."""
    server_name = snake_case(tool_def.server.name)
    return f"{server_name}__{tool_def.name}"


def compile_mcp_tool_definition(server_def: MCPServer) -> List[MCPToolDefinition]:
    """Hydrate an MCPToolDefinition with detailed information from the MCP server.

    We do tool discovery on the MCP server to get the tool definitions.

    Args:
        server_def: The basic MCPToolDefinition to enhance

    Returns:
        MCPToolDefinition with detailed parameters and description
    """
    try:
        mcp_service = MCPService()
        return mcp_service.hydrate_tool_definitions(server_def)
    except Exception:
        return []


def compile_composio_tool_definition(tool_def: ComposioToolDefinition) -> FunctionDefinition:
    """Hydrate a ComposioToolDefinition with detailed information from the Composio API.

    Args:
        tool_def: The basic ComposioToolDefinition to enhance

    Returns:
        FunctionDefinition with detailed parameters and description
    """
    try:
        composio_service = ComposioService()
        tool_details = composio_service.get_tool_by_slug(tool_def.action)

        # Create a FunctionDefinition directly with proper field extraction
        return FunctionDefinition(
            name=tool_def.name,
            description=tool_details.get("description", tool_def.description),
            parameters=tool_details.get("input_parameters", {}),
        )
    except Exception:
        # If hydration fails (including no API key), return basic function definition
        return FunctionDefinition(
            name=tool_def.name,
            description=tool_def.description,
            parameters={},
        )


def compile_vellum_integration_tool_definition(
    tool_def: VellumIntegrationToolDefinition,
    vellum_client: Vellum,
) -> FunctionDefinition:
    """Compile a VellumIntegrationToolDefinition into a FunctionDefinition.

    Args:
        tool_def: The VellumIntegrationToolDefinition to compile
        vellum_client: Vellum client instance

    Returns:
        FunctionDefinition with tool parameters and description
    """
    try:
        service = VellumIntegrationService(vellum_client)
        tool_details = service.get_tool_definition(
            integration=tool_def.integration_name,
            provider=tool_def.provider.value,
            tool_name=tool_def.name,
        )

        return FunctionDefinition(
            name=tool_def.name,
            description=tool_details.description,
            parameters=tool_details.parameters or {},
        )
    except Exception:
        # Fallback for service failures
        return FunctionDefinition(name=tool_def.name, description=tool_def.description, parameters={})


def use_tool_inputs(**inputs):
    """
    Decorator to specify which parameters of a tool function should be provided
    from the parent workflow inputs rather than from the LLM.

    Args:
        **inputs: Mapping of parameter names to parent input references

    Example:
        @use_tool_inputs(
            parent_input=ParentInputs.parent_input,
        )
        def get_string(parent_input: str, user_query: str) -> str:
            return f"Parent: {parent_input}, Query: {user_query}"
    """

    def decorator(func: Callable) -> Callable:
        # Store the inputs mapping on the function for later use
        setattr(func, "__vellum_inputs__", inputs)
        return func

    return decorator
