"""
Decorators for agent observability.

These decorators provide progressive integration for AI agents:
- @agent: Agent identity (Tier 2)
- @session: Session grouping (Tier 3)
- @trace_think/@trace_decide/@trace_act: Decision phases (Tier 4)
- @trace_message/@trace_delegate: Multi-agent (Tier 5)
"""

from __future__ import annotations

import functools
import inspect
from typing import Any, Callable, Dict, Optional, TypeVar, Union, overload

from risicare_core import (
    SemanticPhase,
    SpanKind,
    agent_context,
    async_agent_context,
    generate_agent_id,
    phase_context,
    session_context,
    async_session_context,
)

from risicare.tracer import get_tracer


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


# =============================================================================
# @agent Decorator
# =============================================================================

@overload
def agent(func: F) -> F: ...

@overload
def agent(
    *,
    name: Optional[str] = None,
    role: Optional[str] = None,
    agent_type: Optional[str] = None,
    version: Optional[int] = None,
) -> Callable[[F], F]: ...


def agent(
    func: Optional[F] = None,
    *,
    name: Optional[str] = None,
    role: Optional[str] = None,
    agent_type: Optional[str] = None,
    version: Optional[int] = None,
) -> Union[F, Callable[[F], F]]:
    """
    Decorator to mark a function as an agent.

    Sets up agent context for the duration of the function.
    All spans created within will be attributed to this agent.

    Can be used with or without arguments:
        @agent
        def my_agent(): ...

        @agent(name="researcher", role="worker")
        def research_agent(): ...

    Args:
        func: The function to decorate (when used without arguments).
        name: Agent name (defaults to function name).
        role: Agent role (orchestrator, worker, reviewer, etc.).
        agent_type: Framework type (langgraph, crewai, autogen, custom).
        version: Agent version number.

    Returns:
        Decorated function with agent context.

    Example:
        @agent(name="research_agent", role="worker")
        def research(query: str) -> str:
            return llm.invoke(query)

        @agent(role="orchestrator")
        async def orchestrate(task: str):
            result = await research(task)
            return summarize(result)
    """
    def decorator(fn: F) -> F:
        agent_name = name or fn.__name__
        agent_id = generate_agent_id(prefix=agent_name)

        if inspect.iscoroutinefunction(fn):
            @functools.wraps(fn)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                if not get_tracer().is_enabled:
                    return await fn(*args, **kwargs)
                async with async_agent_context(
                    agent_id=agent_id,
                    agent_name=agent_name,
                    agent_role=role,
                    agent_type=agent_type,
                    version=version,
                ):
                    tracer = get_tracer()
                    with tracer.start_span(
                        name=f"agent:{agent_name}",
                        kind=SpanKind.AGENT,
                    ) as span:
                        return await fn(*args, **kwargs)
            return async_wrapper  # type: ignore
        else:
            @functools.wraps(fn)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                if not get_tracer().is_enabled:
                    return fn(*args, **kwargs)
                with agent_context(
                    agent_id=agent_id,
                    agent_name=agent_name,
                    agent_role=role,
                    agent_type=agent_type,
                    version=version,
                ):
                    tracer = get_tracer()
                    with tracer.start_span(
                        name=f"agent:{agent_name}",
                        kind=SpanKind.AGENT,
                    ) as span:
                        return fn(*args, **kwargs)
            return sync_wrapper  # type: ignore

    if func is not None:
        return decorator(func)
    return decorator


# =============================================================================
# @session Decorator
# =============================================================================

@overload
def session(func: F) -> F: ...

@overload
def session(
    *,
    session_id_arg: str = "session_id",
    user_id_arg: Optional[str] = "user_id",
) -> Callable[[F], F]: ...


def session(
    func: Optional[F] = None,
    *,
    session_id_arg: str = "session_id",
    user_id_arg: Optional[str] = "user_id",
) -> Union[F, Callable[[F], F]]:
    """
    Decorator to create session context from function arguments.

    Extracts session_id and optionally user_id from function arguments
    to establish session context.

    Can be used with or without arguments:
        @session
        def handle_request(session_id: str): ...

        @session(session_id_arg="sid", user_id_arg="uid")
        def handle(sid: str, uid: str): ...

    Args:
        func: The function to decorate (when used without arguments).
        session_id_arg: Name of argument containing session ID.
        user_id_arg: Name of argument containing user ID (optional).

    Returns:
        Decorated function with session context.

    Example:
        @session
        def handle_request(session_id: str, query: str) -> str:
            return process(query)

        @session(session_id_arg="sid")
        async def handle(sid: str, message: str):
            return await process(message)
    """
    def decorator(fn: F) -> F:
        if inspect.iscoroutinefunction(fn):
            @functools.wraps(fn)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                # Extract session_id and user_id from arguments
                sig = inspect.signature(fn)
                bound = sig.bind(*args, **kwargs)
                bound.apply_defaults()

                sid = bound.arguments.get(session_id_arg)
                uid = bound.arguments.get(user_id_arg) if user_id_arg else None

                if sid is None:
                    # No session ID provided, run without session context
                    return await fn(*args, **kwargs)

                async with async_session_context(
                    session_id=str(sid),
                    user_id=str(uid) if uid else None,
                ):
                    return await fn(*args, **kwargs)
            return async_wrapper  # type: ignore
        else:
            @functools.wraps(fn)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                sig = inspect.signature(fn)
                bound = sig.bind(*args, **kwargs)
                bound.apply_defaults()

                sid = bound.arguments.get(session_id_arg)
                uid = bound.arguments.get(user_id_arg) if user_id_arg else None

                if sid is None:
                    return fn(*args, **kwargs)

                with session_context(
                    session_id=str(sid),
                    user_id=str(uid) if uid else None,
                ):
                    return fn(*args, **kwargs)
            return sync_wrapper  # type: ignore

    if func is not None:
        return decorator(func)
    return decorator


# =============================================================================
# @trace Decorator (Generic)
# =============================================================================

def trace(
    name: Optional[str] = None,
    kind: SpanKind = SpanKind.INTERNAL,
    attributes: Optional[dict] = None,
) -> Callable[[F], F]:
    """
    Generic decorator to trace a function with a span.

    Args:
        name: Span name (defaults to function name).
        kind: Span kind.
        attributes: Static attributes to add to span.

    Returns:
        Decorated function.

    Example:
        @trace(name="process_data", kind=SpanKind.INTERNAL)
        def process(data: dict) -> dict:
            return transform(data)
    """
    def decorator(fn: F) -> F:
        span_name = name or fn.__name__

        if inspect.iscoroutinefunction(fn):
            @functools.wraps(fn)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = get_tracer()
                with tracer.start_span(
                    name=span_name,
                    kind=kind,
                    attributes=attributes,
                ) as span:
                    return await fn(*args, **kwargs)
            return async_wrapper  # type: ignore
        else:
            @functools.wraps(fn)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = get_tracer()
                with tracer.start_span(
                    name=span_name,
                    kind=kind,
                    attributes=attributes,
                ) as span:
                    return fn(*args, **kwargs)
            return sync_wrapper  # type: ignore

    return decorator


# =============================================================================
# Phase Decorators (Tier 4)
# =============================================================================

def _phase_decorator(
    phase: SemanticPhase,
    kind: SpanKind,
    name: Optional[str] = None,
    extra_attributes: Optional[Dict[str, Any]] = None,
) -> Callable[[F], F]:
    """Internal factory for phase decorators."""
    def decorator(fn: F) -> F:
        span_name = name or f"{phase.value}:{fn.__name__}"

        if inspect.iscoroutinefunction(fn):
            @functools.wraps(fn)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                if not get_tracer().is_enabled:
                    return await fn(*args, **kwargs)
                with phase_context(phase):
                    tracer = get_tracer()
                    with tracer.start_span(
                        name=span_name,
                        kind=kind,
                    ) as span:
                        span.semantic_phase = phase
                        if extra_attributes:
                            span.set_attributes(extra_attributes)
                        return await fn(*args, **kwargs)
            return async_wrapper  # type: ignore
        else:
            @functools.wraps(fn)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                if not get_tracer().is_enabled:
                    return fn(*args, **kwargs)
                with phase_context(phase):
                    tracer = get_tracer()
                    with tracer.start_span(
                        name=span_name,
                        kind=kind,
                    ) as span:
                        span.semantic_phase = phase
                        if extra_attributes:
                            span.set_attributes(extra_attributes)
                        return fn(*args, **kwargs)
            return sync_wrapper  # type: ignore

    return decorator


@overload
def trace_think(func: F) -> F: ...

@overload
def trace_think(*, name: Optional[str] = None) -> Callable[[F], F]: ...


def trace_think(
    func: Optional[F] = None,
    *,
    name: Optional[str] = None,
) -> Union[F, Callable[[F], F]]:
    """
    Decorator for THINK phase operations.

    Use for reasoning, planning, and analysis operations.

    Example:
        @trace_think
        def analyze(data: dict) -> dict:
            return {"analysis": llm.invoke(f"Analyze: {data}")}

        @trace_think(name="strategic_planning")
        async def plan(goal: str):
            return await llm.ainvoke(f"Plan to achieve: {goal}")
    """
    decorator = _phase_decorator(
        phase=SemanticPhase.THINK,
        kind=SpanKind.THINK,
        name=name,
    )
    if func is not None:
        return decorator(func)
    return decorator


@overload
def trace_decide(func: F) -> F: ...

@overload
def trace_decide(*, name: Optional[str] = None) -> Callable[[F], F]: ...


def trace_decide(
    func: Optional[F] = None,
    *,
    name: Optional[str] = None,
) -> Union[F, Callable[[F], F]]:
    """
    Decorator for DECIDE phase operations.

    Use for decision-making and choice selection operations.

    Example:
        @trace_decide
        def select_tool(options: list) -> str:
            return llm.invoke(f"Select best tool from: {options}")
    """
    decorator = _phase_decorator(
        phase=SemanticPhase.DECIDE,
        kind=SpanKind.DECIDE,
        name=name,
    )
    if func is not None:
        return decorator(func)
    return decorator


@overload
def trace_act(func: F) -> F: ...

@overload
def trace_act(*, name: Optional[str] = None) -> Callable[[F], F]: ...


def trace_act(
    func: Optional[F] = None,
    *,
    name: Optional[str] = None,
) -> Union[F, Callable[[F], F]]:
    """
    Decorator for ACT phase operations.

    Use for action execution (tool calls, API calls, etc.).

    Example:
        @trace_act
        def execute_search(query: str) -> list:
            return search_api.search(query)

        @trace_act(name="database_write")
        async def save(data: dict):
            await db.insert(data)
    """
    decorator = _phase_decorator(
        phase=SemanticPhase.ACT,
        kind=SpanKind.TOOL_CALL,
        name=name,
    )
    if func is not None:
        return decorator(func)
    return decorator


# =============================================================================
# Multi-Agent Decorators (Tier 5)
# =============================================================================

@overload
def trace_message(func: F) -> F: ...

@overload
def trace_message(
    *,
    name: Optional[str] = None,
    target: Optional[str] = None,
    target_name: Optional[str] = None,
) -> Callable[[F], F]: ...


def trace_message(
    func: Optional[F] = None,
    *,
    name: Optional[str] = None,
    target: Optional[str] = None,
    target_name: Optional[str] = None,
) -> Union[F, Callable[[F], F]]:
    """
    Decorator for inter-agent message operations.

    Use for functions that send or receive messages between agents.

    Args:
        name: Custom span name.
        target: Target agent ID for the message recipient.
        target_name: Human-readable name of the target agent.

    Example:
        @trace_message
        def send_to_reviewer(message: str) -> str:
            return reviewer_agent.receive(message)

        @trace_message(target="reviewer", target_name="Code Reviewer")
        def send_review_request(code: str) -> str:
            return reviewer.review(code)
    """
    attrs: Dict[str, Any] = {}
    if target:
        attrs["message.target_agent_id"] = target
    if target_name:
        attrs["message.target_agent_name"] = target_name

    decorator = _phase_decorator(
        phase=SemanticPhase.COMMUNICATE,
        kind=SpanKind.MESSAGE,
        name=name,
        extra_attributes=attrs or None,
    )
    if func is not None:
        return decorator(func)
    return decorator


@overload
def trace_delegate(func: F) -> F: ...

@overload
def trace_delegate(
    *,
    name: Optional[str] = None,
    target: Optional[str] = None,
    target_name: Optional[str] = None,
) -> Callable[[F], F]: ...


def trace_delegate(
    func: Optional[F] = None,
    *,
    name: Optional[str] = None,
    target: Optional[str] = None,
    target_name: Optional[str] = None,
) -> Union[F, Callable[[F], F]]:
    """
    Decorator for task delegation operations.

    Use for functions that delegate work to other agents.

    Args:
        name: Custom span name.
        target: Target agent ID that work is delegated to.
        target_name: Human-readable name of the target agent.

    Example:
        @trace_delegate
        def delegate_research(topic: str) -> str:
            return research_agent.run(topic)

        @trace_delegate(target="research-agent", name="parallel_delegation")
        async def delegate_all(tasks: list):
            return await asyncio.gather(*[agent.run(t) for t in tasks])
    """
    attrs: Dict[str, Any] = {}
    if target:
        attrs["message.target_agent_id"] = target
    if target_name:
        attrs["message.target_agent_name"] = target_name

    decorator = _phase_decorator(
        phase=SemanticPhase.COORDINATE,
        kind=SpanKind.DELEGATION,
        name=name,
        extra_attributes=attrs or None,
    )
    if func is not None:
        return decorator(func)
    return decorator


@overload
def trace_observe(func: F) -> F: ...

@overload
def trace_observe(*, name: Optional[str] = None) -> Callable[[F], F]: ...


def trace_observe(
    func: Optional[F] = None,
    *,
    name: Optional[str] = None,
) -> Union[F, Callable[[F], F]]:
    """
    Decorator for OBSERVE phase operations.

    Use for environment observation, state reading, and information gathering.

    Example:
        @trace_observe
        def check_environment() -> dict:
            return {"status": get_system_status()}

        @trace_observe(name="read_user_state")
        async def get_user_context(user_id: str):
            return await db.get_user(user_id)
    """
    decorator = _phase_decorator(
        phase=SemanticPhase.OBSERVE,
        kind=SpanKind.OBSERVE,
        name=name,
    )
    if func is not None:
        return decorator(func)
    return decorator


@overload
def trace_coordinate(func: F) -> F: ...

@overload
def trace_coordinate(*, name: Optional[str] = None) -> Callable[[F], F]: ...


def trace_coordinate(
    func: Optional[F] = None,
    *,
    name: Optional[str] = None,
) -> Union[F, Callable[[F], F]]:
    """
    Decorator for multi-agent coordination operations.

    Use for consensus building, voting, resource allocation, and synchronization.

    Example:
        @trace_coordinate
        def gather_votes(proposal: str) -> list:
            return [agent.vote(proposal) for agent in self.agents]

        @trace_coordinate(name="consensus_building")
        async def reach_consensus(options: list):
            votes = await asyncio.gather(*[a.vote(options) for a in agents])
            return majority_vote(votes)
    """
    decorator = _phase_decorator(
        phase=SemanticPhase.COORDINATE,
        kind=SpanKind.COORDINATION,
        name=name,
    )
    if func is not None:
        return decorator(func)
    return decorator
