# -*- coding: utf-8 -*-
"""Location: ./tests/integration/test_streamable_http_redis.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

Integration tests for multi-worker Streamable HTTP with RedisEventStore.

Tests that stateful sessions work correctly across multiple gateway workers
using Redis as the shared event store backend.
"""

import asyncio
import os
import uuid
from unittest.mock import AsyncMock

import pytest

from mcpgateway.transports.redis_event_store import RedisEventStore
from mcpgateway.utils.redis_client import get_redis_client


@pytest.fixture
async def redis_cleanup(monkeypatch):
    """Clean up Redis keys before and after tests."""
    # Enable Redis for this test
    monkeypatch.setenv("CACHE_TYPE", "redis")
    monkeypatch.setenv("REDIS_URL", "redis://localhost:6379")

    # Reset Redis client to pick up new settings
    from mcpgateway.utils.redis_client import _reset_client
    from mcpgateway.config import settings

    _reset_client()
    # Force settings reload by directly setting attributes
    settings.cache_type = "redis"
    settings.redis_url = "redis://localhost:6379"

    redis = await get_redis_client()
    key_prefix = f"mcpgw:eventstore:itest:{uuid.uuid4().hex}"

    # Clean before
    if redis:
        keys = await redis.keys(f"{key_prefix}:*")
        if keys:
            await redis.delete(*keys)

    yield key_prefix

    # Clean after
    if redis:
        keys = await redis.keys(f"{key_prefix}:*")
        if keys:
            await redis.delete(*keys)

    # Reset client after test
    _reset_client()


class TestMultiWorkerStatefulSessions:
    """Test stateful sessions across multiple workers."""

    async def test_multi_worker_event_sharing(self, redis_cleanup):
        """
        Simulate multiple workers accessing the same stream.

        Worker 1 stores events, Worker 2 replays them.
        """
        # Create two event stores (simulating two workers)
        worker1_store = RedisEventStore(max_events_per_stream=100, ttl=300, key_prefix=redis_cleanup)
        worker2_store = RedisEventStore(max_events_per_stream=100, ttl=300, key_prefix=redis_cleanup)

        stream_id = "session-abc123"

        # Worker 1: Store initialize event
        init_msg = {"jsonrpc": "2.0", "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}}, "id": 1}
        event1 = await worker1_store.store_event(stream_id, init_msg)

        # Worker 1: Store tools/list request
        tools_req = {"jsonrpc": "2.0", "method": "tools/list", "id": 2}
        event2 = await worker1_store.store_event(stream_id, tools_req)

        # Worker 1: Store tools/list response
        tools_resp = {"jsonrpc": "2.0", "result": {"tools": [{"name": "test_tool", "description": "A test tool"}]}, "id": 2}
        event3 = await worker1_store.store_event(stream_id, tools_resp)

        # Worker 2: Replay from event1 (should see events 2 and 3)
        replayed = []

        async def callback(msg):
            replayed.append(msg)

        result = await worker2_store.replay_events_after(event1, callback)

        # Verify Worker 2 sees all events stored by Worker 1
        assert result == stream_id
        assert len(replayed) == 2
        assert replayed[0] == tools_req
        assert replayed[1] == tools_resp

    async def test_concurrent_worker_writes(self, redis_cleanup):
        """
        Multiple workers writing to different streams simultaneously.
        """
        num_workers = 5
        events_per_worker = 10

        async def worker_task(worker_id):
            """Simulate a worker storing events."""
            store = RedisEventStore(max_events_per_stream=100, ttl=300, key_prefix=redis_cleanup)
            stream_id = f"worker-{worker_id}-stream"
            event_ids = []

            for i in range(events_per_worker):
                msg = {"jsonrpc": "2.0", "method": f"worker_{worker_id}_method_{i}", "id": i}
                event_id = await store.store_event(stream_id, msg)
                event_ids.append(event_id)
                # Small delay to simulate real processing
                await asyncio.sleep(0.01)

            return stream_id, event_ids

        # Run all workers concurrently
        results = await asyncio.gather(*[worker_task(i) for i in range(num_workers)])

        # Verify each worker's stream
        for worker_id, (stream_id, event_ids) in enumerate(results):
            assert len(event_ids) == events_per_worker

            # Create new store to verify persistence
            verify_store = RedisEventStore(max_events_per_stream=100, ttl=300, key_prefix=redis_cleanup)

            # Replay from first event
            replayed = []

            async def callback(msg):
                replayed.append(msg)

            result = await verify_store.replay_events_after(event_ids[0], callback)

            assert result == stream_id
            assert len(replayed) == events_per_worker - 1  # All except first

    async def test_worker_crash_recovery(self, redis_cleanup):
        """
        Worker 1 stores events, crashes, Worker 2 continues from same stream.
        """
        stream_id = "resilient-session"

        # Worker 1: Store some events
        worker1_store = RedisEventStore(max_events_per_stream=100, ttl=300, key_prefix=redis_cleanup)
        msg1 = {"jsonrpc": "2.0", "method": "initialize", "id": 1}
        event1 = await worker1_store.store_event(stream_id, msg1)

        msg2 = {"jsonrpc": "2.0", "method": "tools/list", "id": 2}
        event2 = await worker1_store.store_event(stream_id, msg2)

        # Simulate Worker 1 crash (delete the object)
        del worker1_store

        # Worker 2: Takes over the stream
        worker2_store = RedisEventStore(max_events_per_stream=100, ttl=300, key_prefix=redis_cleanup)

        # Worker 2 can replay existing events
        replayed = []

        async def callback(msg):
            replayed.append(msg)

        result = await worker2_store.replay_events_after(event1, callback)

        assert result == stream_id
        assert len(replayed) == 1
        assert replayed[0] == msg2

        # Worker 2 can continue adding events
        msg3 = {"jsonrpc": "2.0", "result": {"tools": []}, "id": 2}
        event3 = await worker2_store.store_event(stream_id, msg3)

        # Verify all events are there
        replayed.clear()
        result = await worker2_store.replay_events_after(event1, callback)

        assert len(replayed) == 2
        assert replayed[0] == msg2
        assert replayed[1] == msg3

    async def test_round_robin_load_balancing(self, redis_cleanup):
        """
        Simulate round-robin load balancing where requests alternate between workers.
        """
        stream_id = "load-balanced-session"

        # Create 3 workers
        workers = [RedisEventStore(max_events_per_stream=100, ttl=300, key_prefix=redis_cleanup) for _ in range(3)]

        # Simulate 12 requests round-robin across workers
        event_ids = []
        messages = []

        for i in range(12):
            worker = workers[i % 3]  # Round-robin
            msg = {"jsonrpc": "2.0", "method": f"request_{i}", "id": i}
            messages.append(msg)
            event_id = await worker.store_event(stream_id, msg)
            event_ids.append(event_id)

        # Any worker can replay the full history
        for worker_idx, worker in enumerate(workers):
            replayed = []

            async def callback(msg):
                replayed.append(msg)

            result = await worker.replay_events_after(event_ids[0], callback)

            assert result == stream_id
            assert len(replayed) == 11  # All messages after first
            # Verify order is preserved
            for i, msg in enumerate(replayed):
                assert msg == messages[i + 1]

    async def test_eviction_across_workers(self, redis_cleanup):
        """
        Test that eviction works correctly when multiple workers write to same stream.
        """
        stream_id = "eviction-test-stream"
        max_events = 10

        # Create two workers
        worker1 = RedisEventStore(max_events_per_stream=max_events, ttl=300, key_prefix=redis_cleanup)
        worker2 = RedisEventStore(max_events_per_stream=max_events, ttl=300, key_prefix=redis_cleanup)

        # Worker 1 stores 5 events
        event_ids = []
        for i in range(5):
            msg = {"jsonrpc": "2.0", "method": f"worker1_msg_{i}", "id": i}
            event_id = await worker1.store_event(stream_id, msg)
            event_ids.append(event_id)

        # Worker 2 stores 10 more events (total 15, should evict first 5)
        for i in range(5, 15):
            msg = {"jsonrpc": "2.0", "method": f"worker2_msg_{i}", "id": i}
            event_id = await worker2.store_event(stream_id, msg)
            event_ids.append(event_id)

        # Try to replay from first event (should be evicted)
        replayed = []

        async def callback(msg):
            replayed.append(msg)

        result = await worker1.replay_events_after(event_ids[0], callback)
        assert result is None  # First event evicted

        # Replay from 6th event (should exist)
        replayed.clear()
        result = await worker2.replay_events_after(event_ids[5], callback)

        assert result == stream_id
        assert len(replayed) == 9  # Events 7-15

    async def test_session_isolation(self, redis_cleanup):
        """
        Ensure different sessions don't interfere with each other.
        """
        # Create multiple workers
        workers = [RedisEventStore(max_events_per_stream=100, ttl=300, key_prefix=redis_cleanup) for _ in range(3)]

        # Create 3 different sessions
        sessions = ["session-alice", "session-bob", "session-charlie"]

        # Each worker manages one session
        session_events = {}
        for i, (worker, session_id) in enumerate(zip(workers, sessions)):
            events = []
            for j in range(5):
                msg = {"jsonrpc": "2.0", "method": f"session_{i}_msg_{j}", "id": j}
                event_id = await worker.store_event(session_id, msg)
                events.append(event_id)
            session_events[session_id] = events

        # Verify each session has only its own events
        for session_id, event_ids in session_events.items():
            store = RedisEventStore(max_events_per_stream=100, ttl=300, key_prefix=redis_cleanup)
            replayed = []

            async def callback(msg):
                replayed.append(msg)

            result = await store.replay_events_after(event_ids[0], callback)

            assert result == session_id
            assert len(replayed) == 4  # Events after first

            # Verify no cross-contamination
            for msg in replayed:
                # Extract session index from method name
                session_idx = int(msg["method"].split("_")[1])
                assert sessions[session_idx] == session_id

    async def test_high_throughput(self, redis_cleanup):
        """
        Test event store under high throughput scenario.
        """
        stream_id = "high-throughput-stream"
        num_workers = 10
        events_per_worker = 50

        async def worker_task(worker_id):
            """Worker that rapidly stores events."""
            store = RedisEventStore(max_events_per_stream=1000, ttl=300, key_prefix=redis_cleanup)
            for i in range(events_per_worker):
                msg = {"jsonrpc": "2.0", "method": f"w{worker_id}_msg{i}", "id": f"{worker_id}-{i}"}
                await store.store_event(stream_id, msg)

        # Start all workers simultaneously
        await asyncio.gather(*[worker_task(i) for i in range(num_workers)])

        # Verify total event count
        store = RedisEventStore(max_events_per_stream=1000, ttl=300, key_prefix=redis_cleanup)
        redis = await get_redis_client()

        # Get event count
        count = await redis.hget(store._get_stream_meta_key(stream_id), "count")
        assert count is not None
        total_events = int(count)

        # Should have all events (or max if eviction occurred)
        expected_total = num_workers * events_per_worker
        assert total_events == min(expected_total, 1000)
