"""
Real runtime execution test with detailed logging.

This test runs a real Claude Code execution and monitors:
1. Context lifecycle
2. No "no_execution_id_available" warnings
3. Delayed cleanup
4. Memory stability

Run with: python tests/test_real_runtime_execution.py
"""

import asyncio
import sys
import os
import time
from unittest.mock import MagicMock, AsyncMock

# Add parent directory to path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

# Import structlog to capture logs
import structlog
from io import StringIO

from control_plane_api.worker.runtimes.claude_code.runtime import ClaudeCodeRuntime
from control_plane_api.worker.runtimes.base import RuntimeExecutionContext
from control_plane_api.worker.runtimes.claude_code.litellm_proxy import _context_store


class LogCapture:
    """Capture logs for analysis."""
    def __init__(self):
        self.logs = []

    def __call__(self, logger, method_name, event_dict):
        self.logs.append(event_dict)
        return event_dict


async def test_real_execution_with_monitoring():
    """Run a real execution and monitor the entire lifecycle."""

    print("\n" + "="*70)
    print("  REAL WORKER EXECUTION TEST WITH CONTEXT MONITORING")
    print("="*70 + "\n")

    # Set up log capture (don't reconfigure structlog, just capture)
    log_capture = LogCapture()
    # structlog is already configured, just capture logs

    # Create runtime with mocked dependencies
    mock_control_plane = MagicMock()
    mock_control_plane.cache_metadata = MagicMock()
    mock_cancellation_manager = MagicMock()

    runtime = ClaudeCodeRuntime(
        control_plane_client=mock_control_plane,
        cancellation_manager=mock_cancellation_manager,
    )

    print("✓ Runtime initialized\n")

    # Create execution context
    execution_id = f"test-exec-{int(time.time())}"
    context = RuntimeExecutionContext(
        execution_id=execution_id,
        agent_id="test-agent",
        organization_id="test-org",
        prompt="What is 2+2? Just say the number.",
        system_prompt="You are a helpful assistant. Be concise.",
        model_id="claude-sonnet-4",
        skills=[],
        user_metadata={
            "user_id": "test-user-123",
            "session_id": "test-session-456",
        }
    )

    print(f"📋 Execution ID: {execution_id}")
    print(f"📋 User ID: test-user-123")
    print(f"📋 Session ID: test-session-456")
    print(f"📋 Model: claude-sonnet-4\n")

    # Monitor context before execution
    print("⏱️  BEFORE EXECUTION:")
    ctx = _context_store.get_context(execution_id)
    print(f"   Context exists: {ctx is not None}")
    print()

    # Start monitoring task
    monitoring_active = True
    context_checks = []

    async def monitor_context():
        """Monitor context availability during execution."""
        while monitoring_active:
            ctx = _context_store.get_context(execution_id)
            timestamp = time.time()
            context_checks.append({
                'time': timestamp,
                'available': ctx is not None
            })
            await asyncio.sleep(0.5)

    monitor_task = asyncio.create_task(monitor_context())

    # Check for ANTHROPIC_API_KEY
    if not os.getenv("ANTHROPIC_API_KEY"):
        print("⚠️  ANTHROPIC_API_KEY not set - using mock execution\n")

        # Mock execution
        print("🔄 Running MOCK execution...")

        # Simulate setting context (runtime does this)
        from control_plane_api.worker.runtimes.claude_code.litellm_proxy import set_execution_context
        set_execution_context(execution_id, {
            "user_id": "test-user-123",
            "session_id": "test-session-456",
            "organization_id": "test-org",
            "agent_id": "test-agent",
        })

        print("   ✓ Context set")

        # Check context is available
        ctx = _context_store.get_context(execution_id)
        print(f"   ✓ Context available during execution: {ctx is not None}")

        # Simulate execution work
        await asyncio.sleep(1)

        # Schedule cleanup (this is what the runtime does)
        from control_plane_api.worker.runtimes.claude_code.litellm_proxy import clear_execution_context
        print("\n   ⏰ Scheduling cleanup with 5-second delay...")
        clear_execution_context(execution_id, immediate=False, delay_seconds=5.0)

        # Context should still be available immediately
        ctx = _context_store.get_context(execution_id)
        print(f"   ✓ Context still available after scheduling cleanup: {ctx is not None}")

        # Wait 2 seconds - context should still be available
        await asyncio.sleep(2)
        ctx = _context_store.get_context(execution_id)
        print(f"   ✓ Context still available after 2s: {ctx is not None}")

        # Wait for cleanup
        print("\n   ⏳ Waiting for cleanup (5s delay)...")
        await asyncio.sleep(3.5)

        ctx = _context_store.get_context(execution_id)
        print(f"   ✓ Context cleared after delay: {ctx is None}")

    else:
        print("🔄 Running REAL execution with Anthropic API...\n")

        try:
            # Run real execution
            start_time = time.time()
            result = await runtime.execute(context)
            execution_time = time.time() - start_time

            print(f"\n✅ Execution completed in {execution_time:.2f}s")
            print(f"   Success: {result.success}")
            print(f"   Response length: {len(result.response)} chars")

            # Wait a bit for cleanup
            print("\n   ⏳ Monitoring context cleanup...")
            await asyncio.sleep(7)

        except Exception as e:
            print(f"\n❌ Execution failed: {e}")
            import traceback
            traceback.print_exc()

    # Stop monitoring
    monitoring_active = False
    await monitor_task

    # Analyze context availability
    print("\n" + "="*70)
    print("  CONTEXT AVAILABILITY TIMELINE")
    print("="*70)

    if context_checks:
        start = context_checks[0]['time']
        for check in context_checks:
            elapsed = check['time'] - start
            status = "✓ AVAILABLE" if check['available'] else "✗ CLEARED"
            print(f"   {elapsed:6.2f}s: {status}")

    # Analyze logs for warnings
    print("\n" + "="*70)
    print("  LOG ANALYSIS")
    print("="*70)

    # Note: Log capture is complex with structlog, but we can verify by checking
    # the console output above for any warnings
    print(f"\n   ✓ Check console output above for any 'no_execution_id_available' warnings")
    print(f"   ✓ You should see 'execution_context_stored' and 'execution_context_cleared' logs")

    # Final verification
    print("\n" + "="*70)
    print("  VERIFICATION RESULTS")
    print("="*70 + "\n")

    checks = {
        "Context available during execution": any(c['available'] for c in context_checks[:3]) if context_checks else False,
        "Context cleared after delay": not context_checks[-1]['available'] if context_checks else False,
        "Timeline has multiple checks": len(context_checks) > 5,
    }

    all_passed = all(checks.values())

    for check, passed in checks.items():
        symbol = "✅" if passed else "❌"
        print(f"   {symbol} {check}")

    print("\n" + "="*70)
    if all_passed:
        print("  ✅ ALL CHECKS PASSED - FIXES WORKING CORRECTLY!")
    else:
        print("  ❌ SOME CHECKS FAILED - REVIEW RESULTS ABOVE")
    print("="*70 + "\n")

    return 0 if all_passed else 1


if __name__ == "__main__":
    exit_code = asyncio.run(test_real_execution_with_monitoring())
    sys.exit(exit_code)
