from unittest import mock

from mlflow.genai.judges.adapters.databricks_managed_judge_adapter import (
    DatabricksManagedJudgeAdapter,
)
from mlflow.genai.judges.adapters.databricks_serving_endpoint_adapter import (
    DatabricksServingEndpointAdapter,
)
from mlflow.genai.judges.adapters.gateway_adapter import GatewayAdapter
from mlflow.genai.judges.adapters.litellm_adapter import LiteLLMAdapter
from mlflow.types.llm import ChatMessage

# Test cases: (model_uri, prompt_type, expected_without_litellm, expected_with_litellm)
TEST_CASES = [
    ("databricks", "string", DatabricksManagedJudgeAdapter, DatabricksManagedJudgeAdapter),
    ("databricks", "list", DatabricksManagedJudgeAdapter, DatabricksManagedJudgeAdapter),
    ("databricks:/my-endpoint", "string", DatabricksServingEndpointAdapter, DatabricksServingEndpointAdapter),
    ("databricks:/my-endpoint", "list", DatabricksServingEndpointAdapter, DatabricksServingEndpointAdapter),
    ("endpoints:/my-endpoint", "string", DatabricksServingEndpointAdapter, DatabricksServingEndpointAdapter),
    ("endpoints:/my-endpoint", "list", DatabricksServingEndpointAdapter, DatabricksServingEndpointAdapter),
    ("openai:/gpt-4", "string", GatewayAdapter, LiteLLMAdapter),
    ("openai:/gpt-4", "list", None, LiteLLMAdapter),
    ("anthropic:/claude-3-5-sonnet-20241022", "string", GatewayAdapter, LiteLLMAdapter),
    ("anthropic:/claude-3-5-sonnet-20241022", "list", None, LiteLLMAdapter),
]


def get_applicable_adapter(model_uri: str, prompt):
    adapters = [
        DatabricksManagedJudgeAdapter,
        DatabricksServingEndpointAdapter,
        LiteLLMAdapter,  # Checked BEFORE GatewayAdapter
        GatewayAdapter,
    ]

    for adapter in adapters:
        if adapter.is_applicable(model_uri, prompt):
            return adapter
    return None


def test_scenario(scenario_name: str, litellm_available: bool):
    print(f"\n{'='*70}")
    print(f"{scenario_name}")
    print(f"{'='*70}\n")

    string_prompt = "test input"
    list_prompt = [ChatMessage(role="user", content="test")]

    passed = 0
    failed = 0

    for model_uri, prompt_type, expected_without, expected_with in TEST_CASES:
        expected_adapter = expected_with if litellm_available else expected_without
        prompt = string_prompt if prompt_type == "string" else list_prompt

        with mock.patch(
            "mlflow.genai.judges.adapters.litellm_adapter._is_litellm_available",
            return_value=litellm_available,
        ):
            actual_adapter = get_applicable_adapter(model_uri, prompt)

        status = "✓" if actual_adapter == expected_adapter else "✗"
        expected_name = expected_adapter.__name__ if expected_adapter else "None"
        actual_name = actual_adapter.__name__ if actual_adapter else "None"

        if actual_adapter == expected_adapter:
            passed += 1
            print(f"{status} {model_uri:40} {prompt_type:6} -> {actual_name}")
        else:
            failed += 1
            print(
                f"{status} {model_uri:40} {prompt_type:6} -> Expected: {expected_name}, Got: {actual_name}"
            )

    print(f"\nResults: {passed} passed, {failed} failed")
    return failed == 0

success1 = test_scenario("Scenario 1: WITHOUT LiteLLM (fallback to GatewayAdapter)", False)
success2 = test_scenario("Scenario 2: WITH LiteLLM (LiteLLMAdapter handles all)", True)

print(f"\n{'='*70}")
if success1 and success2:
    print("✅ All adapter selection tests passed!")
    return 0
else:
    print("⚠️  Some tests failed!")
    return 1