import inspect
import logging
import socket
from contextlib import contextmanager
from functools import wraps
from typing import Any, Callable, Dict, Iterator, Optional

from openinference.semconv.trace import (
    OpenInferenceMimeTypeValues,
    OpenInferenceSpanKindValues,
    SpanAttributes,
)
from opentelemetry import context
from opentelemetry import trace as trace_api
from opentelemetry.trace import INVALID_SPAN, StatusCode, Tracer
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from opentelemetry.util.types import Attributes
from starlette.background import BackgroundTasks
from starlette.responses import StreamingResponse

from ._with_span import _WithSpan
from .helpers import _Encoder, safe_json_dumps

YTRACING_SPAN_DECORATED_ATTR = "__ytracing_span_decorated__"
ML_TRACEPARENT_HEADER_NAME = "x-ml-traceparent"

logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())


def process_input(bound_args: inspect.BoundArguments):
    try:
        input_value = safe_json_dumps(bound_args.arguments, cls=_Encoder)
        return input_value
    except BaseException as e:
        logger.exception(str(e))


def get_input_attribute(func: Any, args: Any, kwargs: Any, get_input: Optional[Callable[[Any], str]] = None) -> str:
    bound_args = inspect.signature(func).bind(*args, **kwargs)
    return get_input(*args, **kwargs) if get_input else process_input(bound_args)


def process_output(response: Any):
    try:
        output_value = safe_json_dumps(response, cls=_Encoder)
        return output_value
    except BaseException as e:
        logger.exception(str(e))


def get_output_attribute(response: Any, get_output: Optional[Callable[[Any], str]]) -> str:
    return get_output(response) if get_output else process_output(response)


async def get_output_attribute_async(response: Any, get_output: Optional[Callable[[Any], str]]) -> str:
    return get_output(response) if get_output else process_output(response)


@contextmanager
def _start_as_current_span(
    tracer: Tracer,
    span_name: str,
    attributes: Dict[str, Any],
    traceparent: Optional[str],
) -> Iterator[_WithSpan]:
    try:
        carrier = {}
        ctx = None

        logger.debug(f"Starting new span '{span_name}'. traceparent: {traceparent}")
        if traceparent:
            carrier["traceparent"] = traceparent
            ctx = TraceContextTextMapPropagator().extract(carrier=carrier)
            attributes["request.traceparent"] = traceparent
        else:
            # Use current context if no explicit traceparent is provided
            # This ensures child spans are properly linked to parent spans
            # This is especially important when LiveKit's session.start() creates
            # a long-lived "agent_session" span - our @traceable() functions should
            # create child spans under it
            ctx = context.get_current()
            # Check if there's an active span in the current context
            current_span = trace_api.get_current_span()
            if current_span and current_span != trace_api.INVALID_SPAN:
                span_context = current_span.get_span_context()
                if span_context.is_valid:
                    logger.debug(
                        f"Creating child span '{span_name}' under parent span (trace_id: {span_context.trace_id})"
                    )
                else:
                    logger.debug(f"Creating root span '{span_name}' (no valid parent context)")

        span = tracer.start_span(name=span_name, attributes=dict(attributes), context=ctx)
    except Exception:
        logger.exception("Failed to start span")
        span = INVALID_SPAN
    with trace_api.use_span(
        span,
        end_on_exit=False,
        record_exception=False,
        set_status_on_exception=False,
    ) as span:
        yield _WithSpan(
            span=span,
        )


def parse_ml_traceparent_header(
    request_headers: Optional[Dict[str, Any]] = None,
) -> Optional[str]:
    if request_headers is None:
        return None

    traceparent_header = request_headers.get(ML_TRACEPARENT_HEADER_NAME)

    if traceparent_header is None:
        return None

    return traceparent_header


async def _end_background_span(span: _WithSpan, output_attributes: Attributes, ctx: context.Context | None) -> None:
    if ctx is None or span.is_finished:
        return

    token = context.attach(ctx)
    try:
        span.finish_tracing(
            status=trace_api.Status(status_code=trace_api.StatusCode.OK),
            attributes=output_attributes,
        )
    finally:
        context.detach(token)


def traceable(
    name: Optional[str] = None,
    kind: Optional[str] = OpenInferenceSpanKindValues.CHAIN.value,
    get_input: Optional[Callable[[Any], str]] = None,
    get_output: Optional[Callable[[Any], str]] = None,
    get_request_headers: Optional[Callable[[Any], Dict[str, Any]]] = None,
):
    """
    A decorator that adds tracing functionality to a function.

    The `@traceable` decorator wraps a function with tracing logic, allowing for the automatic creation
    of a trace span when the function is called.

    Parameters
    ----------
    name : Optional[str], optional
        The name of the span to be created.
        If not provided, the function name will be used as the span name.

    kind : Optional[str], optional
        OpenInference span kind.
        Uses "CHAIN" by default.

    get_input : Optional[Callable[[Any], str]], optional
        A function that takes the function's input arguments and returns a string representation
        to be logged as part of the trace.
        If not provided, will use json-serialized function arguments.

    get_output : Optional[Callable[[Any], str]], optional
        A function that takes the function's return value and returns a string representation
        to be logged as part of the trace.
        If not provided, will use json-serialized response.

    get_request_headers: Optional[Callable[[Any], Dict[str, Any]]], optional
        A function that returns request headers dictionary.
        This function is used to extract `traceparent` header to support distributed tracing

    Returns
    -------
    Callable
        A decorated function that wraps the original function with tracing logic.

    Example
    -------
    .. code-block:: python

        from ytracing import traceable


        @traceable(name="my_custom_span")
        def my_function(x, y):
            return x + y


        @traceable(
            get_input=lambda a, b: f"Inputs: {a}, {b}",
            get_output=lambda result: f"Result: {result}"
        )
        def another_function(a, b):
            return a * b


        from fastapi import request

        @traceable(
            get_input=lambda a, b: f"Inputs: {a}, {b}",
            get_output=lambda result: f"Result: {result}",
            get_request_headers=lambda: dict(request.headers)
        )
        def another_function(a, b):
            return a * b

        # For async generators in routes with streaming responses
        @traceable(name="route_handler")
        async def stream_endpoint():
            return StreamingResponse(generate_stream())

        @traceable(name="stream_generator")
        async def generate_stream():
            for i in range(10):
                yield f"data: {i}\\n\\n"
    """

    def decorator(f):
        if hasattr(f, YTRACING_SPAN_DECORATED_ATTR):
            return f
        setattr(f, YTRACING_SPAN_DECORATED_ATTR, True)

        def get_span_attributes(f, args, kwargs):
            input_value = get_input_attribute(f, args, kwargs, get_input)
            return {
                SpanAttributes.OPENINFERENCE_SPAN_KIND: kind,
                SpanAttributes.INPUT_VALUE: input_value,
                SpanAttributes.INPUT_MIME_TYPE: (
                    OpenInferenceMimeTypeValues.TEXT.value if get_input else OpenInferenceMimeTypeValues.JSON.value
                ),
                "pod_name": str(socket.gethostname()),
            }

        @wraps(f)
        async def async_gen_wrapper(*args, **kwargs):
            tracer = trace_api.get_tracer(f.__module__ or __name__)
            span_name = name or f"{f.__qualname__}"
            request_headers = get_request_headers(*args, **kwargs) if get_request_headers else None
            yielded_values = []

            with _start_as_current_span(
                tracer=tracer,
                span_name=span_name,
                attributes=get_span_attributes(f, args, kwargs),
                traceparent=parse_ml_traceparent_header(request_headers),
            ) as with_span:
                try:
                    async_gen = f(*args, **kwargs)

                    async for item in async_gen:
                        try:
                            output = get_output_attribute(item, get_output)
                            yielded_values.append(output)
                            with_span.add_event("yielded_value")
                            yield item
                        except Exception:
                            logger.exception(f"Failed to process yielded value of type {type(item)}")
                except Exception as e:
                    with_span.record_exception(e)
                    status = trace_api.Status(
                        status_code=StatusCode.ERROR,
                        description=str(e),
                    )
                    with_span.finish_tracing(status=status)
                    raise
                finally:
                    output_attributes = {
                        SpanAttributes.OUTPUT_VALUE: yielded_values,
                        SpanAttributes.OUTPUT_MIME_TYPE: (
                            OpenInferenceMimeTypeValues.TEXT.value
                            if get_output
                            else OpenInferenceMimeTypeValues.JSON.value
                        ),
                    }
                    with_span.finish_tracing(
                        status=trace_api.Status(status_code=trace_api.StatusCode.OK),
                        attributes=output_attributes,
                    )

        @wraps(f)
        def gen_wrapper(*args, **kwargs):
            tracer = trace_api.get_tracer(f.__module__ or __name__)
            span_name = name or f"{f.__qualname__}"
            request_headers = get_request_headers(*args, **kwargs) if get_request_headers else None
            yielded_values = []

            with _start_as_current_span(
                tracer=tracer,
                span_name=span_name,
                attributes=get_span_attributes(f, args, kwargs),
                traceparent=parse_ml_traceparent_header(request_headers),
            ) as with_span:
                try:
                    gen = f(*args, **kwargs)
                    for item in gen:
                        try:
                            output = get_output_attribute(item, get_output)
                            yielded_values.append(output)
                            with_span.add_event("yielded_value")
                            yield item
                        except Exception:
                            logger.exception(f"Failed to process yielded value of type {type(item)}")
                except Exception as e:
                    with_span.record_exception(e)
                    status = trace_api.Status(
                        status_code=StatusCode.ERROR,
                        description=str(e),
                    )
                    with_span.finish_tracing(status=status)
                    raise
                finally:
                    output_attributes = {
                        SpanAttributes.OUTPUT_VALUE: yielded_values,
                        SpanAttributes.OUTPUT_MIME_TYPE: (
                            OpenInferenceMimeTypeValues.TEXT.value
                            if get_output
                            else OpenInferenceMimeTypeValues.JSON.value
                        ),
                    }
                    with_span.finish_tracing(
                        status=trace_api.Status(status_code=trace_api.StatusCode.OK),
                        attributes=output_attributes,
                    )

        @wraps(f)
        async def async_wrapper(*args, **kwargs):
            tracer = trace_api.get_tracer(f.__module__ or __name__)
            span_name = name or f"{f.__qualname__}"
            request_headers = get_request_headers(*args, **kwargs) if get_request_headers else None

            ctx = None

            with _start_as_current_span(
                tracer=tracer,
                span_name=span_name,
                attributes=get_span_attributes(f, args, kwargs),
                traceparent=parse_ml_traceparent_header(request_headers),
            ) as with_span:
                try:
                    ctx = context.get_current()
                    response = await f(*args, **kwargs)
                except Exception as e:
                    with_span.record_exception(e)
                    status = trace_api.Status(
                        status_code=StatusCode.ERROR,
                        description=str(e),
                    )
                    with_span.finish_tracing(status=status)
                    raise
                try:
                    output = get_output_attribute(response, get_output)
                    output_attributes = {
                        SpanAttributes.OUTPUT_VALUE: output,
                        SpanAttributes.OUTPUT_MIME_TYPE: (
                            OpenInferenceMimeTypeValues.TEXT.value
                            if get_output
                            else OpenInferenceMimeTypeValues.JSON.value
                        ),
                    }

                    if isinstance(response, StreamingResponse):
                        original = response.body_iterator

                        async def generator_with_ctx():
                            token = context.attach(ctx)
                            try:
                                async for chunk in original:
                                    yield chunk
                            finally:
                                context.detach(token)

                        response.body_iterator = generator_with_ctx()

                        if response.background is None:
                            response.background = BackgroundTasks()
                        response.background.add_task(_end_background_span, with_span, output_attributes, ctx)
                    else:
                        with_span.finish_tracing(
                            status=trace_api.Status(status_code=trace_api.StatusCode.OK),
                            attributes=output_attributes,
                        )

                    return response
                except Exception:
                    logger.exception(f"Failed to finalize response of type {type(response)}")
                    with_span.finish_tracing()

        @wraps(f)
        def wrapper(*args, **kwargs):
            tracer = trace_api.get_tracer(f.__module__ or __name__)
            span_name = name or f"{f.__qualname__}"
            request_headers = get_request_headers(*args, **kwargs) if get_request_headers else None

            with _start_as_current_span(
                tracer=tracer,
                span_name=span_name,
                attributes=get_span_attributes(f, args, kwargs),
                traceparent=parse_ml_traceparent_header(request_headers),
            ) as with_span:
                try:
                    response = f(*args, **kwargs)
                except Exception as e:
                    with_span.record_exception(e)
                    status = trace_api.Status(
                        status_code=StatusCode.ERROR,
                        description=str(e),
                    )
                    with_span.finish_tracing(status=status)
                    raise
                try:
                    output = get_output_attribute(response, get_output)
                    output_attributes = {
                        SpanAttributes.OUTPUT_VALUE: output,
                        SpanAttributes.OUTPUT_MIME_TYPE: (
                            OpenInferenceMimeTypeValues.TEXT.value
                            if get_output
                            else OpenInferenceMimeTypeValues.JSON.value
                        ),
                    }
                    with_span.finish_tracing(
                        status=trace_api.Status(status_code=trace_api.StatusCode.OK),
                        attributes=output_attributes,
                    )

                    return response
                except Exception:
                    logger.exception(f"Failed to finalize response of type {type(response)}")
                    with_span.finish_tracing()

        if inspect.isasyncgenfunction(f):
            return async_gen_wrapper
        elif inspect.isgeneratorfunction(f):
            return gen_wrapper
        elif inspect.iscoroutinefunction(f):
            return async_wrapper
        else:
            return wrapper

    return decorator
