"""
Universal decorator factory for all span types
"""

from __future__ import annotations

import functools
import inspect
from typing import Any, Callable, TypeVar, ParamSpec, TYPE_CHECKING

from agenx.trace.spans.models import SpanType
from agenx.trace.spans.span import LiveSpan

if TYPE_CHECKING:
    from agenx.tracer import Tracer

P = ParamSpec('P')
R = TypeVar('R')

_SENSITIVE_PATTERNS = frozenset([
    'password', 'passwd', 'pwd',
    'secret', 'secret_key', 'api_key', 'apikey', 'api-key',
    'token', 'access_token', 'refresh_token', 'auth_token',
    'credential', 'auth', 'authorization', 'bearer',
    'private_key', 'privatekey', 'private-key',
])

def _get_span_name(func: Callable, name: str | None) -> str:
    """Determine span name: explicit name or function name"""
    return name if name else func.__name__


def _truncate(value: str, max_len: int = 1000) -> str:
    """Truncate string to max_len"""
    if len(value) <= max_len:
        return value
    return value[:max_len - 3] + '...'


def _is_sensitive(name: str) -> bool:
    """Check if parameter name suggests sensitive data."""
    name_lower = name.lower().replace('-', '_')
    return any(pattern in name_lower for pattern in _SENSITIVE_PATTERNS)


def _safe_repr(obj: Any, depth: int = 3, seen: set | None = None) -> Any:
    """Safely serialize object to JSON-serializable representation"""

    if seen is None:
        seen = set()
    
    obj_id = id(obj)
    if obj_id in seen:
        return "<circular>"
    seen.add(obj_id)
    
    try:
        if obj is None:
            return None
        
        if isinstance(obj, (str, int, float, bool)):
            s = str(obj)
            return _truncate(s)
        
        if isinstance(obj, bytes):
            return f"<bytes:len={len(obj)}>"
        
        if isinstance(obj, (list, tuple)):
            if depth <= 0:
                return f"<{type(obj).__name__}:len={len(obj)}>"
            return [_safe_repr(item, depth - 1, seen) for item in obj[:10]]
        
        if isinstance(obj, dict):
            if depth <= 0:
                return f"<dict:len={len(obj)}>"
            return {
                k: "[REDACTED]" if _is_sensitive(str(k)) else _safe_repr(v, depth - 1, seen)
                for k, v in list(obj.items())[:10]
            }
        
        return _truncate(repr(obj))
    
    except Exception:
        return f"<{type(obj).__name__}>"


def _extract_token_usage(result: Any) -> dict | None:
    """Try to extract token usage from various LLM response formats."""
    # OpenAI format
    if hasattr(result, 'usage'):
        return {
            'input': getattr(result.usage, 'prompt_tokens', 0),
            'output': getattr(result.usage, 'completion_tokens', 0),
            'total': getattr(result.usage, 'total_tokens', 0),
        }
    
    # Anthropic format
    if hasattr(result, 'usage'):
        return {
            'input': getattr(result.usage, 'input_tokens', 0),
            'output': getattr(result.usage, 'output_tokens', 0),
            'total': getattr(result.usage, 'input_tokens', 0) + getattr(result.usage, 'output_tokens', 0),
        }
    
    if isinstance(result, dict):
        if 'usage' in result:
            u = result['usage']
            return {
                'input': u.get('prompt_tokens', u.get('input_tokens', 0)),
                'output': u.get('completion_tokens', u.get('output_tokens', 0)),
                'total': u.get('total_tokens', 0),
            }
    
    return None


def _count_iterations(result: Any) -> int:
    """Count iterations if result suggests an iterative process."""
    if hasattr(result, 'iterations'):
        return int(result.iterations)
    
    if isinstance(result, dict) and 'iterations' in result:
        return int(result['iterations'])
    
    if isinstance(result, tuple) and len(result) == 2:
        try:
            return int(result[1])
        except (TypeError, ValueError):
            pass
    
    return 1 


def _capture_exception(span: LiveSpan, exc: BaseException) -> None:
    """Capture exception details in metadata."""
    try:
        if hasattr(span._span.attributes, 'metadata'):
            span._span.attributes.metadata['exception_type'] = type(exc).__name__
            span._span.attributes.metadata['exception_message'] = _truncate(str(exc))
    except Exception:
        pass


def _capture_input(
    span: LiveSpan,
    span_type: SpanType,
    args: tuple,
    kwargs: dict,
) -> None:
    """Capture function arguments as span attributes"""
    try:
        if span_type in (SpanType.STEP, SpanType.ROOT):
            input_data = {
                "args": _safe_repr(args),
                "kwargs": _safe_repr(kwargs),
            }
            span.step.input = input_data
        
        elif span_type == SpanType.LLM:
            prompt = args[0] if args else kwargs.get('prompt', kwargs.get('messages', ''))
            if isinstance(prompt, list):
                span.llm.prompt = _truncate(str(prompt[:5]))        # If list, storing only 5 images for v1 release
            else:
                span.llm.prompt = _truncate(str(prompt))
        
        elif span_type == SpanType.AGENT:
            input_data = args[0] if args else kwargs
            span.agent.input = _safe_repr(input_data)
        
        elif span_type == SpanType.TOOL:
            input_data = args[0] if args else kwargs
            span.tool.input = _safe_repr(input_data)
        
        elif span_type == SpanType.RETRIEVER:
            query = args[0] if args else kwargs.get('query', '')
            span.retriever.query = _truncate(str(query))
        
        elif span_type == SpanType.GUARD:
            input_data = args[0] if args else kwargs
            span.guard.input = _safe_repr(input_data)
    
    except Exception:
        pass  


def _capture_output(
    span: LiveSpan,
    span_type: SpanType,
    result: Any
) -> None:
    """Capture function return value as span attributes"""
    
    try:
        if span_type in (SpanType.STEP, SpanType.ROOT):
            span.step.output = _safe_repr(result)

        elif span_type == SpanType.LLM:
            if hasattr(result, 'content'):
                span.llm.output = _truncate(str(result.content))
            elif isinstance(result, str):
                span.llm.output = _truncate(result)
            elif isinstance(result, dict) and 'content' in result:
                span.llm.output = _truncate(str(result['content']))
            else:
                span.llm.output = _truncate(str(result)[:200])

            tokens = _extract_token_usage(result)
            if tokens:
                span.llm.tokens_input = tokens['input']
                span.llm.tokens_output = tokens['output']
                span.llm.tokens_total = tokens['total']
        
        elif span_type == SpanType.AGENT:
            span.agent.output = _safe_repr(result)
        
        elif span_type == SpanType.TOOL:
            span.tool.output = _safe_repr(result)
        
        elif span_type == SpanType.RETRIEVER:
            if isinstance(result, (list, tuple)):
                span.retriever.documents = _safe_repr(result[:10])   # Storing only 10 for v1 release
                span.retriever.returned = len(result)
            else:
                span.retriever.ouput = _safe_repr(result)

        elif span_type == SpanType.GUARD:
            span.guard.passed = True

            if isinstance(result, tuple) and len(result) == 2:
                span.guard.passed = bool(result[0])
                span.guard.reason = str(result[1]) if result[1] else ""

    except Exception:
        pass


def _capture_guard_result(span: LiveSpan, result: Any) -> None:
    """Special handling for guard decorators that return (passed, reason)."""
    try:
        if isinstance(result, tuple) and len(result) >= 2:
            span.guard.passed = bool(result[0])
            if result[1]:
                span.guard.reason = _truncate(str(result[1]))
    except Exception:
        pass


def _create_sync_wrapper(
    func: Callable[P, R],
    tracer: Tracer,
    span_type: SpanType,
    name: str,
    capture_args: bool,
    capture_result: bool,
    **decorator_options: Any,
 )-> Callable[P,R]:
    """Create a synchronous wrapper for the decorated function"""

    @functools.wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        span_name = _get_span_name(func, name)

        with tracer.span(span_name, type=span_type) as span:
            for key, value in decorator_options.items():
                if hasattr(span._span.attributes, key):
                    setattr(span._span.attributes, key, value)
                elif hasattr(span._span.attributes, 'metadata'):
                    span._span.attributes.metadata[key] = value

            if capture_args:
                _capture_input(span, span_type, args, kwargs)

            if span_type == SpanType.LLM and 'model' in decorator_options:
                span.llm.model = decorator_options.get('model')
                if 'provider' in decorator_options:
                    span.llm.provider = decorator_options.get('provider')
                if 'temperature' in decorator_options:
                    span.llm.temperature = decorator_options.get('temperature')

            elif span_type == SpanType.AGENT and 'tools' in decorator_options:
                span.agent.tools = list(decorator_options.get('tools'))

            elif span_type == SpanType.TOOL and 'tool_name' in decorator_options:
                span.tool.tool_name = decorator_options.get('tool_name')

            elif span_type == SpanType.RETRIEVER and 'top_k' in decorator_options:
                span.retriever.top_k = decorator_options.get('top_k')

            elif span_type == SpanType.GUARD and 'guard_name' in decorator_options:
                span.guard.guard_name = decorator_options.get('guard_name')

            try:
                result = func(*args, **kwargs)

                if capture_result:
                    _capture_output(span, span_type, result)

                if span_type == SpanType.AGENT:
                    span.agent.iterations = _count_iterations(result)

                if span_type == SpanType.GUARD:
                    _capture_guard_result(span, result)

                return result

            except Exception as e:
                _capture_exception(span, e)
                raise

    return wrapper           


def _create_async_wrapper(
    func: Callable[P, R],
    tracer: Tracer,
    span_type: SpanType,
    name: str | None,
    capture_args: bool,
    capture_result: bool,
    **decorator_options: Any,
) -> Callable[P, R]:
    """Create an async wrapper for the decorated function."""
    
    @functools.wraps(func)
    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        span_name = _get_span_name(func, name)

        async with tracer.aspan(span_name, type=span_type) as span:
            for key, value in decorator_options.items():
                if hasattr(span._span.attributes, key):
                    setattr(span._span.attributes, key, value)
                elif hasattr(span._span.attributes, 'metadata'):
                    span._span.attributes.metadata[key] = value

            if capture_args:
                _capture_input(span, span_type, args, kwargs)

            if span_type == SpanType.LLM and 'model' in decorator_options:
                span.llm.model = decorator_options['model']
                if 'provider' in decorator_options:
                    span.llm.provider = decorator_options['provider']
                if 'temperature' in decorator_options:
                    span.llm.temperature = decorator_options['temperature']

            elif span_type == SpanType.AGENT and 'tools' in decorator_options:
                span.agent.tools = list(decorator_options['tools'])

            elif span_type == SpanType.TOOL and 'tool_name' in decorator_options:
                span.tool.tool_name = decorator_options['tool_name']

            elif span_type == SpanType.RETRIEVER and 'top_k' in decorator_options:
                span.retriever.top_k = decorator_options['top_k']

            elif span_type == SpanType.GUARD and 'guard_name' in decorator_options:
                span.guard.guard_name = decorator_options['guard_name']

            try:
                result = await func(*args, **kwargs)

                if capture_result:
                    _capture_output(span, span_type, result)

                if span_type == SpanType.AGENT:
                    span.agent.iterations = _count_iterations(result)

                if span_type == SpanType.GUARD:
                    _capture_guard_result(span, result)

                return result

            except BaseException as e:
                _capture_exception(span, e)
                raise

    return wrapper


def create_decorator(
    func: Callable[P, R],
    name: str | None = None,
    span_type: SpanType = SpanType.STEP,
    tracer: Tracer | None = None,
    capture_args: bool = True,
    capture_result: bool = True,
    **options: Any
) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
    """
    Universal decorator factory for all span types.
    
    Handles three call patterns:
    1. @trace.step                        -> func is not None
    2. @trace.step("name")                -> func is None, name is str
    3. @trace.step(capture_args=False)    -> func is None, options provided
    
    Returns:
        If func is None: A decorator function
        If func is provided: A wrapped function
    """

    if tracer is None:
        from agenx import tracer as default_tracer
        tracer = default_tracer

    """Case 1: Direct decoration (@decorator)"""
    if func is not None:
        if inspect.iscoroutinefunction(func):
            return _create_async_wrapper(
                func, tracer, span_type, name, capture_args, capture_result, **options
            )
        else:
            return _create_sync_wrapper(
                func, tracer, span_type, name, capture_args, capture_result, **options
            )
        
    """Case 2 & 3: Called with arguments (@decorator() or @decorator("name"))"""
    def decorator(actual_func: Callable[P, R]) -> Callable[P, R]:
        if inspect.iscoroutinefunction(actual_func):
            return _create_async_wrapper(
                actual_func, tracer, span_type, name, capture_args, capture_result, **options
            )
        else:
            return _create_sync_wrapper(
                actual_func, tracer, span_type, name, capture_args, capture_result, **options
            )
    
    return decorator
    
