"""
Shared wrapping helpers used by all instrumentors

These sit on top of `wrapt` and enforce agenx conventions:
- Always open a span before calling original
- Always close span on return OR exception
- Capture input/output on the span according to span type
- Propagate exceptions cleanly (never swallow)
- Work for both sync and async methods
"""

from typing import Callable, Any, Optional
import wrapt

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

_patched: set[str] = set()

def _set_input(span: LiveSpan, span_type: SpanType, value: Any) -> None:
    if span_type in (SpanType.ROOT, SpanType.STEP):
        span.step.input = value
    elif span_type == SpanType.LLM:
        span.llm.prompt = str(value) if value is not None else None
    elif span_type == SpanType.AGENT:
        span.agent.input = value
    elif span_type == SpanType.TOOL:
        span.tool.input = value
    elif span_type == SpanType.RETRIEVER:
        span.retriever.query = str(value) if value is not None else None
    elif span_type == SpanType.GUARD:
        span.guard.input = value      

def _set_output(span: LiveSpan, span_type: SpanType, value: Any) -> None:
    if span_type in (SpanType.ROOT, SpanType.STEP):
        span.step.output = value
    elif span_type == SpanType.LLM:
        span.llm.output = str(value) if value is not None else None
    elif span_type == SpanType.AGENT:
        span.agent.output = value
    elif span_type == SpanType.TOOL:
        span.tool.output = value
    elif span_type == SpanType.RETRIEVER:
        span.retriever.documents = value if isinstance(value, list) else [value]
        span.retriever.returned = len(span.retriever.documents) 
    elif span_type == SpanType.GUARD:
        if isinstance(value, tuple) and len(value) == 2:
            span.guard.passed, span.guard.reason = value
        else:
            span.guard.passed = bool(value)   

def wrap_sync(
    module_path: str,
    class_and_method: str,
    span_name: str,
    span_type: SpanType,
    extract_input: Callable[..., Any],
    extract_output: Callable[..., Any],
    enrich_span: Optional[Callable[[LiveSpan, Any, tuple, dict], None]] = None,
) -> None:
    """
    Patch a synchronous method with an agenx span wrapper.

    Args:
        module_path:        dotted module string, e.g. "langchain_core.runnables.base"
        class_and_method:   "ClassName.method_name"
        span_name:          name that appears in traces, e.g. "langchain.chain"
        span_type:          SpanType
        extract_input:      callable that derives the span input from call args
        extract_output:     callable that derives the span output from the return value
        enrich_span:        optional callback to set extra typed attributes
    """
    key = f"{module_path}.{class_and_method}"
    if key in _patched:
        return
    _patched.add(key)

    @wrapt.patch_function_wrapper(module_path, class_and_method)
    def _wrapper(wrapped, instance, args, kwargs):
        with trace.span(span_name, type=span_type) as span:
            try:
                span_input = extract_input(args, kwargs, instance)
                _set_input(span, span_type, span_input)
            except Exception:
                pass 

            if enrich_span:
                try:
                    enrich_span(span, instance, args, kwargs)
                except Exception:
                    pass

            result = wrapped(*args, **kwargs)

            try:
                span_output = extract_output(result, instance)
                _set_output(span, span_type, span_output)
            except Exception:
                pass

            return result
        
def wrap_async(
    module_path: str,
    class_and_method: str,
    span_name: str,
    span_type: SpanType,
    extract_input: Callable[..., Any],
    extract_output: Callable[..., Any],
    enrich_span: Optional[Callable[[LiveSpan, Any, tuple, dict], None]] = None,
) -> None:
    """Same as wrap_sync but for coroutine methods"""

    key = f"{module_path}.{class_and_method}"
    if key in _patched:
        return
    _patched.add(key)

    @wrapt.patch_function_wrapper(module_path, class_and_method)
    async def _wrapper(wrapped, instance, args, kwargs):
        async with trace.aspan(span_name, type=span_type) as span:
            try:
                span_input = extract_input(args, kwargs, instance)
                _set_input(span, span_type, span_input)
            except Exception:
                pass

            if enrich_span:
                try:
                    enrich_span(span, instance, args, kwargs)
                except Exception:
                    pass

            result = await wrapped(*args, **kwargs)

            try:
                span_output = extract_output(result, instance)
                _set_output(span, span_type, span_output)
            except Exception:
                pass

            return result