"""
Safe tool wrapper for MCP tools.

Provides a decorator that wraps MCP tool functions with:
- Consistent error handling (no raw tracebacks to LLM)
- Timing instrumentation (_timing_ms in every response)
- Structured error format: {"status":"error","tool":"name","error":"message","_timing_ms":N}
"""

import functools
import json
import logging
import time
import traceback

logger = logging.getLogger("databridge.tools")


def safe_tool(tool_name: str):
    """Decorator that wraps an MCP tool with error handling and timing.

    Usage::

        @mcp.tool()
        @safe_tool("get_smart_recommendations")
        def get_smart_recommendations(...) -> str:
            ...

    Args:
        tool_name: Name used in error payloads and log messages.

    Behaviour:
        - On success: if the return value is a JSON string, injects ``_timing_ms``;
          if it's a dict, adds ``_timing_ms`` and re-serialises.
        - On error: logs the full traceback and returns a sanitised JSON error.
    """

    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            start = time.perf_counter()
            try:
                result = func(*args, **kwargs)
                elapsed_ms = round((time.perf_counter() - start) * 1000, 1)

                # Inject timing into the response
                if isinstance(result, dict):
                    result["_timing_ms"] = elapsed_ms
                    return result
                elif isinstance(result, str):
                    try:
                        parsed = json.loads(result)
                        if isinstance(parsed, dict):
                            parsed["_timing_ms"] = elapsed_ms
                            return json.dumps(parsed, indent=2, default=str)
                    except (json.JSONDecodeError, TypeError):
                        pass
                    # Non-JSON string — return as-is with timing note
                    return result

                # Other return types — pass through
                return result

            except Exception as exc:
                elapsed_ms = round((time.perf_counter() - start) * 1000, 1)
                logger.error(
                    "Tool %s failed: %s\n%s",
                    tool_name,
                    exc,
                    traceback.format_exc(),
                )
                error_payload = {
                    "status": "error",
                    "tool": tool_name,
                    "error": str(exc),
                    "_timing_ms": elapsed_ms,
                }
                # Return dict for tools that return Dict, string for str-returning tools
                if func.__annotations__.get("return") in (dict, "Dict[str, Any]"):
                    return error_payload
                return json.dumps(error_payload, indent=2)

        return wrapper

    return decorator
