from __future__ import annotations
from functools import wraps
from typing import Callable, TypeVar, Any
from collections.abc import Mapping
import inspect
import json
import types

from opentelemetry import trace as otel_trace
from opentelemetry.trace import Status, StatusCode
from .context import push_entity, pop_entity, current_entity_path

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

class ImpactAttributes:
    """Impact-specific attributes created via @trace."""

    TRACE_TYPE = "impact.trace.type"
    TRACE_NAME = "impact.trace.name"
    TRACE_PATH = "impact.trace.path"
    TRACE_INPUT = "impact.trace.input"
    TRACE_OUTPUT = "impact.trace.output"
    TRACE_PREFIX = "impact.trace."

IA = ImpactAttributes

def _start_span(span_name: str, span_type: str, attributes: Mapping[str, Any] | None = None):
    """Start a span as the current span."""
    tracer = otel_trace.get_tracer("impact")
    attrs: dict[str, Any] = dict(attributes or {})
    attrs[IA.TRACE_TYPE] = span_type
    return tracer.start_as_current_span(span_name, attributes=attrs)

def _capture_io(span, args, kwargs, result_marker=None):
    try:
        span.set_attribute(IA.TRACE_INPUT, json.dumps({"args": args, "kwargs": kwargs}, default=str))
        if result_marker is not None:
            span.set_attribute(IA.TRACE_OUTPUT, json.dumps(result_marker, default=str))
    except Exception:
        # best-effort only
        pass

def _wrap_body(span_name: str, span_type: str, func: F, custom_attributes: Mapping[str, Any] | None = None):
    is_async = inspect.iscoroutinefunction(func)
    is_async_gen = inspect.isasyncgenfunction(func)
    is_gen = inspect.isgeneratorfunction(func)

    # Prefix custom attributes under impact.trace.*.
    attrs: dict[str, Any] = {}
    if custom_attributes:
        for k, v in custom_attributes.items():
            # Prevent overriding reserved decorator keys.
            if k in ("type", "name", "path", "input", "output"):
                continue
            attrs[f"{IA.TRACE_PREFIX}{k}"] = v

    # Enforce the name attribute (useful for UIs even if span.name changes).
    attrs[IA.TRACE_NAME] = span_name

    if is_async_gen:
        @wraps(func)
        async def async_gen_wrapper(*args, **kwargs):
            with _start_span(span_name, span_type, attributes=attrs) as span:
                token = push_entity(span_name)
                # Set the path attribute after pushing entity so it includes current span
                path = current_entity_path()
                if path is not None:
                    span.set_attribute(IA.TRACE_PATH, path)

                _capture_io(span, args, kwargs)
                try:
                    async for item in func(*args, **kwargs):
                        yield item
                    _capture_io(span, args, kwargs, result_marker={"generator": True})
                except Exception as e:
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    raise
                finally:
                    pop_entity(token)
        return async_gen_wrapper  # type: ignore[misc]

    if is_async:
        @wraps(func)
        async def async_wrapper(*args, **kwargs):
            with _start_span(span_name, span_type, attributes=attrs) as span:
                token = push_entity(span_name)
                # Set the path attribute after pushing entity so it includes current span
                path = current_entity_path()
                if path is not None:
                    span.set_attribute(IA.TRACE_PATH, path)

                _capture_io(span, args, kwargs)
                try:
                    result = await func(*args, **kwargs)
                    _capture_io(span, args, kwargs, result_marker=result)
                    return result
                except Exception as e:
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    raise
                finally:
                    pop_entity(token)
        return async_wrapper  # type: ignore[misc]

    if is_gen:
        @wraps(func)
        def gen_wrapper(*args, **kwargs):
            with _start_span(span_name, span_type, attributes=attrs) as span:
                token = push_entity(span_name)
                # Set the path attribute after pushing entity so it includes current span
                path = current_entity_path()
                if path is not None:
                    span.set_attribute(IA.TRACE_PATH, path)

                _capture_io(span, args, kwargs)
                try:
                    gen = func(*args, **kwargs)
                    assert isinstance(gen, types.GeneratorType)
                    for item in gen:
                        yield item
                    _capture_io(span, args, kwargs, result_marker={"generator": True})
                except Exception as e:
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    raise
                finally:
                    pop_entity(token)
        return gen_wrapper  # type: ignore[misc]

    @wraps(func)
    def sync_wrapper(*args, **kwargs):
        with _start_span(span_name, span_type, attributes=attrs) as span:
            token = push_entity(span_name)
            # Set the path attribute after pushing entity so it includes current span
            path = current_entity_path()
            if path is not None:
                span.set_attribute(IA.TRACE_PATH, path)

            _capture_io(span, args, kwargs)
            try:
                result = func(*args, **kwargs)
                _capture_io(span, args, kwargs, result_marker=result)
                return result
            except Exception as e:
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.record_exception(e)
                raise
            finally:
                pop_entity(token)
    return sync_wrapper  # type: ignore[misc]


def trace(func: F | None = None, type: str = "span", *, name: str | None = None, attributes: Mapping[str, Any] | None = None):
    """Trace a function with an OpenTelemetry span."""
    def decorator(fn: F) -> F:
        span_name = name or fn.__qualname__
        return _wrap_body(span_name, type, fn, attributes)
    if func is not None:
        return decorator(func)
    return decorator
