"""Unit tests for thread safety in cache_helper module."""

import unittest
import threading
import time
from suno_recs.worker.retrievals.cache_helper import (
    try_cache_get,
    try_cache_put,
)
from suno_recs.worker import constants


class MockRedisClient:
    """Mock Redis client that simulates real Redis behavior."""

    def __init__(self, name):
        self.name = name
        self.data = {}
        self.call_count = {"get": 0, "setex": 0}

    def get(self, key):
        """Get value from cache. Returns bytes like real Redis."""
        self.call_count["get"] += 1
        value = self.data.get(key)
        if value:
            return value.encode("utf-8")
        return None

    def setex(self, key, ttl, value):
        """Set value with TTL."""
        self.call_count["setex"] += 1
        self.data[key] = value
        return True


class TestCacheHelperThreadSafety(unittest.TestCase):
    """Test thread safety of cache_helper module."""

    def setUp(self):
        """Set up test fixtures."""
        # Enable caching for test buckets
        self.original_cache_config = constants.CACHE_CONFIG.copy()
        for i in range(20):
            constants.CACHE_CONFIG[f"test_bucket_{i}"] = {"enabled": True, "ttl_minutes": 60}

    def tearDown(self):
        """Clean up after tests."""
        # Restore original cache config
        constants.CACHE_CONFIG.clear()
        constants.CACHE_CONFIG.update(self.original_cache_config)

    def test_concurrent_client_operations(self):
        """Test that multiple threads can safely use different Redis clients."""
        results = []
        errors = []

        def worker(thread_id, redis_client):
            """Worker function that uses client and performs operations."""
            try:
                # Small delay to increase chance of race conditions
                time.sleep(0.001)

                # Try to write some data using the provided client
                test_data = [{"id": f"thread_{thread_id}", "value": thread_id}]
                try_cache_put(f"test_bucket_{thread_id}", test_data, redis_client=redis_client)

                # Try to read it back
                cached = try_cache_get(f"test_bucket_{thread_id}", redis_client=redis_client)

                results.append(
                    {"thread_id": thread_id, "client_name": redis_client.name, "cached_data": cached}
                )

            except Exception as e:
                errors.append({"thread_id": thread_id, "error": str(e)})

        # Create multiple Redis clients
        clients = [MockRedisClient(f"client_{i}") for i in range(5)]

        # Create and start threads
        threads = []
        for i in range(20):  # 20 threads using 5 different clients
            client = clients[i % 5]
            thread = threading.Thread(target=worker, args=(i, client))
            threads.append(thread)
            thread.start()

        # Wait for all threads to complete
        for thread in threads:
            thread.join()

        # Verify no errors occurred
        self.assertEqual(len(errors), 0, f"Thread safety test failed with errors: {errors}")
        self.assertEqual(len(results), 20, f"Expected 20 results, got {len(results)}")

        # Verify each thread got its own data back
        for result in results:
            thread_id = result["thread_id"]
            cached_data = result["cached_data"]
            self.assertIsNotNone(cached_data, f"Thread {thread_id} got None cached data")
            self.assertEqual(len(cached_data), 1, f"Thread {thread_id} got wrong data length")
            self.assertEqual(
                cached_data[0]["id"], f"thread_{thread_id}", f"Thread {thread_id} got wrong data"
            )

    def test_concurrent_operations_same_client(self):
        """Test concurrent operations using the same client."""
        # This test verifies thread safety by checking that operations
        # complete successfully even under high concurrency

        shared_client = MockRedisClient("shared_client")
        operation_count = 0

        def rapid_operations():
            """Perform rapid Redis operations."""
            nonlocal operation_count
            for i in range(50):
                # Use different clients for different operations
                client = shared_client if i % 2 == 0 else MockRedisClient(f"temp_client_{i}")

                # Perform cache operations
                try_cache_put(f"test_bucket_{i % 5}", [{"id": i}], redis_client=client)
                result = try_cache_get(f"test_bucket_{i % 5}", redis_client=client)

                if result:
                    operation_count += 1

        # Run multiple threads doing rapid operations
        threads = []
        for _ in range(10):
            thread = threading.Thread(target=rapid_operations)
            threads.append(thread)
            thread.start()

        # Wait for completion
        for thread in threads:
            thread.join()

        # If lock is working properly, we should have completed many operations
        # without any crashes or data corruption
        self.assertGreater(operation_count, 0, "Should have completed some operations successfully")

    def test_no_client_provided(self):
        """Test behavior when no Redis client is provided."""
        # Operations should gracefully handle missing client
        result = try_cache_get("test_bucket_0", redis_client=None)
        self.assertIsNone(result, "Should return None when no client is provided")

        # Put should also handle gracefully (no exception)
        try_cache_put("test_bucket_0", [{"test": "data"}], redis_client=None)  # Should not raise


if __name__ == "__main__":
    unittest.main()