# Copyright 2025 Tenro.ai
# SPDX-License-Identifier: Apache-2.0

"""Decorator-based linking for agents, LLMs, and tools.

Context-aware decorators that register functions with the Tenro system:
- Active Construct: Creates spans for testing and verification.
- No Construct: Executes the function without span tracking.

Examples:
    >>> from tenro import Construct, link_agent
    >>>
    >>> @link_agent("Manager")
    ... def manager_agent(task: str) -> str:
    ...     return worker_agent(task)
    >>>
    >>> with Construct() as construct:
    ...     result = manager_agent("Build feature")
    ...     # Linked to construct - can be simulated, observed, verified
"""

from __future__ import annotations

import inspect
from collections.abc import Callable
from contextvars import ContextVar
from functools import wraps
from typing import TYPE_CHECKING, Any, TypeVar

if TYPE_CHECKING:
    from tenro.construct.construct import Construct

# Type variable for decorated functions
F = TypeVar("F", bound=Callable[..., Any])

# Global registry for active construct (async-safe via contextvars)
_active_construct: ContextVar[Construct | None] = ContextVar("active_construct", default=None)


def _get_active_construct() -> Construct | None:
    """Get the currently active construct from context.

    Returns:
        The active construct, or `None` if no construct is active.
    """
    return _active_construct.get()


def _set_active_construct(construct: Construct | None) -> None:
    """Set the active construct in context.

    Args:
        `construct`: The construct to set as active, or `None` to clear.
    """
    _active_construct.set(construct)


def link_agent(name: str | None = None) -> Callable[[F], F]:
    """Decorator to register agent functions with the Tenro system.

    When a Construct is active, the decorator records an agent span. Otherwise,
    the function executes normally.
    Supports both sync and async functions.

    Args:
        `name`: Agent name for the span. If `None`, uses function name.

    Returns:
        Decorated function that registers with active Construct.

    Examples:
        >>> @link_agent("PlannerBot")
        ... def plan_trip(destination: str) -> str:
        ...     return agent.run(destination)
        >>>
        >>> with Construct() as construct:
        ...     result = plan_trip("Paris")
    """

    def decorator(func: F) -> F:
        agent_name = name or func.__name__

        # Check if function is async
        if inspect.iscoroutinefunction(func):

            @wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                construct = _get_active_construct()

                # Priority 1: Construct context
                if construct:
                    with construct.link_agent(agent_name, input_data=args) as span:
                        result = await func(*args, **kwargs)
                        span.output_data = result
                        return result

                # No construct: run function normally
                return await func(*args, **kwargs)

            # Mark wrapper with metadata for validation
            async_wrapper._tenro_linked_type = "agent"  # type: ignore[attr-defined]
            async_wrapper._tenro_linked_name = agent_name  # type: ignore[attr-defined]
            return async_wrapper  # type: ignore[return-value]
        else:

            @wraps(func)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                construct = _get_active_construct()

                # Priority 1: Construct context
                if construct:
                    with construct.link_agent(agent_name, input_data=args) as span:
                        result = func(*args, **kwargs)
                        span.output_data = result
                        return result

                # No construct: run function normally
                return func(*args, **kwargs)

            # Mark wrapper with metadata for validation
            sync_wrapper._tenro_linked_type = "agent"  # type: ignore[attr-defined]
            sync_wrapper._tenro_linked_name = agent_name  # type: ignore[attr-defined]
            return sync_wrapper  # type: ignore[return-value]

    return decorator


def link_llm(provider: str, model: str | None = None) -> Callable[[F], F]:
    """Decorator to register LLM call functions with the Tenro system.

    When a Construct is active, the decorator records an LLM span. Otherwise,
    the function executes normally.
    Supports both sync and async functions.

    Args:
        `provider`: LLM provider (e.g., "openai", "anthropic").
        `model`: Model identifier (e.g., "gpt-4", "claude-3").

    Returns:
        Decorated function that registers with active Construct.

    Examples:
        >>> @link_llm("openai", model="gpt-4")
        ... def call_llm(prompt: str) -> str:
        ...     return client.chat.completions.create(...)
        >>>
        >>> with Construct() as construct:
        ...     result = call_llm("Hello")
    """

    def decorator(func: F) -> F:
        # Capture caller info for error messages
        sig = inspect.signature(func)
        caller_signature = f"{func.__qualname__}{sig}"
        try:
            file = inspect.getsourcefile(func) or inspect.getfile(func)
            line = inspect.getsourcelines(func)[1]
            caller_location: str | None = f"{file}:{line}"
        except (OSError, TypeError):
            caller_location = None

        # Check if function is async
        if inspect.iscoroutinefunction(func):

            @wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                construct = _get_active_construct()

                # Priority 1: Construct context
                if construct:
                    with construct.link_llm(
                        provider,
                        model=model,
                        caller_signature=caller_signature,
                        caller_location=caller_location,
                    ) as span:
                        span.messages = [{"role": "user", "content": str(args)}]
                        result = await func(*args, **kwargs)
                        span.response = str(result)
                        return result

                # No construct: run function normally
                return await func(*args, **kwargs)

            return async_wrapper  # type: ignore[return-value]
        else:

            @wraps(func)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                construct = _get_active_construct()

                # Priority 1: Construct context
                if construct:
                    with construct.link_llm(
                        provider,
                        model=model,
                        caller_signature=caller_signature,
                        caller_location=caller_location,
                    ) as span:
                        span.messages = [{"role": "user", "content": str(args)}]
                        result = func(*args, **kwargs)
                        span.response = str(result)
                        return result

                # No construct: run function normally
                return func(*args, **kwargs)

            return sync_wrapper  # type: ignore[return-value]

    return decorator


def link_tool(name: str | None = None) -> Callable[[F], F]:
    """Decorator to register tool functions with the Tenro system.

    When a Construct is active, the decorator records a tool span. Otherwise,
    the function executes normally.
    Supports both sync and async functions.

    Args:
        `name`: Tool name for the span. If `None`, uses function name.

    Returns:
        Decorated function that registers with active Construct.

    Examples:
        >>> @link_tool("search")
        ... def search(query: str) -> list[str]:
        ...     return ["result1", "result2"]
        >>>
        >>> with Construct() as construct:
        ...     results = search("AI agents")
    """

    def decorator(func: F) -> F:
        tool_name = name or func.__name__

        # Check if function is async
        if inspect.iscoroutinefunction(func):

            @wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                construct = _get_active_construct()

                # Priority 1: Construct context
                if construct:
                    with construct.link_tool(tool_name) as span:
                        span.args = args
                        span.kwargs = kwargs
                        result = await func(*args, **kwargs)
                        span.result = result
                        return result

                # No construct: run function normally
                return await func(*args, **kwargs)

            # Mark wrapper with metadata for validation
            async_wrapper._tenro_linked_type = "tool"  # type: ignore[attr-defined]
            async_wrapper._tenro_linked_name = tool_name  # type: ignore[attr-defined]
            return async_wrapper  # type: ignore[return-value]
        else:

            @wraps(func)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                construct = _get_active_construct()

                # Priority 1: Construct context
                if construct:
                    with construct.link_tool(tool_name) as span:
                        span.args = args
                        span.kwargs = kwargs
                        result = func(*args, **kwargs)
                        span.result = result
                        return result

                # No construct: run function normally
                return func(*args, **kwargs)

            # Mark wrapper with metadata for validation
            sync_wrapper._tenro_linked_type = "tool"  # type: ignore[attr-defined]
            sync_wrapper._tenro_linked_name = tool_name  # type: ignore[attr-defined]
            return sync_wrapper  # type: ignore[return-value]

    return decorator


__all__ = [
    "link_agent",
    "link_llm",
    "link_tool",
    "_get_active_construct",
    "_set_active_construct",
]
