"""
Test HITL (Human-in-the-Loop) Trace Continuation

Tests the v0.6.3 feature where HITL interrupts continue the same trace
instead of creating new traces.

Key scenarios:
1. Initial invocation creates new trace and stores mapping
2. Resume invocation (same thread_id) continues existing trace
3. HITL markers are present in continued observations
"""

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


class TestHitlTraceContinuation:
    """Test suite for HITL trace continuation feature"""

    @pytest.mark.asyncio
    async def test_store_thread_trace_mapping(self):
        """Test that thread->trace mapping is stored in Redis"""
        mock_redis = AsyncMock()
        mock_redis.setex = AsyncMock(return_value=True)
        
        with patch('swisper_studio_sdk.tracing.redis_publisher._redis_client', mock_redis):
            with patch('swisper_studio_sdk.tracing.redis_publisher._project_id', 'test-project'):
                with patch('swisper_studio_sdk.tracing.redis_publisher._initialized', True):
                    with patch('swisper_studio_sdk.tracing.redis_publisher._initialization_failed', False):
                        from swisper_studio_sdk.tracing.redis_publisher import store_thread_trace_mapping
                        
                        result = await store_thread_trace_mapping(
                            thread_id="chat-123",
                            trace_id="trace-456"
                        )
                        
                        assert result == True
                        mock_redis.setex.assert_called_once()
                        call_args = mock_redis.setex.call_args
                        # Check key format
                        assert "swisper_studio:thread_trace:test-project:chat-123" in str(call_args)

    @pytest.mark.asyncio
    async def test_get_trace_for_thread_found(self):
        """Test retrieving existing trace for thread"""
        mock_redis = AsyncMock()
        mock_redis.get = AsyncMock(return_value=b"existing-trace-id")
        
        with patch('swisper_studio_sdk.tracing.redis_publisher._redis_client', mock_redis):
            with patch('swisper_studio_sdk.tracing.redis_publisher._project_id', 'test-project'):
                with patch('swisper_studio_sdk.tracing.redis_publisher._initialized', True):
                    with patch('swisper_studio_sdk.tracing.redis_publisher._initialization_failed', False):
                        from swisper_studio_sdk.tracing.redis_publisher import get_trace_for_thread
                        
                        result = await get_trace_for_thread(thread_id="chat-123")
                        
                        assert result == "existing-trace-id"

    @pytest.mark.asyncio
    async def test_get_trace_for_thread_not_found(self):
        """Test when no trace exists for thread"""
        mock_redis = AsyncMock()
        mock_redis.get = AsyncMock(return_value=None)
        
        with patch('swisper_studio_sdk.tracing.redis_publisher._redis_client', mock_redis):
            with patch('swisper_studio_sdk.tracing.redis_publisher._project_id', 'test-project'):
                with patch('swisper_studio_sdk.tracing.redis_publisher._initialized', True):
                    with patch('swisper_studio_sdk.tracing.redis_publisher._initialization_failed', False):
                        from swisper_studio_sdk.tracing.redis_publisher import get_trace_for_thread
                        
                        result = await get_trace_for_thread(thread_id="new-chat")
                        
                        assert result is None

    @pytest.mark.asyncio
    async def test_clear_thread_trace_mapping(self):
        """Test clearing thread->trace mapping"""
        mock_redis = AsyncMock()
        mock_redis.delete = AsyncMock(return_value=1)
        
        with patch('swisper_studio_sdk.tracing.redis_publisher._redis_client', mock_redis):
            with patch('swisper_studio_sdk.tracing.redis_publisher._project_id', 'test-project'):
                with patch('swisper_studio_sdk.tracing.redis_publisher._initialized', True):
                    with patch('swisper_studio_sdk.tracing.redis_publisher._initialization_failed', False):
                        from swisper_studio_sdk.tracing.redis_publisher import clear_thread_trace_mapping
                        
                        result = await clear_thread_trace_mapping(thread_id="chat-123")
                        
                        assert result == True
                        mock_redis.delete.assert_called_once()


class TestHitlContinuationDetection:
    """Test HITL continuation detection in graph wrapper"""

    @pytest.mark.asyncio
    async def test_detect_hitl_continuation_from_config(self):
        """Test that HITL continuation is detected from thread_id in config"""
        # This test validates the logic flow without full LangGraph setup
        from swisper_studio_sdk.tracing.graph_wrapper import _handle_top_level_agent
        
        # Mock the necessary functions
        mock_ainvoke = AsyncMock(return_value={"result": "success"})
        
        config = {
            "configurable": {
                "thread_id": "chat-123"
            }
        }
        input_state = {
            "user_message": "Yes, send the email",
            "chat_id": "chat-123"
        }
        
        # Mock get_trace_for_thread to return existing trace
        with patch('swisper_studio_sdk.tracing.graph_wrapper.get_trace_for_thread') as mock_get_trace:
            mock_get_trace.return_value = "existing-trace-id"
            
            with patch('swisper_studio_sdk.tracing.graph_wrapper._handle_hitl_continuation') as mock_hitl:
                mock_hitl.return_value = {"result": "success"}
                
                with patch('swisper_studio_sdk.tracing.graph_wrapper.is_tracing_initialized', return_value=True):
                    with patch('swisper_studio_sdk.tracing.graph_wrapper.get_project_id', return_value='test-project'):
                        with patch('swisper_studio_sdk.tracing.graph_wrapper._is_redis_healthy', return_value=True):
                            with patch('swisper_studio_sdk.tracing.graph_wrapper.is_tracing_enabled_for_project', return_value=True):
                                with patch('swisper_studio_sdk.tracing.graph_wrapper._start_heartbeat_task'):
                                    # Execute
                                    result = await _handle_top_level_agent(
                                        mock_ainvoke,
                                        input_state,
                                        config,
                                        {},
                                        "test_agent"
                                    )
                                    
                                    # Verify HITL continuation was detected and called
                                    mock_get_trace.assert_called_once_with("chat-123")
                                    mock_hitl.assert_called_once()


class TestHitlObservationMarkers:
    """Test that HITL observations have proper markers"""

    @pytest.mark.asyncio
    async def test_hitl_continuation_markers_in_observation(self):
        """Test that HITL continuation observations have proper markers"""
        published_events = []
        
        async def mock_publish(event_type, trace_id, observation_id=None, data=None, timeout=None):
            published_events.append({
                "event_type": event_type,
                "trace_id": trace_id,
                "observation_id": observation_id,
                "data": data
            })
            return True
        
        mock_ainvoke = AsyncMock(return_value={"result": "email sent"})
        
        with patch('swisper_studio_sdk.tracing.graph_wrapper.publish_event', side_effect=mock_publish):
            with patch('swisper_studio_sdk.tracing.graph_wrapper.set_current_trace'):
                with patch('swisper_studio_sdk.tracing.graph_wrapper.set_current_observation', return_value="token"):
                    from swisper_studio_sdk.tracing.graph_wrapper import _handle_hitl_continuation
                    
                    result = await _handle_hitl_continuation(
                        mock_ainvoke,
                        {"user_message": "Yes, send it", "chat_id": "chat-123"},
                        {"configurable": {"thread_id": "chat-123"}},
                        {},
                        "productivity_agent",
                        "existing-trace-id",
                        "chat-123"
                    )
                    
                    # Verify observation_start has HITL markers
                    start_events = [e for e in published_events if e["event_type"] == "observation_start"]
                    assert len(start_events) == 1
                    
                    start_data = start_events[0]["data"]
                    assert start_data["is_hitl_continuation"] == True
                    assert "HITL Continuation" in start_data["name"]
                    assert start_data["thread_id"] == "chat-123"
                    
                    # Verify trace_id is the existing one
                    assert start_events[0]["trace_id"] == "existing-trace-id"
                    
                    # Verify observation_end also has HITL marker
                    end_events = [e for e in published_events if e["event_type"] == "observation_end"]
                    assert len(end_events) == 1
                    assert end_events[0]["data"]["is_hitl_continuation"] == True


# Run tests if executed directly
if __name__ == "__main__":
    pytest.main([__file__, "-v"])
