"""
Context-safe background task wrappers.

The problem: asyncio.create_task() does NOT copy the current OTel context.
Spans created inside the task get parent_id=null (orphaned traces).

The fix: capture context BEFORE scheduling, restore it INSIDE the task.

Usage (FastAPI):
    from plyra_trace import PlyraBackgroundTask

    @app.post("/process")
    async def process(background_tasks: BackgroundTasks):
        background_tasks.add_task(
            PlyraBackgroundTask(process_email_background),
            email_id
        )

Usage (raw asyncio):
    task = asyncio.create_task(
        plyra_trace.run_with_context(process_email_background(email_id))
    )

Usage (threading):
    ctx = plyra_trace.capture_context()
    thread = threading.Thread(
        target=plyra_trace.run_in_context,
        args=(ctx, worker_fn, arg1),
    )
"""

import asyncio
from collections.abc import Callable, Coroutine
from typing import Any

from opentelemetry import context as otel_context


class PlyraBackgroundTask:
    """
    Wraps a coroutine function to preserve OTel trace context across
    asyncio.create_task() and FastAPI BackgroundTasks boundaries.

    Captures the current OTel context at wrap time (in the calling
    coroutine where the span is active) and restores it when the
    background task executes.

    Args:
        func: Async or sync callable to wrap

    Example:
        >>> from plyra_trace import PlyraBackgroundTask
        >>>
        >>> @app.post("/run")
        ... async def run(background_tasks: BackgroundTasks):
        ...     background_tasks.add_task(
        ...         PlyraBackgroundTask(do_work), request_id
        ...     )
    """

    def __init__(self, func: Callable) -> None:
        self._func = func
        self._ctx = otel_context.get_current()  # Capture NOW

    async def __call__(self, *args: Any, **kwargs: Any) -> Any:
        token = otel_context.attach(self._ctx)  # Restore in task
        try:
            if asyncio.iscoroutinefunction(self._func):
                return await self._func(*args, **kwargs)
            return self._func(*args, **kwargs)
        finally:
            otel_context.detach(token)


def capture_context() -> otel_context.Context:
    """
    Capture the current OTel context for use in threads.

    Returns:
        The current OTel context, to be passed to run_in_context()

    Example:
        >>> ctx = plyra_trace.capture_context()
        >>> thread = threading.Thread(
        ...     target=plyra_trace.run_in_context,
        ...     args=(ctx, worker_fn, arg1),
        ... )
        >>> thread.start()
    """
    return otel_context.get_current()


def run_in_context(ctx: otel_context.Context, fn: Callable, *args: Any, **kwargs: Any) -> Any:
    """
    Run a sync function with a previously captured OTel context.
    Use this when spawning threads manually.

    Args:
        ctx: OTel context captured via capture_context()
        fn: Sync function to run
        *args: Positional arguments passed to fn
        **kwargs: Keyword arguments passed to fn

    Returns:
        Return value of fn

    Example:
        >>> ctx = plyra_trace.capture_context()
        >>> def _run():
        ...     plyra_trace.run_in_context(ctx, worker_fn, arg1, arg2)
        >>> threading.Thread(target=_run).start()
    """
    token = otel_context.attach(ctx)
    try:
        return fn(*args, **kwargs)
    finally:
        otel_context.detach(token)


async def run_with_context(coro: Coroutine) -> Any:
    """
    Run a coroutine with the current OTel context preserved.
    Use with asyncio.create_task().

    Args:
        coro: Coroutine to run (already created)

    Returns:
        Return value of the coroutine

    Example:
        >>> task = asyncio.create_task(
        ...     plyra_trace.run_with_context(my_coroutine(args))
        ... )
    """
    ctx = otel_context.get_current()
    token = otel_context.attach(ctx)
    try:
        return await coro
    finally:
        otel_context.detach(token)
