"""Tool-related types and execution utilities for the Swarm framework."""

import asyncio
import inspect
import json
import logging as _logging
from enum import Enum
from functools import wraps
from typing import (
    Any,
    Awaitable,
    Callable,
    Dict,
    List,
    Optional,
    Type,
    Union,
    cast,
    overload,
)
from uuid import uuid4

from openai.types.chat import ChatCompletionMessageToolCall, ChatCompletionToolParam
from openai.types.shared_params import FunctionDefinition
from pydantic import BaseModel, Field, ValidationError

from swarm.agent import ToolFunctionType
from swarm.config.logging import get_structured_logger
from swarm.context import ContextProtocol, RunContext
from swarm.message import SwarmToolMessage
from swarm.retry import ToolRetryExecutor
from swarm.types import ContentType

logger = get_structured_logger(__name__)


# ------------------------------------------------------------
# Model Types
# ------------------------------------------------------------
class ContextAction(Enum):
    """How to handle context variables in tool results."""

    PRESERVE = "preserve"  # Keep existing context
    CLEAR = "clear"  # Clear all context
    UPDATE = "update"  # Update with new values


class ErrorType(str, Enum):
    """Classification of tool execution errors.

    Used to determine retry behavior and error message sanitization:
    - INTERNAL: Bugs/unexpected errors - retry once, sanitize before LLM sees
    - TRANSIENT: Network/timeout errors - retry with backoff
    - VALIDATION: LLM provided bad arguments - no retry, LLM should fix
    - BUSINESS: Expected business logic errors - no retry, explain to user
    - USER_INPUT: User needs to provide different input - no retry
    """

    INTERNAL = "internal"
    TRANSIENT = "transient"
    VALIDATION = "validation"
    BUSINESS = "business"
    USER_INPUT = "user_input"


class ResultMetadata(BaseModel):
    """
    Typed metadata for Result objects to control autonomous agent execution.

    Attributes:
        agent_completed: Signal that autonomous agent has finished its work
        agent_paused: Signal that autonomous agent needs to pause for user input
        completion_summary: Summary of what was completed (if agent_completed=True)
        pause_reason: Reason for pausing (if agent_paused=True)
    """

    agent_completed: bool = False
    agent_paused: bool = False
    completion_summary: Optional[str] = None
    pause_reason: Optional[str] = None


class ToolExecutionResponse(BaseModel):
    """Response from tool execution containing messages, agent handoffs, context updates, and metadata signals."""

    messages: List["SwarmToolMessage"] = []
    agent: Optional[Any] = None  # Avoid circular import
    context_variables: Dict[str, Any] = {}
    metadata: ResultMetadata = Field(
        default_factory=ResultMetadata
    )  # Aggregated metadata signals from all tool results
    stopped: bool = False  # True if execution was interrupted by stop signal


class Result(BaseModel):
    """
    Encapsulates the possible return values for an agent function.

    A result can be one of three types:
    - NATIVE: Simple string response
    - INTERACTIVE: JSON-serialized InteractiveValue
    - RICH: JSON-serialized RichContentValue

    Attributes:
        value (str): The result value, either plain text or JSON string
        agent: The agent instance, if applicable
        context_variables (dict): A dictionary of context variables
        context_action (ContextAction): How to handle context variables
        content_type (ContentType): The type of content (text/interactive/rich)
        component_type (ComponentType): UI component type for interactive/rich results
        is_error (bool): Whether this result represents an error

    Usage:
        Result.create(value="Hello, world!")
        Result(value="Some text", content_type=ContentType.TEXT)
    """

    value: str = ""
    agent: Optional[Any] = None  # Avoid circular import, use Any
    context_variables: Dict[str, Any] = {}
    context_action: ContextAction = ContextAction.UPDATE
    content_type: ContentType = ContentType.TEXT
    component_type: Optional[Any] = None  # App-specific component type
    is_error: bool = False
    error_type: Optional[ErrorType] = None  # Classification for retry/sanitization
    metadata: ResultMetadata = ResultMetadata()  # Agent control signals

    # Factory methods for creating different types of results
    @classmethod
    def create(cls, value: str, **kwargs: Any) -> "Result":
        """Create a native result with just a string value."""
        return cls(value=value, content_type=ContentType.TEXT, **kwargs)


# ------------------------------------------------------------
# Tool Execution Constants
# ------------------------------------------------------------
INTERNAL_ERROR_USER_MESSAGE = "An internal error occurred. Please try again."


# ------------------------------------------------------------
# Tool Execution
# ------------------------------------------------------------
async def execute_tools(
    tool_calls: List[ChatCompletionMessageToolCall],
    tools: List[ToolFunctionType],
    run_context: RunContext[ContextProtocol],
    current_agent_name: Optional[str] = None,
    *,
    hooks: Optional[List[Any]] = None,
    agent: Optional[Any] = None,
    check_stop_fn: Optional[Callable[[], Awaitable[bool]]] = None,
) -> ToolExecutionResponse:
    """Execute tool calls and return responses.

    Args:
        tool_calls: List of tool calls from the LLM
        tools: List of available tool functions
        run_context: Runtime context with all dependencies
        current_agent_name: Name of the currently active agent (for agent switch tracking)
        hooks: Lifecycle hooks to fire per-call (passed through from runner)
        agent: Agent reference for hook fire helpers
        check_stop_fn: Optional async callable returning True if stop was signalled
    """
    from swarm.hooks import fire_hook_on_tool_end, fire_hook_on_tool_start

    tool_execution_response = ToolExecutionResponse(
        messages=[], agent=None, context_variables={}
    )
    tools_map = {tool.__name__: tool for tool in tools}

    # Extract for internal use
    stream_manager = run_context.stream_manager

    for tool_call in tool_calls:
        # ── Stop check between individual tool calls ──────────────────
        if check_stop_fn is not None and await check_stop_fn():
            _logging.getLogger(__name__).info(
                "execute_tools.stopped_between_tools",
                extra={
                    "completed_tools": len(tool_execution_response.messages),
                    "total_tools": len(tool_calls),
                },
            )
            tool_execution_response.stopped = True
            return tool_execution_response
        # Malformed tool calls (missing function/name) — skip, no hooks
        if not tool_call.function or not tool_call.function.name:
            continue

        tool_name = tool_call.function.name
        tool_id = tool_call.id

        # Validate tool_id exists — skip, no hooks (no meaningful correlation id)
        if not tool_id:
            logger.warning(
                "Tool call missing ID for {tool_name}",
                extra={
                    "tool_name": tool_name,
                    "location": "swarm.tool.execute_tools",
                },
            )
            continue

        # --- Valid tool call from here: hooks fire for all paths below ---

        # Fast O(1) lookup
        func = tools_map.get(tool_name)
        if not func:
            # Unknown tool — fire hooks with error
            unknown_err = ValueError(f"Unknown function: {tool_name}")
            if agent is not None:
                await fire_hook_on_tool_start(agent, tool_id, tool_name, {})
                await fire_hook_on_tool_end(
                    agent, tool_id, tool_name, None, error=unknown_err
                )
            tool_execution_response.messages.append(
                SwarmToolMessage(
                    id=str(uuid4()),
                    tool_call_id=tool_id,
                    tool_name=tool_name,
                    content=f"Unknown function: {tool_name}",
                    is_error=True,
                )
            )
            continue

        # Parse arguments
        try:
            args = json.loads(tool_call.function.arguments or "{}")
        except json.JSONDecodeError as e:
            # Invalid JSON — fire hooks with error
            if agent is not None:
                await fire_hook_on_tool_start(agent, tool_id, tool_name, {})
                await fire_hook_on_tool_end(agent, tool_id, tool_name, None, error=e)
            tool_execution_response.messages.append(
                SwarmToolMessage(
                    id=str(uuid4()),
                    tool_call_id=tool_id,
                    tool_name=tool_name,
                    content=[{"type": "text", "text": f"Invalid arguments: {e}"}],
                    is_error=True,
                )
            )
            continue

        # Fire on_tool_start before execution
        if agent is not None:
            await fire_hook_on_tool_start(agent, tool_id, tool_name, args)

        # Execute tool with retry for transient errors
        await stream_manager.publish_status("Running ...")

        # Use ToolRetryExecutor for retry logic (policy required - pure DI)
        if run_context.retry_policy is None:
            raise ValueError(
                "retry_policy is required in RunContext. "
                "Ensure the caller provides a RetryPolicy instance."
            )
        retry_executor = ToolRetryExecutor(run_context.retry_policy)

        exec_error: BaseException | None = None
        try:
            result = await retry_executor.execute_with_retry(
                execute_tool,
                func,
                args,
                run_context,
                tool_name=tool_name,
            )
        except BaseException as e:
            exec_error = e
            raise
        finally:
            if agent is not None:
                await fire_hook_on_tool_end(
                    agent,
                    tool_id,
                    tool_name,
                    result if exec_error is None else None,
                    error=exec_error,
                )

        # Sanitize internal errors before they reach the LLM
        display_value = result.value
        if result.is_error and result.error_type == ErrorType.INTERNAL:
            display_value = INTERNAL_ERROR_USER_MESSAGE

        # Create response with properly formatted content
        tool_execution_response.messages.append(
            SwarmToolMessage(
                id=str(uuid4()),
                tool_call_id=tool_id,
                tool_name=tool_name,
                content=[
                    {"type": "text", "text": display_value if display_value else ""}
                ],
                component_type=result.component_type,
                content_type=result.content_type,
                is_error=result.is_error,
            )
        )

        # Update context (safely)
        match result.context_action:
            case ContextAction.CLEAR:
                tool_execution_response.context_variables.clear()
            case ContextAction.UPDATE:
                if result.context_variables:
                    tool_execution_response.context_variables.update(
                        result.context_variables
                    )
            case ContextAction.PRESERVE:
                pass

        # Aggregate metadata signals (OR logic - if any tool signals, propagate it)
        if result.metadata.agent_completed:
            tool_execution_response.metadata.agent_completed = True
            tool_execution_response.metadata.completion_summary = (
                result.metadata.completion_summary
            )
        if result.metadata.agent_paused:
            tool_execution_response.metadata.agent_paused = True
            tool_execution_response.metadata.pause_reason = result.metadata.pause_reason

        # Extract agent handoff signal if present
        if result.agent:
            tool_execution_response.agent = result.agent
            # Stamp agent switch info on the tool response message
            last_msg = tool_execution_response.messages[-1]
            last_msg.metadata = last_msg.metadata or {}
            last_msg.metadata["agent_switch"] = {
                "from_agent": current_agent_name or "unknown",
                "to_agent": result.agent.name,
            }

    return tool_execution_response


async def execute_tool(
    func: Callable[..., Any],
    args: Dict[str, Any],
    run_context: RunContext[ContextProtocol],
) -> Result:
    """Execute a tool with proper error handling and context management.

    Args:
        func: The tool function to execute
        args: Arguments to pass to the function
        run_context: Runtime context with all dependencies

    Returns:
        Result object with tool execution outcome
    """
    try:
        # Check function parameters and inject dependencies
        func_params = inspect.signature(func).parameters

        # Inject RunContext if tool expects ctx parameter
        if "ctx" in func_params:
            args["ctx"] = run_context

        # Execute function and await any nested coroutines
        raw_result = await _execute_and_await(func, args)

        return raw_result

    except ValidationError as e:
        # Pydantic validation error - LLM provided bad arguments
        error_message = f"Validation error in {func.__name__}: {str(e)}"
        logger.warning(error_message)
        return Result(
            value=error_message, is_error=True, error_type=ErrorType.VALIDATION
        )

    except (TimeoutError, asyncio.TimeoutError, ConnectionError, OSError) as e:
        # Network/transient errors - can retry
        error_message = f"Transient error in {func.__name__}: {str(e)}"
        logger.warning(error_message)
        return Result(
            value=error_message, is_error=True, error_type=ErrorType.TRANSIENT
        )

    except Exception as e:
        # Internal/unexpected errors - log full trace, sanitize for LLM
        error_message = f"Error executing tool {func.__name__}: {str(e)}"
        logger.error(error_message, exc_info=True)
        return Result(value=error_message, is_error=True, error_type=ErrorType.INTERNAL)


async def _execute_and_await(func: Callable[..., Any], args: Dict[str, Any]) -> Result:
    """Execute function and await any nested coroutines iteratively."""
    # Initial execution
    if asyncio.iscoroutinefunction(func):
        result = await func(**args)
    else:
        result = func(**args)

    # Keep awaiting until we get a non-coroutine
    while asyncio.iscoroutine(result):
        result = await result

    if not isinstance(result, Result):
        logger.warning(
            "Tool returned non-Result object",
            extra={"result": result},
        )
        return Result(value=str(result))

    return result


# ------------------------------------------------------------
# Tool Pre-Completion processing
# ------------------------------------------------------------
def prepare_tools_for_completion(
    tools: List[Callable[..., Any]],
) -> List[Dict[str, Any]]:
    """Convert agent functions to OpenAI tool format.

    Args:
        tools: List of callable functions

    Returns:
        List of tool definitions in OpenAI format
    """
    # Special parameter names that should not be exposed to the model
    special_param_names = {"ctx"}

    def _strip_special_params(schema: Dict[str, Any]) -> Dict[str, Any]:
        """Remove special runtime parameters from tool schema."""
        try:
            fn = schema.get("function", {})
            params = fn.get("parameters", {})
            properties = params.get("properties", {})
            required = params.get("required", [])

            if isinstance(properties, dict):
                for name in special_param_names:
                    properties.pop(name, None)

            if isinstance(required, list):
                # Remove special params from required list
                params["required"] = [
                    name for name in required if name not in special_param_names
                ]
        except Exception:
            # Best-effort cleanup; do not fail tool preparation
            pass
        return schema

    prepared: List[Dict[str, Any]] = []
    for func in tools:
        # Build base schema from the function metadata
        schema = function_to_json(func)
        # Ensure special runtime-only params are not exposed
        cleaned_schema = _strip_special_params(schema)
        prepared.append(cleaned_schema)

    return prepared


# ------------------------------------------------------------
# Tool Creation Utilities
# ------------------------------------------------------------


def clean_schema_for_openai(
    schema: Union[Dict[str, Any], List[Any], str, bool, None],
) -> Union[Dict[str, Any], List[Any], str, bool, None]:
    """
    Clean up a JSON schema to match OpenAI's requirements.

    Handles:
    - Converting allOf to anyOf
    - Ensuring additionalProperties: false for objects
    - Removing unsupported keywords
    - Ensuring proper enum handling
    - Setting all properties as required
    """
    if not isinstance(schema, dict):
        return schema

    # List of supported types
    SUPPORTED_TYPES = {
        "string",
        "number",
        "boolean",
        "integer",
        "object",
        "array",
        "null",
    }

    # Keywords not supported by OpenAI
    UNSUPPORTED_KEYWORDS = {
        "minLength",
        "maxLength",
        "pattern",
        "format",  # string
        "minimum",
        "maximum",
        "multipleOf",  # number
        "patternProperties",
        "unevaluatedProperties",
        "propertyNames",  # object
        "minProperties",
        "maxProperties",
        "unevaluatedItems",
        "contains",
        "minContains",
        "maxContains",  # array
        "minItems",
        "maxItems",
        "uniqueItems",
    }

    # Handle non-dict schemas
    if not isinstance(schema, dict):
        return schema

    cleaned: Dict[str, Any] = {}

    for key, value in schema.items():
        # Skip unsupported keywords
        if key in UNSUPPORTED_KEYWORDS:
            continue

        # Convert allOf to anyOf
        if key == "allOf":
            cleaned["anyOf"] = [clean_schema_for_openai(item) for item in value]
            continue

        # Handle nested objects and arrays
        if isinstance(value, dict):
            cleaned[key] = clean_schema_for_openai(value)
        elif isinstance(value, list) and key != "enum":
            cleaned[key] = [clean_schema_for_openai(item) for item in value]
        else:
            cleaned[key] = value

    # Handle type restrictions
    if "type" in cleaned:
        type_value = cleaned["type"]
        # Convert unknown types to string
        if isinstance(type_value, str) and type_value not in SUPPORTED_TYPES:
            cleaned["type"] = "string"
        # Handle array of types
        elif isinstance(type_value, list):
            # If 'null' is in types, convert to anyOf
            if "null" in type_value:
                non_null_types = [t for t in type_value if t != "null"]
                if len(non_null_types) == 1:
                    cleaned["type"] = non_null_types[0]
                else:
                    cleaned["anyOf"] = [{"type": t} for t in non_null_types]
                    del cleaned["type"]  # Remove type when using anyOf
            else:
                # Use the first supported type
                cleaned["type"] = next(
                    (t for t in type_value if t in SUPPORTED_TYPES), "string"
                )

    # Handle enum type special case from Pydantic
    if "enum" in schema:
        if "type" not in cleaned:
            cleaned["type"] = "string"
        cleaned["enum"] = schema["enum"]

    # Ensure additionalProperties is false for objects
    if schema.get("type") == "object":
        cleaned["additionalProperties"] = False

        # Make all properties required for objects
        if "properties" in cleaned:
            cleaned["required"] = list(cleaned["properties"].keys())

    return cleaned


def create_function_schema(
    func: Callable[..., Any],
    input_model: Optional[Type[BaseModel]],
    name: str,
    description: str,
) -> ChatCompletionToolParam:
    """
    Creates an OpenAI function schema from a Python function and its input model.
    """
    # Handle case where input model is None
    if input_model is None:
        parameters = {
            "type": "object",
            "properties": {},
            "required": [],
            "additionalProperties": False,
        }
    else:
        raw_schema = input_model.model_json_schema()
        parameters = clean_schema_for_openai(raw_schema)  # type: ignore

    # Create the function definition
    function_def = FunctionDefinition(
        name=name,
        description=description,
        parameters=parameters,
        strict=True,
    )

    return ChatCompletionToolParam(type="function", function=function_def)


@overload
def function_tool(
    func: Callable[..., Any],
    *,
    name: Optional[str] = None,
    description: Optional[str] = None,
    input_model: Optional[Type[BaseModel]] = None,
) -> Callable[..., Any]:
    """Overload for usage as @function_tool (no parentheses)."""
    ...


@overload
def function_tool(
    *,
    name: Optional[str] = None,
    description: Optional[str] = None,
    input_model: Optional[Type[BaseModel]] = None,
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """Overload for usage as @function_tool(...)."""
    ...


def function_tool(
    func: Optional[Callable[..., Any]] = None,
    *,
    name: Optional[str] = None,
    description: Optional[str] = None,
    input_model: Optional[Type[BaseModel]] = None,
) -> Union[Callable[..., Any], Callable[[Callable[..., Any]], Callable[..., Any]]]:
    """
    Simplified function tool decorator that replaces @function_schema.

    Args:
        func: The function to wrap
        name: Tool name (defaults to function name)
        description: Tool description (defaults to docstring)
        input_model: Optional Pydantic model for input validation

    Returns:
        Decorated function with schema metadata
    """

    def decorator(target_func: Callable[..., Any]) -> Callable[..., Any]:
        # Extract metadata
        tool_name = name or target_func.__name__
        tool_description = description or target_func.__doc__ or ""

        # Store metadata on function
        tf = cast(Any, target_func)
        tf.__input_model__ = input_model
        tf.__tool_name__ = tool_name
        tf.__tool_description__ = tool_description

        @wraps(target_func)
        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
            if input_model is None:
                return await target_func(*args, **kwargs)

            # Extract validation kwargs (exclude ctx which is passed separately)
            validation_kwargs = {k: v for k, v in kwargs.items() if k != "ctx"}

            try:
                # Creates and validates Pydantic model
                validated_data = input_model(**validation_kwargs)
                # Pass model directly as first parameter, then ctx if present
                call_kwargs = {}
                if "ctx" in kwargs:
                    call_kwargs["ctx"] = kwargs["ctx"]
                # Pass validated model as first positional arg
                return await target_func(validated_data, **call_kwargs)
            except ValidationError as e:
                return Result(value=f"Invalid input: {str(e)}", is_error=True)
            except Exception as e:
                return Result(value=f"Error executing tool: {str(e)}", is_error=True)

        @wraps(target_func)
        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
            if input_model is None:
                return target_func(*args, **kwargs)

            # Extract validation kwargs (exclude ctx which is passed separately)
            validation_kwargs = {k: v for k, v in kwargs.items() if k != "ctx"}

            try:
                validated_data = input_model(**validation_kwargs)
                # Pass model directly as first parameter, then ctx if present
                call_kwargs = {}
                if "ctx" in kwargs:
                    call_kwargs["ctx"] = kwargs["ctx"]
                # Pass validated model as first positional arg
                return target_func(validated_data, **call_kwargs)
            except ValidationError as e:
                return Result(value=f"Invalid input: {str(e)}", is_error=True)
            except Exception as e:
                return Result(value=f"Error executing tool: {str(e)}", is_error=True)

        # Choose appropriate wrapper based on function type
        wrapper: Callable[..., Any] = (
            async_wrapper if asyncio.iscoroutinefunction(target_func) else sync_wrapper
        )

        # Copy metadata to wrapper
        w = cast(Any, wrapper)
        w.__input_model__ = input_model
        w.__tool_name__ = tool_name
        w.__tool_description__ = tool_description

        return wrapper

    # Handle both @function_tool and @function_tool(...) usage
    if func is not None:
        return decorator(func)

    return decorator


def function_to_json(func: Callable[..., Any]) -> Dict[str, Any]:
    """
    Converts a function with @function_tool decorator to JSON schema format.
    """
    if hasattr(func, "__input_model__") and hasattr(func, "__tool_name__"):
        # Function has @function_tool decorator
        schema = create_function_schema(
            func,
            cast(Any, func).__input_model__,
            cast(Any, func).__tool_name__,
            cast(Any, func).__tool_description__,
        )
        return dict(schema)

    # Fallback to inspection-based approach for undecorated functions
    type_map = {
        str: "string",
        int: "integer",
        float: "number",
        bool: "boolean",
        list: "array",
        dict: "object",
        type(None): "null",
    }

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

    parameters = {}
    for param in signature.parameters.values():
        try:
            param_type = type_map.get(param.annotation, "string")
        except KeyError as e:
            raise KeyError(
                f"Unknown type annotation {param.annotation} for "
                f"parameter {param.name}: {str(e)}"
            ) from e
        parameters[param.name] = {"type": param_type}

    required = [
        param.name
        for param in signature.parameters.values()
        if param.default == inspect._empty
    ]

    return {
        "type": "function",
        "function": {
            "name": func.__name__,
            "description": func.__doc__ or "",
            "parameters": {
                "type": "object",
                "properties": parameters,
                "required": required,
            },
        },
    }
