"""Decorator for automatic governance policy evaluation.

Provides a ``@governance`` decorator analogous to ``@trace`` that wraps
any function with pre/post governance evaluation against the Beacon API.

Usage:
    from lumenova_beacon.governance import governance

    # Without parentheses (requires stack_ids/tags set elsewhere or will
    # raise at call time):
    @governance
    def my_tool(query: str) -> str:
        ...

    # With configuration:
    @governance(stack_ids=['my-stack'], fail_open=False)
    def send_email(to: str, body: str) -> bool:
        ...

    # Async support:
    @governance(tags=['env:prod'])
    async def async_tool(data: dict) -> dict:
        ...
"""

from __future__ import annotations

import functools
import inspect
from collections.abc import Callable
from typing import Any, TypeVar, cast, overload

from lumenova_beacon.governance.client import GovernanceClient
from lumenova_beacon.governance._helpers import (
    aevaluate_or_fail_open,
    enforce,
    evaluate_or_fail_open,
    get_sdk_version,
    truncate,
)
from lumenova_beacon.tracing.trace import get_current_trace_id, get_current_span

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

_DEFAULT_HOOKS = frozenset({'pre_tool', 'post_tool', 'pre_llm', 'post_llm'})


def _get_trace_context() -> tuple[str | None, str | None]:
    """Read current trace_id and span_id from ContextVars."""
    trace_id = get_current_trace_id()
    span = get_current_span()
    return trace_id, span.span_id if span is not None else None


def _build_source(
    agent_id: str | None,
    agent_name: str | None,
) -> dict[str, Any]:
    """Build the ``input.source`` dict."""
    source: dict[str, Any] = {'sdk_version': get_sdk_version()}
    if agent_id:
        source['agent_id'] = agent_id
    if agent_name:
        source['agent_name'] = agent_name
    return source


def _build_context(
    session_id: str | None,
    environment: str | None,
    metadata: dict[str, Any] | None,
    user_id: str | None = None,
    user_email: str | None = None,
    user_name: str | None = None,
) -> dict[str, Any]:
    """Build the ``input.context`` dict."""
    ctx: dict[str, Any] = {}
    if session_id:
        ctx['session_id'] = session_id
    if environment:
        ctx['environment'] = environment
    if user_id:
        ctx['user_id'] = user_id
    if user_email:
        ctx['user_email'] = user_email
    if user_name:
        ctx['user_name'] = user_name
    if metadata:
        ctx['metadata'] = metadata
    return ctx


# Type overloads for proper type hints in both usage patterns.
@overload
def governance(_func: F) -> F: ...


@overload
def governance(
    _func: None = None,
    *,
    name: str | None = None,
    type_: str = 'tool',
    stack_ids: list[str] | None = None,
    tags: list[str] | None = None,
    fail_open: bool = True,
    enabled_hooks: set[str] | None = None,
    result_parser: Callable[[dict[str, Any]], bool] | None = None,
    timeout: float | None = None,
    agent_id: str | None = None,
    agent_name: str | None = None,
    session_id: str | None = None,
    environment: str | None = None,
    metadata: dict[str, Any] | None = None,
) -> Callable[[F], F]: ...


def governance(
    _func: Callable | None = None,
    *,
    name: str | None = None,
    type_: str = 'tool',
    stack_ids: list[str] | None = None,
    tags: list[str] | None = None,
    fail_open: bool = True,
    enabled_hooks: set[str] | None = None,
    result_parser: Callable[[dict[str, Any]], bool] | None = None,
    timeout: float | None = None,
    agent_id: str | None = None,
    agent_name: str | None = None,
    session_id: str | None = None,
    environment: str | None = None,
    metadata: dict[str, Any] | None = None,
) -> Callable[[F], F] | F:
    """Decorator to enforce governance policies on function calls.

    Evaluates governance policies before and/or after the decorated function
    executes.  Raises
    :class:`~lumenova_beacon.exceptions.GovernanceViolationError` if a
    policy blocks the action.

    This decorator supports three usage patterns:

    - ``@governance`` — bare decorator (no parentheses)
    - ``@governance()`` — empty parentheses
    - ``@governance(stack_ids=['s1'], ...)`` — with configuration

    Args:
        _func: The function being decorated (when used without parentheses).
        name: Custom name sent to the governance API.  Defaults to the
            decorated function's ``__name__``.
        type_: ``'tool'`` (default) or ``'llm'``.  Determines the payload
            shape and which ``enabled_hooks`` keys apply.
        stack_ids: Explicit policy stack IDs.  At least one of ``stack_ids``
            or ``tags`` must be provided.
        tags: Tag-based policy stack discovery.
        fail_open: Allow the action when the governance API is unreachable.
        enabled_hooks: Which hooks to evaluate.  Relevant keys depend on
            ``type_``: ``'pre_tool'``/``'post_tool'`` for tools,
            ``'pre_llm'``/``'post_llm'`` for LLM calls.
        result_parser: Custom ``(result_dict) -> bool`` parser.
        timeout: Override request timeout in seconds.
        agent_id: Agent identifier, included in ``input.source``.
        agent_name: Human-readable agent name, included in ``input.source``.
        session_id: Override session ID. Defaults to ``BeaconClient.config.session_id``.
        environment: Override environment. Defaults to ``BeaconClient.config.environment``.
        metadata: Custom metadata included in ``input.context.metadata``.

    Returns:
        The decorated function (or a decorator if called with arguments).
    """
    hooks = set(enabled_hooks) if enabled_hooks is not None else set(_DEFAULT_HOOKS)

    def decorator(func: F) -> F:
        func_name = name or func.__name__
        client = GovernanceClient()

        # Resolve defaults from BeaconClient config
        from lumenova_beacon.core.client import get_client
        beacon_client = get_client()
        resolved_session_id = session_id if session_id is not None else beacon_client.config.session_id
        resolved_environment = environment if environment is not None else beacon_client.config.environment

        source = _build_source(agent_id, agent_name)
        context = _build_context(
            resolved_session_id, resolved_environment, metadata,
            user_id=beacon_client.config.user_id,
            user_email=beacon_client.config.user_email,
            user_name=beacon_client.config.user_name,
        )

        pre_hook_name = f'pre_{type_}'
        post_hook_name = f'post_{type_}'

        @functools.wraps(func)
        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
            parameters = kwargs.copy()

            # Pre-execution check.
            if pre_hook_name in hooks:
                extra: dict[str, Any] = {
                    type_: {'name': func_name, 'parameters': parameters},
                    'source': source,
                    'context': context,
                }
                trace_id, span_id = _get_trace_context()
                resp = evaluate_or_fail_open(
                    client, 'pre', type_,
                    stack_ids, tags, extra, timeout, fail_open,
                    trace_id=trace_id, span_id=span_id,
                )
                enforce(resp, 'pre', f"{type_} '{func_name}'", result_parser)

            result = func(*args, **kwargs)

            # Post-execution check.
            if post_hook_name in hooks:
                post_extra: dict[str, Any] = {
                    type_: {
                        'name': func_name,
                        'parameters': parameters,
                        'output': truncate(str(result)),
                    },
                    'source': source,
                    'context': context,
                }
                trace_id, span_id = _get_trace_context()
                resp = evaluate_or_fail_open(
                    client, 'post', type_,
                    stack_ids, tags, post_extra, timeout, fail_open,
                    trace_id=trace_id, span_id=span_id,
                )
                enforce(resp, 'post', f"{type_} '{func_name}'", result_parser)

            return result

        @functools.wraps(func)
        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
            parameters = kwargs.copy()

            # Pre-execution check.
            if pre_hook_name in hooks:
                extra: dict[str, Any] = {
                    type_: {'name': func_name, 'parameters': parameters},
                    'source': source,
                    'context': context,
                }
                trace_id, span_id = _get_trace_context()
                resp = await aevaluate_or_fail_open(
                    client, 'pre', type_,
                    stack_ids, tags, extra, timeout, fail_open,
                    trace_id=trace_id, span_id=span_id,
                )
                enforce(resp, 'pre', f"{type_} '{func_name}'", result_parser)

            result = await func(*args, **kwargs)

            # Post-execution check.
            if post_hook_name in hooks:
                post_extra: dict[str, Any] = {
                    type_: {
                        'name': func_name,
                        'parameters': parameters,
                        'output': truncate(str(result)),
                    },
                    'source': source,
                    'context': context,
                }
                trace_id, span_id = _get_trace_context()
                resp = await aevaluate_or_fail_open(
                    client, 'post', type_,
                    stack_ids, tags, post_extra, timeout, fail_open,
                    trace_id=trace_id, span_id=span_id,
                )
                enforce(resp, 'post', f"{type_} '{func_name}'", result_parser)

            return result

        if inspect.iscoroutinefunction(func):
            return cast(F, async_wrapper)
        return cast(F, sync_wrapper)

    # Detect how the decorator is being used.
    if _func is None:
        # Called with parentheses: @governance() or @governance(stack_ids=[...])
        return decorator
    else:
        # Called without parentheses: @governance
        return decorator(cast(F, _func))
