import asyncio
import functools
import inspect
import json
import time
import traceback
from collections import OrderedDict
from typing import Any, Callable, Optional, TypeVar, cast

try:
    from pipecat.services.llm_service import FunctionCallParams
except ImportError:
    FunctionCallParams = None

F = TypeVar('F', bound=Callable[..., Any])


class ToolExecutionContext:
    def __init__(self):
        self.tool_name: Optional[str] = None
        self.tool_call_id: Optional[str] = None
        self.arguments: Optional[dict] = None
        self.start_time: Optional[float] = None
        self.end_time: Optional[float] = None
        self.duration_ms: Optional[float] = None
        self.result: Any = None
        self.error: Optional[Exception] = None
        self.error_traceback: Optional[str] = None
        self.success: bool = True
        self.timeout_enabled: bool = False
        self.timeout_seconds: Optional[float] = None
        self.timeout_triggered: bool = False
        self.callback_called: bool = False

    def to_dict(self) -> dict:
        return {
            "tool_name": self.tool_name,
            "tool_call_id": self.tool_call_id,
            "arguments": self.arguments,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration_ms": self.duration_ms,
            "success": self.success,
            "error": str(self.error) if self.error else None,
            "error_type": type(self.error).__name__ if self.error else None,
            "error_traceback": self.error_traceback,
        }


class ToolObserverRegistry:
    _instance: Optional['ToolObserverRegistry'] = None
    
    def __init__(self, max_age_seconds: float = 300.0):
        self._executions: OrderedDict[str, tuple[ToolExecutionContext, float]] = OrderedDict()
        self._lock = asyncio.Lock()
        self._max_age = max_age_seconds
    
    @classmethod
    def get_instance(cls) -> 'ToolObserverRegistry':
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance
    
    async def register_execution(self, context: ToolExecutionContext) -> None:
        if not context.tool_call_id:
            return
        async with self._lock:
            self._executions[context.tool_call_id] = (context, time.time())
            if len(self._executions) > 100:
                await self._cleanup_old_executions()
    
    async def get_execution(self, tool_call_id: str) -> Optional[ToolExecutionContext]:
        async with self._lock:
            entry = self._executions.get(tool_call_id)
            if entry:
                return entry[0]
            return None
    
    async def clear_execution(self, tool_call_id: str) -> None:
        async with self._lock:
            self._executions.pop(tool_call_id, None)
    
    async def _cleanup_old_executions(self) -> None:
        now = time.time()
        expired = [
            tool_call_id for tool_call_id, (_, created_at) in self._executions.items()
            if (now - created_at) > self._max_age
        ]
        for tool_call_id in expired:
            self._executions.pop(tool_call_id, None)
    
    def register_execution_sync(self, context: ToolExecutionContext) -> None:
        if not context.tool_call_id:
            return
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                asyncio.create_task(self.register_execution(context))
            else:
                loop.run_until_complete(self.register_execution(context))
        except Exception:
            pass


def _extract_tool_name(tool_name: Optional[str], args: tuple, func: Callable) -> str:
    if tool_name:
        return tool_name
    if args and hasattr(args[0], 'schema'):
        return getattr(args[0].schema, 'name', func.__name__)
    return func.__name__


def _extract_params(args: tuple, kwargs: dict) -> Optional[FunctionCallParams]:
    if FunctionCallParams is None:
        return None
    for arg in args:
        if isinstance(arg, FunctionCallParams):
            return arg
    for value in kwargs.values():
        if isinstance(value, FunctionCallParams):
            return value
    return None


def _detect_error_in_result(result: Any) -> bool:
    if result is None:
        return False
    
    try:
        if isinstance(result, dict):
            if any(key in result for key in ['error', 'exception', 'failure', 'failed']):
                return True
            if result.get('success') is False:
                return True
            result_str = json.dumps(result).lower()
        elif isinstance(result, str):
            result_str = result.lower()
        else:
            result_str = str(result).lower()
        
        error_patterns = ['error', 'exception', 'failed', 'failure']
        return any(pattern in result_str for pattern in error_patterns)
    except Exception:
        return False


def _wrap_callback(
    original_callback: Callable,
    context: ToolExecutionContext,
    include_result: bool
) -> Callable:
    async def wrapped_callback(result):
        context.callback_called = True
        if include_result and context.result is None:
            context.result = result
            if result is None:
                context.success = False
                context.error = Exception("Tool returned None (null)")
            elif _detect_error_in_result(result):
                context.success = False
                context.error = Exception(f"Error detected in result: {result}")
        return await original_callback(result)
    return wrapped_callback


async def _execute_with_timeout(
    func: Callable,
    args: tuple,
    kwargs: dict,
    context: ToolExecutionContext,
    original_callback: Optional[Callable],
) -> Any:
    """Execute function with timeout using asyncio.wait_for for proper cancellation."""
    try:
        result = await asyncio.wait_for(
            func(*args, **kwargs),
            timeout=context.timeout_seconds
        )
        return result
    except asyncio.CancelledError:
        if original_callback and not context.callback_called:
            context.success = False
            context.error = Exception("Tool execution was cancelled")
            cancel_payload = {
                "success": False,
                "error": "Tool execution was cancelled"
            }
            try:
                await original_callback(cancel_payload)
                context.callback_called = True
            except Exception:
                pass
        raise
    except asyncio.TimeoutError:
        context.timeout_triggered = True
        context.success = False
        context.error = TimeoutError(f"Tool execution timeout after {context.timeout_seconds}s")
        
        if original_callback and not context.callback_called:
            timeout_payload = {
                "success": False,
                "error": f"Tool execution timeout after {context.timeout_seconds} seconds"
            }
            try:
                await original_callback(timeout_payload)
                context.callback_called = True
            except Exception:
                pass
        
        raise


def _process_result(context: ToolExecutionContext, result: Any, include_result: bool) -> None:
    if include_result and context.result is None:
        context.result = result
        if result is None:
            context.success = False
            context.error = Exception("Tool returned None (null)")
        elif _detect_error_in_result(result):
            context.success = False
            context.error = Exception(f"Error detected in result: {result}")
    
    if context.success and context.error is None:
        context.success = True


async def _handle_missing_callback(
    context: ToolExecutionContext,
    original_callback: Optional[Callable]
) -> None:
    if not (original_callback and not context.callback_called):
        return
    
    try:
        if context.result is not None:
            await original_callback(context.result)
        else:
            error_payload = {
                "success": False,
                "error": "Tool returned None (null) and did not call result_callback"
            }
            await original_callback(error_payload)
        context.callback_called = True
    except Exception:
        pass
    
    if context.success:
        context.success = False
        if context.result is None:
            context.error = Exception("Tool returned None (null) and did not call result_callback")
        else:
            context.error = Exception("Tool has result_callback available but did not call it")


def _handle_missing_callback_sync(
    context: ToolExecutionContext,
    original_callback: Optional[Callable]
) -> None:
    if not (original_callback and callable(original_callback) and not context.callback_called):
        return
    
    async def send_result_via_callback():
        try:
            if context.result is not None:
                await original_callback(context.result)
            else:
                error_payload = {
                    "success": False,
                    "error": "Tool returned None (null) and did not call result_callback"
                }
                await original_callback(error_payload)
            context.callback_called = True
        except Exception:
            pass
    
    try:
        loop = asyncio.get_event_loop()
        if loop.is_running():
            asyncio.create_task(send_result_via_callback())
        else:
            loop.run_until_complete(send_result_via_callback())
    except Exception:
        pass
    
    if context.success:
        context.success = False
        if context.result is None:
            context.error = Exception("Tool returned None (null) and did not call result_callback")
        else:
            context.error = Exception("Tool has result_callback available but did not call it")


def observe_tool(
    tool_name: Optional[str] = None,
    include_result: bool = True,
    include_traceback: bool = True,
    enable_timeout: bool = True,
    timeout_seconds: float = 10.0,
) -> Callable[[F], F]:
    if FunctionCallParams is None:
        def passthrough_decorator(func: F) -> F:
            return func
        return passthrough_decorator
    
    def decorator(func: F) -> F:
        registry = ToolObserverRegistry.get_instance()
        
        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs) -> Any:
            context = ToolExecutionContext()
            context.start_time = time.time()
            context.tool_name = _extract_tool_name(tool_name, args, func)
            
            params = _extract_params(args, kwargs)
            original_callback = None
            
            if enable_timeout:
                context.timeout_enabled = True
                context.timeout_seconds = timeout_seconds
            
            if params:
                context.tool_call_id = getattr(params, 'tool_call_id', None)
                if hasattr(params, 'arguments') and isinstance(params.arguments, dict):
                    context.arguments = params.arguments.copy()
                
                if context.tool_call_id:
                    await registry.register_execution(context)
                
                original_callback = getattr(params, 'result_callback', None)
                if original_callback and callable(original_callback):
                    params.result_callback = _wrap_callback(original_callback, context, include_result)
            
            try:
                if context.timeout_enabled and original_callback:
                    result = await _execute_with_timeout(func, args, kwargs, context, original_callback)
                else:
                    result = await func(*args, **kwargs)
                _process_result(context, result, include_result)
                return result
            except asyncio.CancelledError:
                context.success = False
                if not context.error:
                    context.error = Exception("Tool execution was cancelled")
                
                if params and original_callback and not context.callback_called:
                    cancel_payload = {
                        "success": False,
                        "error": "Tool execution was cancelled"
                    }
                    try:
                        await original_callback(cancel_payload)
                        context.callback_called = True
                    except Exception:
                        pass
                
                raise
            except asyncio.TimeoutError:
                raise
            except Exception as e:
                context.success = False
                context.error = e
                
                if include_traceback:
                    context.error_traceback = traceback.format_exc()
                
                if params and original_callback and not context.callback_called:
                    error_payload = {
                        "success": False,
                        "error": str(e),
                        "error_type": type(e).__name__,
                    }
                    try:
                        await original_callback(error_payload)
                        context.callback_called = True
                    except Exception:
                        pass
                
                raise
            finally:
                context.end_time = time.time()
                context.duration_ms = (context.end_time - context.start_time) * 1000
                
                await _handle_missing_callback(context, original_callback)
        
        @functools.wraps(func)
        def sync_wrapper(*args, **kwargs) -> Any:
            context = ToolExecutionContext()
            context.start_time = time.time()
            context.tool_name = _extract_tool_name(tool_name, args, func)
            
            params = _extract_params(args, kwargs)
            
            if enable_timeout:
                context.timeout_enabled = True
                context.timeout_seconds = timeout_seconds
            
            if params:
                context.tool_call_id = getattr(params, 'tool_call_id', None)
                if hasattr(params, 'arguments') and isinstance(params.arguments, dict):
                    context.arguments = params.arguments.copy()
                
                if context.tool_call_id:
                    registry.register_execution_sync(context)
            
            original_callback = None
            if params:
                original_callback = getattr(params, 'result_callback', None)
            
            timeout_task = None
            if context.timeout_enabled and params and original_callback and callable(original_callback):
                async def timeout_handler():
                    await asyncio.sleep(context.timeout_seconds)
                    if context.end_time is None:
                        context.timeout_triggered = True
                        context.success = False
                        context.error = TimeoutError(f"Tool execution timeout after {context.timeout_seconds}s")
                        timeout_payload = {
                            "success": False,
                            "error": f"Tool execution timeout after {context.timeout_seconds} seconds"
                        }
                        await original_callback(timeout_payload)
                        context.callback_called = True
                
                try:
                    loop = asyncio.get_event_loop()
                    if loop.is_running():
                        timeout_task = asyncio.create_task(timeout_handler())
                    else:
                        timeout_task = loop.create_task(timeout_handler())
                except RuntimeError:
                    pass
            
            try:
                result = func(*args, **kwargs)
                _process_result(context, result, include_result)
                return result
            except Exception as e:
                context.success = False
                context.error = e
                
                if include_traceback:
                    context.error_traceback = traceback.format_exc()
                
                raise
            finally:
                context.end_time = time.time()
                context.duration_ms = (context.end_time - context.start_time) * 1000
                
                _handle_missing_callback_sync(context, original_callback)
                
                if timeout_task and not timeout_task.done():
                    try:
                        timeout_task.cancel()
                    except Exception:
                        pass
        
        if inspect.iscoroutinefunction(func):
            return cast(F, async_wrapper)
        else:
            return cast(F, sync_wrapper)
    
    return decorator
