"""
Tool executor — schema-validated tool dispatch with audit logging.

execute_tool() is the single entry point for all tool invocations. It:
  1. Looks up the tool in the registry
  2. Parses the tool call arguments (LLM output — always untrusted JSON)
  3. Validates args against the tool's JSON Schema (rejects invalid input)
  4. If requires_confirmation=True, returns a confirmation request without executing
  5. Calls the tool handler
  6. Logs the invocation (success or failure) to the audit trail
  7. Returns the result string

CRITICAL: Tool arguments come from LLM output, which is untrusted. Schema
validation MUST happen before any handler is called. Invalid args are rejected
with an error message — never silently coerced.

CRITICAL: All tool execution happens within the single Celery task context.
Never dispatch separate Celery tasks for tool execution.
"""

from __future__ import annotations

import json
import logging
import time
import uuid
from typing import TYPE_CHECKING, Any

import jsonschema

if TYPE_CHECKING:
    from orchestrator.audit.logger import AuditLogger
    from orchestrator.tools.registry import ToolDefinition

logger = logging.getLogger(__name__)

_CONFIRMATION_MESSAGE_TEMPLATE = (
    "This action requires your approval before I proceed:\n\n"
    "**Tool:** {tool_name}\n"
    "**Arguments:** {args_summary}\n\n"
    "Please reply **yes** to confirm or **no** to cancel."
)


async def execute_tool(
    tool_call: dict[str, Any],
    registry: dict[str, "ToolDefinition"],
    tenant_id: uuid.UUID,
    agent_id: uuid.UUID,
    audit_logger: "AuditLogger",
) -> str:
    """
    Execute a tool call from LLM output.

    Args:
        tool_call:    LLM tool call dict. Must contain:
                      {"function": {"name": "...", "arguments": "{ JSON string }"}}
        registry:     Tool registry dict (name → ToolDefinition).
        tenant_id:    Tenant UUID for audit logging.
        agent_id:     Agent UUID for audit logging.
        audit_logger: AuditLogger instance for recording the invocation.

    Returns:
        Tool result as a string. Returns an error message string (not an
        exception) for invalid args, unknown tools, or handler errors.
    """
    # Extract tool name and raw arguments
    function_data = tool_call.get("function", {})
    tool_name: str = function_data.get("name", "")
    raw_arguments: str = function_data.get("arguments", "{}")

    # ------------------------------------------------------------------
    # 1. Look up tool in registry
    # ------------------------------------------------------------------
    if tool_name not in registry:
        error_msg = f"Unknown tool: '{tool_name}'. Available tools: {', '.join(registry.keys())}"
        logger.warning("execute_tool: %s", error_msg)
        return error_msg

    tool = registry[tool_name]

    # ------------------------------------------------------------------
    # 2. Parse JSON arguments (LLM output is untrusted)
    # ------------------------------------------------------------------
    try:
        args: dict[str, Any] = json.loads(raw_arguments)
        if not isinstance(args, dict):
            raise ValueError("Tool arguments must be a JSON object")
    except (json.JSONDecodeError, ValueError) as exc:
        error_msg = f"Invalid tool arguments for '{tool_name}': {exc}"
        logger.warning("execute_tool: %s", error_msg)
        await _log_failure(audit_logger, tool_name, {}, error_msg, tenant_id, agent_id)
        return f"Error: {error_msg}"

    # ------------------------------------------------------------------
    # 3. Schema validation — reject invalid args before calling handler
    # ------------------------------------------------------------------
    validation_error: str | None = _validate_args(args, tool.parameters)
    if validation_error:
        error_msg = f"Invalid arguments for '{tool_name}': {validation_error}"
        logger.warning("execute_tool schema validation failed: %s", error_msg)
        await _log_failure(audit_logger, tool_name, args, error_msg, tenant_id, agent_id)
        return f"Error: {error_msg}"

    # ------------------------------------------------------------------
    # 4. Confirmation gate — pause if tool requires user approval
    # ------------------------------------------------------------------
    if tool.requires_confirmation:
        try:
            args_summary = json.dumps(args, ensure_ascii=False)
        except Exception:
            args_summary = repr(args)

        confirmation_msg = _CONFIRMATION_MESSAGE_TEMPLATE.format(
            tool_name=tool_name,
            args_summary=args_summary,
        )
        # Don't log a tool invocation — this is a confirmation request, not an execution
        return confirmation_msg

    # ------------------------------------------------------------------
    # 5. Execute the handler
    # ------------------------------------------------------------------
    start_ms = time.monotonic()
    try:
        result: str = await tool.handler(**args)
        latency_ms = int((time.monotonic() - start_ms) * 1000)

        await audit_logger.log_tool_call(
            tool_name=tool_name,
            args=args,
            result=result,
            tenant_id=tenant_id,
            agent_id=agent_id,
            latency_ms=latency_ms,
            error=None,
        )
        return result

    except Exception as exc:
        latency_ms = int((time.monotonic() - start_ms) * 1000)
        error_msg = f"{type(exc).__name__}: {exc}"
        logger.exception("Tool handler '%s' raised exception: %s", tool_name, error_msg)

        await audit_logger.log_tool_call(
            tool_name=tool_name,
            args=args,
            result=None,
            tenant_id=tenant_id,
            agent_id=agent_id,
            latency_ms=latency_ms,
            error=error_msg,
        )
        return f"Tool '{tool_name}' encountered an error: {error_msg}"


def _validate_args(args: dict[str, Any], schema: dict[str, Any]) -> str | None:
    """
    Validate args against a JSON Schema.

    Returns None if valid, or an error message string if invalid.
    """
    try:
        jsonschema.validate(instance=args, schema=schema)
        return None
    except jsonschema.ValidationError as exc:
        # Return a concise validation error message
        return exc.message
    except jsonschema.SchemaError as exc:
        logger.error("Tool schema is invalid: %s", exc.message)
        return f"Schema error: {exc.message}"


async def _log_failure(
    audit_logger: "AuditLogger",
    tool_name: str,
    args: dict[str, Any],
    error_msg: str,
    tenant_id: uuid.UUID,
    agent_id: uuid.UUID,
) -> None:
    """Log a tool invocation failure to the audit trail."""
    try:
        await audit_logger.log_tool_call(
            tool_name=tool_name,
            args=args,
            result=None,
            tenant_id=tenant_id,
            agent_id=agent_id,
            latency_ms=0,
            error=error_msg,
        )
    except Exception:
        logger.exception("Failed to write audit log for tool failure: %s", tool_name)
