"""
Tests for SDK decorators.

Tests cover:
- @agent decorator (sync and async)
- @session decorator (sync and async)
- @trace decorator
- Phase decorators (@trace_think, @trace_decide, @trace_act, @trace_observe)
- Multi-agent decorators (@trace_message, @trace_delegate, @trace_coordinate)
"""

import asyncio
import pytest
from unittest.mock import MagicMock, patch

from risicare_core import (
    get_current_agent,
    get_current_session,
    get_current_phase,
    SemanticPhase,
    _reset_context_for_testing,
)

from risicare import (
    agent,
    session,
    trace,
    trace_think,
    trace_decide,
    trace_act,
    trace_observe,
    trace_message,
    trace_delegate,
    trace_coordinate,
    _reset_client_for_testing,
    _reset_tracer_for_testing,
)


# =============================================================================
# Fixtures
# =============================================================================

@pytest.fixture(autouse=True)
def reset_all():
    """Reset all state before and after each test."""
    _reset_context_for_testing()
    _reset_client_for_testing()
    _reset_tracer_for_testing()
    yield
    _reset_context_for_testing()
    _reset_client_for_testing()
    _reset_tracer_for_testing()


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

class TestAgentDecorator:
    def test_agent_sync_basic(self):
        """@agent decorator works on sync functions."""
        @agent
        def my_agent():
            ctx = get_current_agent()
            return ctx.agent_id if ctx else None

        result = my_agent()
        assert result == "my_agent"

    def test_agent_sync_with_args(self):
        """@agent decorator with explicit arguments."""
        @agent(name="custom_name", role="worker")
        def my_agent():
            ctx = get_current_agent()
            return (ctx.agent_name, ctx.agent_role) if ctx else None

        result = my_agent()
        assert result == ("custom_name", "worker")

    def test_agent_preserves_function_args(self):
        """@agent decorator preserves function arguments."""
        @agent
        def add_numbers(a, b, multiplier=1):
            return (a + b) * multiplier

        assert add_numbers(2, 3) == 5
        assert add_numbers(2, 3, multiplier=2) == 10

    def test_agent_preserves_return_value(self):
        """@agent decorator preserves return value."""
        @agent
        def get_data():
            return {"key": "value", "count": 42}

        result = get_data()
        assert result == {"key": "value", "count": 42}

    @pytest.mark.asyncio
    async def test_agent_async(self):
        """@agent decorator works on async functions."""
        @agent(name="async_agent")
        async def my_async_agent():
            await asyncio.sleep(0)
            ctx = get_current_agent()
            return ctx.agent_name if ctx else None

        result = await my_async_agent()
        assert result == "async_agent"

    @pytest.mark.asyncio
    async def test_agent_async_preserves_args(self):
        """@agent decorator on async preserves arguments."""
        @agent
        async def process(value: int) -> int:
            await asyncio.sleep(0)
            return value * 2

        result = await process(21)
        assert result == 42

    def test_agent_context_not_leaked(self):
        """Agent context doesn't leak outside decorated function."""
        @agent(name="scoped_agent")
        def my_agent():
            return get_current_agent().agent_name

        assert get_current_agent() is None
        result = my_agent()
        assert result == "scoped_agent"
        assert get_current_agent() is None


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

class TestSessionDecorator:
    def test_session_sync_basic(self):
        """@session decorator extracts session_id from args."""
        @session
        def handle_request(session_id: str):
            ctx = get_current_session()
            return ctx.session_id if ctx else None

        result = handle_request("sess-123")
        assert result == "sess-123"

    def test_session_with_user_id(self):
        """@session decorator extracts user_id."""
        @session
        def handle_request(session_id: str, user_id: str):
            ctx = get_current_session()
            return (ctx.session_id, ctx.user_id) if ctx else None

        result = handle_request("sess-123", "user-456")
        assert result == ("sess-123", "user-456")

    def test_session_custom_arg_names(self):
        """@session decorator with custom argument names."""
        @session(session_id_arg="sid", user_id_arg="uid")
        def handle(sid: str, uid: str, data: dict):
            ctx = get_current_session()
            return (ctx.session_id, ctx.user_id) if ctx else None

        result = handle("s-1", "u-1", {"key": "value"})
        assert result == ("s-1", "u-1")

    def test_session_missing_id_no_context(self):
        """@session without session_id doesn't create context."""
        @session(session_id_arg="sid")
        def handle(data: dict):  # No sid argument
            return get_current_session()

        result = handle({"key": "value"})
        assert result is None

    @pytest.mark.asyncio
    async def test_session_async(self):
        """@session decorator works on async functions."""
        @session
        async def async_handle(session_id: str):
            await asyncio.sleep(0)
            ctx = get_current_session()
            return ctx.session_id if ctx else None

        result = await async_handle("async-sess")
        assert result == "async-sess"

    def test_session_context_not_leaked(self):
        """Session context doesn't leak outside decorated function."""
        @session
        def handle(session_id: str):
            return get_current_session().session_id

        assert get_current_session() is None
        result = handle("test-sess")
        assert result == "test-sess"
        assert get_current_session() is None


# =============================================================================
# @trace Decorator Tests
# =============================================================================

class TestTraceDecorator:
    def test_trace_basic(self):
        """@trace decorator creates span."""
        @trace(name="my_operation")
        def my_func():
            return "result"

        # Should not raise, should return value
        result = my_func()
        assert result == "result"

    def test_trace_default_name(self):
        """@trace uses function name as default span name."""
        @trace()
        def named_function():
            return 42

        result = named_function()
        assert result == 42

    @pytest.mark.asyncio
    async def test_trace_async(self):
        """@trace decorator works on async functions."""
        @trace(name="async_op")
        async def async_func():
            await asyncio.sleep(0)
            return "async_result"

        result = await async_func()
        assert result == "async_result"


# =============================================================================
# Phase Decorator Tests
# =============================================================================

class TestPhaseDecorators:
    def test_trace_think_sets_phase(self):
        """@trace_think sets THINK phase."""
        @trace_think
        def analyze():
            return get_current_phase()

        result = analyze()
        assert result == SemanticPhase.THINK

    def test_trace_decide_sets_phase(self):
        """@trace_decide sets DECIDE phase."""
        @trace_decide
        def choose():
            return get_current_phase()

        result = choose()
        assert result == SemanticPhase.DECIDE

    def test_trace_act_sets_phase(self):
        """@trace_act sets ACT phase."""
        @trace_act
        def execute():
            return get_current_phase()

        result = execute()
        assert result == SemanticPhase.ACT

    def test_trace_observe_sets_phase(self):
        """@trace_observe sets OBSERVE phase."""
        @trace_observe
        def watch():
            return get_current_phase()

        result = watch()
        assert result == SemanticPhase.OBSERVE

    def test_phase_decorators_with_name(self):
        """Phase decorators accept custom name."""
        @trace_think(name="custom_thinking")
        def my_think():
            return get_current_phase()

        result = my_think()
        assert result == SemanticPhase.THINK

    @pytest.mark.asyncio
    async def test_phase_decorators_async(self):
        """Phase decorators work on async functions."""
        @trace_think
        async def async_think():
            await asyncio.sleep(0)
            return get_current_phase()

        @trace_decide
        async def async_decide():
            await asyncio.sleep(0)
            return get_current_phase()

        @trace_act
        async def async_act():
            await asyncio.sleep(0)
            return get_current_phase()

        assert await async_think() == SemanticPhase.THINK
        assert await async_decide() == SemanticPhase.DECIDE
        assert await async_act() == SemanticPhase.ACT

    def test_phase_not_leaked(self):
        """Phase context doesn't leak."""
        @trace_think
        def think():
            return get_current_phase()

        assert get_current_phase() is None
        result = think()
        assert result == SemanticPhase.THINK
        assert get_current_phase() is None


# =============================================================================
# Multi-Agent Decorator Tests
# =============================================================================

class TestMultiAgentDecorators:
    def test_trace_message_sets_phase(self):
        """@trace_message sets COMMUNICATE phase."""
        @trace_message
        def send_message():
            return get_current_phase()

        result = send_message()
        assert result == SemanticPhase.COMMUNICATE

    def test_trace_delegate_sets_phase(self):
        """@trace_delegate sets COORDINATE phase."""
        @trace_delegate
        def delegate_task():
            return get_current_phase()

        result = delegate_task()
        assert result == SemanticPhase.COORDINATE

    def test_trace_coordinate_sets_phase(self):
        """@trace_coordinate sets COORDINATE phase."""
        @trace_coordinate
        def coordinate_agents():
            return get_current_phase()

        result = coordinate_agents()
        assert result == SemanticPhase.COORDINATE

    @pytest.mark.asyncio
    async def test_multi_agent_decorators_async(self):
        """Multi-agent decorators work on async functions."""
        @trace_message
        async def async_message():
            await asyncio.sleep(0)
            return get_current_phase()

        @trace_delegate
        async def async_delegate():
            await asyncio.sleep(0)
            return get_current_phase()

        @trace_coordinate
        async def async_coordinate():
            await asyncio.sleep(0)
            return get_current_phase()

        assert await async_message() == SemanticPhase.COMMUNICATE
        assert await async_delegate() == SemanticPhase.COORDINATE
        assert await async_coordinate() == SemanticPhase.COORDINATE


# =============================================================================
# Nested Decorator Tests
# =============================================================================

class TestNestedDecorators:
    def test_session_with_agent(self):
        """@session and @agent can be combined."""
        @session
        @agent(name="nested_agent")
        def handle(session_id: str):
            sess = get_current_session()
            agt = get_current_agent()
            return (sess.session_id, agt.agent_name)

        result = handle("sess-nested")
        assert result == ("sess-nested", "nested_agent")

    def test_agent_with_phase(self):
        """@agent and phase decorators can be combined."""
        @agent(name="thinking_agent")
        @trace_think
        def think_and_act():
            agt = get_current_agent()
            phase = get_current_phase()
            return (agt.agent_name, phase)

        result = think_and_act()
        assert result == ("thinking_agent", SemanticPhase.THINK)

    @pytest.mark.asyncio
    async def test_nested_async_decorators(self):
        """Nested decorators work with async."""
        @session
        @agent(name="async_nested")
        @trace_decide
        async def complex_handler(session_id: str):
            await asyncio.sleep(0)
            sess = get_current_session()
            agt = get_current_agent()
            phase = get_current_phase()
            return (sess.session_id, agt.agent_name, phase)

        result = await complex_handler("async-sess")
        assert result == ("async-sess", "async_nested", SemanticPhase.DECIDE)
