import threading
import numpy as np
import pytest
import uuid
from typing import Dict

from chromadb.segment.impl.manager.cache.cache import SegmentLRUCache
from chromadb.types import Segment, SegmentScope


def new_segment(collection_id: uuid.UUID) -> Segment:
    return Segment(
        id=uuid.uuid4(),
        type="test",
        scope=SegmentScope.VECTOR,
        collection=collection_id,
        metadata=None,
        file_paths={},
    )


@pytest.fixture
def cache_setup():
    def _get_segment_disk_size(key: uuid.UUID) -> int:
        return np.random.randint(1, 100)

    def callback_cache_evict(segment: Segment) -> None:
        pass

    class SetupData:
        def __init__(self):
            self.cache = SegmentLRUCache(
                capacity=1000,
                callback=lambda k, v: callback_cache_evict(v),
                size_func=lambda k: _get_segment_disk_size(k),
            )
            self.iterations = 10000
            self.num_threads = 50
            self.errors: Dict[str, int] = {
                "concurrency_error": 0
            }
            self.lock = threading.Lock()

    return SetupData()


def test_thread_safety(cache_setup):
    """Test that demonstrates thread safety issues in the LRU cache"""

    def worker():
        """Worker that performs multiple cache operations"""
        iter =0
        try:
            while cache_setup.errors["concurrency_error"] <=0:
                iter += 1
                key = uuid.uuid4()
                segment = new_segment(key)
                cache_setup.cache.set(key, segment)
                if iter >= cache_setup.iterations:
                    print(f"Stopping due to max iterations: {iter} reached")
                    break

        except Exception as e:
            if "dictionary changed size during iteration" in str(e):
                with cache_setup.lock:
                    cache_setup.errors["concurrency_error"] += 1

    # Create and start threads
    threads = []
    for _ in range(cache_setup.num_threads):
        t = threading.Thread(target=worker)
        threads.append(t)
        t.start()

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

    # Assert that we found thread safety issues
    assert cache_setup.errors["concurrency_error"] > 0, "Expected to find thread safety issues but none were detected"