"""
Telemetry decorator for Unity MCP tools
"""

import functools
import time
import inspect
import logging
from typing import Callable, Any
from telemetry import record_tool_usage, record_milestone, MilestoneType

_log = logging.getLogger("unity-mcp-telemetry")
_decorator_log_count = 0

def telemetry_tool(tool_name: str):
    """Decorator to add telemetry tracking to MCP tools"""
    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def _sync_wrapper(*args, **kwargs) -> Any:
            start_time = time.time()
            success = False
            error = None
            # Extract sub-action (e.g., 'get_hierarchy') from bound args when available
            sub_action = None
            try:
                sig = inspect.signature(func)
                bound = sig.bind_partial(*args, **kwargs)
                bound.apply_defaults()
                sub_action = bound.arguments.get("action")
            except Exception:
                sub_action = None
            try:
                global _decorator_log_count
                if _decorator_log_count < 10:
                    _log.info(f"telemetry_decorator sync: tool={tool_name}")
                    _decorator_log_count += 1
                result = func(*args, **kwargs)
                success = True
                if tool_name == "manage_script" and kwargs.get("action") == "create":
                    record_milestone(MilestoneType.FIRST_SCRIPT_CREATION)
                elif tool_name.startswith("manage_scene"):
                    record_milestone(MilestoneType.FIRST_SCENE_MODIFICATION)
                record_milestone(MilestoneType.FIRST_TOOL_USAGE)
                return result
            except Exception as e:
                error = str(e)
                raise
            finally:
                duration_ms = (time.time() - start_time) * 1000
                record_tool_usage(tool_name, success, duration_ms, error, sub_action=sub_action)

        @functools.wraps(func)
        async def _async_wrapper(*args, **kwargs) -> Any:
            start_time = time.time()
            success = False
            error = None
            # Extract sub-action (e.g., 'get_hierarchy') from bound args when available
            sub_action = None
            try:
                sig = inspect.signature(func)
                bound = sig.bind_partial(*args, **kwargs)
                bound.apply_defaults()
                sub_action = bound.arguments.get("action")
            except Exception:
                sub_action = None
            try:
                global _decorator_log_count
                if _decorator_log_count < 10:
                    _log.info(f"telemetry_decorator async: tool={tool_name}")
                    _decorator_log_count += 1
                result = await func(*args, **kwargs)
                success = True
                if tool_name == "manage_script" and kwargs.get("action") == "create":
                    record_milestone(MilestoneType.FIRST_SCRIPT_CREATION)
                elif tool_name.startswith("manage_scene"):
                    record_milestone(MilestoneType.FIRST_SCENE_MODIFICATION)
                record_milestone(MilestoneType.FIRST_TOOL_USAGE)
                return result
            except Exception as e:
                error = str(e)
                raise
            finally:
                duration_ms = (time.time() - start_time) * 1000
                record_tool_usage(tool_name, success, duration_ms, error, sub_action=sub_action)

        return _async_wrapper if inspect.iscoroutinefunction(func) else _sync_wrapper
    return decorator