from core.entities.provider_entities import ModelSettings
from core.model_runtime.entities.model_entities import ModelType
from core.model_runtime.model_providers import model_provider_factory
from core.provider_manager import ProviderManager
from models.provider import LoadBalancingModelConfig, ProviderModelSetting


def test__to_model_settings(mocker):
    # Get all provider entities
    provider_entities = model_provider_factory.get_providers()

    provider_entity = None
    for provider in provider_entities:
        if provider.provider == "openai":
            provider_entity = provider

    # Mocking the inputs
    provider_model_settings = [
        ProviderModelSetting(
            id="id",
            tenant_id="tenant_id",
            provider_name="openai",
            model_name="gpt-4",
            model_type="text-generation",
            enabled=True,
            load_balancing_enabled=True,
        )
    ]
    load_balancing_model_configs = [
        LoadBalancingModelConfig(
            id="id1",
            tenant_id="tenant_id",
            provider_name="openai",
            model_name="gpt-4",
            model_type="text-generation",
            name="__inherit__",
            encrypted_config=None,
            enabled=True,
        ),
        LoadBalancingModelConfig(
            id="id2",
            tenant_id="tenant_id",
            provider_name="openai",
            model_name="gpt-4",
            model_type="text-generation",
            name="first",
            encrypted_config='{"openai_api_key": "fake_key"}',
            enabled=True,
        ),
    ]

    mocker.patch(
        "core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}
    )

    provider_manager = ProviderManager()

    # Running the method
    result = provider_manager._to_model_settings(provider_entity, provider_model_settings, load_balancing_model_configs)

    # Asserting that the result is as expected
    assert len(result) == 1
    assert isinstance(result[0], ModelSettings)
    assert result[0].model == "gpt-4"
    assert result[0].model_type == ModelType.LLM
    assert result[0].enabled is True
    assert len(result[0].load_balancing_configs) == 2
    assert result[0].load_balancing_configs[0].name == "__inherit__"
    assert result[0].load_balancing_configs[1].name == "first"


def test__to_model_settings_only_one_lb(mocker):
    # Get all provider entities
    provider_entities = model_provider_factory.get_providers()

    provider_entity = None
    for provider in provider_entities:
        if provider.provider == "openai":
            provider_entity = provider

    # Mocking the inputs
    provider_model_settings = [
        ProviderModelSetting(
            id="id",
            tenant_id="tenant_id",
            provider_name="openai",
            model_name="gpt-4",
            model_type="text-generation",
            enabled=True,
            load_balancing_enabled=True,
        )
    ]
    load_balancing_model_configs = [
        LoadBalancingModelConfig(
            id="id1",
            tenant_id="tenant_id",
            provider_name="openai",
            model_name="gpt-4",
            model_type="text-generation",
            name="__inherit__",
            encrypted_config=None,
            enabled=True,
        )
    ]

    mocker.patch(
        "core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}
    )

    provider_manager = ProviderManager()

    # Running the method
    result = provider_manager._to_model_settings(provider_entity, provider_model_settings, load_balancing_model_configs)

    # Asserting that the result is as expected
    assert len(result) == 1
    assert isinstance(result[0], ModelSettings)
    assert result[0].model == "gpt-4"
    assert result[0].model_type == ModelType.LLM
    assert result[0].enabled is True
    assert len(result[0].load_balancing_configs) == 0


def test__to_model_settings_lb_disabled(mocker):
    # Get all provider entities
    provider_entities = model_provider_factory.get_providers()

    provider_entity = None
    for provider in provider_entities:
        if provider.provider == "openai":
            provider_entity = provider

    # Mocking the inputs
    provider_model_settings = [
        ProviderModelSetting(
            id="id",
            tenant_id="tenant_id",
            provider_name="openai",
            model_name="gpt-4",
            model_type="text-generation",
            enabled=True,
            load_balancing_enabled=False,
        )
    ]
    load_balancing_model_configs = [
        LoadBalancingModelConfig(
            id="id1",
            tenant_id="tenant_id",
            provider_name="openai",
            model_name="gpt-4",
            model_type="text-generation",
            name="__inherit__",
            encrypted_config=None,
            enabled=True,
        ),
        LoadBalancingModelConfig(
            id="id2",
            tenant_id="tenant_id",
            provider_name="openai",
            model_name="gpt-4",
            model_type="text-generation",
            name="first",
            encrypted_config='{"openai_api_key": "fake_key"}',
            enabled=True,
        ),
    ]

    mocker.patch(
        "core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}
    )

    provider_manager = ProviderManager()

    # Running the method
    result = provider_manager._to_model_settings(provider_entity, provider_model_settings, load_balancing_model_configs)

    # Asserting that the result is as expected
    assert len(result) == 1
    assert isinstance(result[0], ModelSettings)
    assert result[0].model == "gpt-4"
    assert result[0].model_type == ModelType.LLM
    assert result[0].enabled is True
    assert len(result[0].load_balancing_configs) == 0