import random
import string
from typing import Optional, List, Tuple, Any
from unittest.mock import patch

from chromadb.config import System, Settings
from chromadb.quota import QuotaEnforcer, Resource
import pytest


def generate_random_string(size: int) -> str:
    return ''.join(random.choices(string.ascii_letters + string.digits, k=size))

def mock_get_for_subject(self, resource: Resource, subject: Optional[str] = "", tier: Optional[str] = "") -> Optional[
    int]:
    """Mock function to simulate quota retrieval."""
    return 10


def run_static_checks(enforcer: QuotaEnforcer, test_cases: List[Tuple[Any, Optional[str]]], data_key: str):
    """Generalized function to run static checks on different types of data."""
    for test_case in test_cases:
        data, expected_error = test_case if len(test_case) == 2 else (test_case[0], None)
        args = {data_key: [data]}
        if expected_error:
            with pytest.raises(Exception) as exc_info:
                enforcer.static_check(**args)
            assert expected_error in str(exc_info.value.resource)
        else:
            enforcer.static_check(**args)



@pytest.fixture(scope="module")
def enforcer() -> QuotaEnforcer:
    settings = Settings(
        chroma_quota_provider_impl =  "chromadb.quota.test_provider.QuotaProviderForTest"
    )
    system = System(settings)
    return system.require(QuotaEnforcer)

@patch('chromadb.quota.test_provider.QuotaProviderForTest.get_for_subject', mock_get_for_subject)
def test_static_enforcer_metadata(enforcer):
    test_cases = [
        ({generate_random_string(20): generate_random_string(5)}, "METADATA_KEY_LENGTH"),
        ({generate_random_string(5): generate_random_string(5)}, None),
        ({generate_random_string(5): generate_random_string(20)}, "METADATA_VALUE_LENGTH"),
        ({generate_random_string(5): generate_random_string(5)}, None)
    ]
    run_static_checks(enforcer, test_cases, 'metadatas')


@patch('chromadb.quota.test_provider.QuotaProviderForTest.get_for_subject', mock_get_for_subject)
def test_static_enforcer_documents(enforcer):
    test_cases = [
        (generate_random_string(20), "DOCUMENT_SIZE"),
        (generate_random_string(5), None)
    ]
    run_static_checks(enforcer, test_cases, 'documents')

@patch('chromadb.quota.test_provider.QuotaProviderForTest.get_for_subject', mock_get_for_subject)
def test_static_enforcer_embeddings(enforcer):
    test_cases = [
        (random.sample(range(1, 101), 100), "EMBEDDINGS_DIMENSION"),
        (random.sample(range(1, 101), 5), None)
    ]
    run_static_checks(enforcer, test_cases, 'embeddings')

# Should not raise an error if no quota provider is present
def test_enforcer_without_quota_provider():
    test_cases = [
        (random.sample(range(1, 101), 1), None),
        (random.sample(range(1, 101), 5), None)
    ]
    settings = Settings()
    system = System(settings)
    enforcer = system.require(QuotaEnforcer)
    run_static_checks(enforcer, test_cases, 'embeddings')
