"""
AgentBase Decorators - @observe and @heal for function-level tracing.
"""

import functools
import time
import logging
import asyncio

from agentbase_sdk.tracer import Tracer, Span

logger = logging.getLogger("agentbase")


def observe(name=None, span_type="custom", capture_input=True, capture_output=True, as_trace=False):
    """Decorator that automatically traces a function."""
    def decorator(func):
        span_name = name or f"{func.__module__}.{func.__qualname__}"

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            ctx = Tracer(span_name) if as_trace else Span(span_name, span_type=span_type)
            with ctx as s:
                if capture_input:
                    s.set_input(_fmt(args, kwargs))
                try:
                    result = func(*args, **kwargs)
                    if capture_output and result is not None:
                        s.set_output(str(result)[:5000])
                    return result
                except Exception as e:
                    s.set_error(type(e).__name__, str(e))
                    raise

        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            ctx = Tracer(span_name) if as_trace else Span(span_name, span_type=span_type)
            with ctx as s:
                if capture_input:
                    s.set_input(_fmt(args, kwargs))
                try:
                    result = await func(*args, **kwargs)
                    if capture_output and result is not None:
                        s.set_output(str(result)[:5000])
                    return result
                except Exception as e:
                    s.set_error(type(e).__name__, str(e))
                    raise

        return async_wrapper if asyncio.iscoroutinefunction(func) else wrapper
    return decorator


def heal(retry=3, retry_delay=1.0, fallback=None, on_error=None, circuit_break_after=10):
    """Decorator that adds auto-heal with retry, backoff, and circuit breaker."""
    state = {"failures": 0, "circuit_open": False, "circuit_opened_at": 0.0}

    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            if state["circuit_open"]:
                if time.time() - state["circuit_opened_at"] < 60.0:
                    if fallback is not None:
                        return fallback() if callable(fallback) else fallback
                    raise RuntimeError(f"Circuit breaker open for {func.__qualname__}")
                state["circuit_open"] = False
                state["failures"] = 0

            last_error = None
            for attempt in range(1, retry + 1):
                with Span(f"heal.{func.__qualname__}.attempt_{attempt}", span_type="custom") as s:
                    try:
                        result = func(*args, **kwargs)
                        state["failures"] = 0
                        return result
                    except Exception as e:
                        last_error = e
                        state["failures"] += 1
                        s.set_error(type(e).__name__, str(e))
                        if on_error:
                            on_error(e, attempt)
                        if state["failures"] >= circuit_break_after:
                            state["circuit_open"] = True
                            state["circuit_opened_at"] = time.time()
                            break
                        if attempt < retry:
                            time.sleep(retry_delay * (2 ** (attempt - 1)))

            if fallback is not None:
                return fallback() if callable(fallback) else fallback
            raise last_error

        return wrapper
    return decorator


def _fmt(args, kwargs):
    parts = [str(a)[:500] for a in args]
    parts += [f"{k}={str(v)[:500]}" for k, v in kwargs.items()]
    return ", ".join(parts)[:5000]
