import json
from unittest import mock

import pytest

from mlflow.exceptions import MlflowException
from mlflow.utils.providers import (
    _fetch_remote_provider,
    _flatten_catalog_entry,
    _get_remote_cache,
    _list_provider_names,
    _load_bundled_provider,
    _load_provider,
    _normalize_provider,
    cost_per_token,
    get_all_providers,
    get_models,
    get_provider_config_response,
)


def test_normalize_provider_normalizes_vertex_ai_variants():
    assert _normalize_provider("vertex_ai") == "vertex_ai"
    assert _normalize_provider("vertex_ai-anthropic") == "vertex_ai"
    assert _normalize_provider("vertex_ai-llama_models") == "vertex_ai"
    assert _normalize_provider("vertex_ai-mistral") == "vertex_ai"


def test_normalize_provider_does_not_normalize_other_providers():
    assert _normalize_provider("openai") == "openai"
    assert _normalize_provider("anthropic") == "anthropic"
    assert _normalize_provider("bedrock") == "bedrock"
    assert _normalize_provider("gemini") == "gemini"


def test_list_provider_names_returns_bundled_providers():
    _list_provider_names.cache_clear()
    providers = _list_provider_names()
    assert len(providers) > 0
    assert "openai" in providers
    assert "anthropic" in providers
    assert "bedrock" in providers


def test_list_provider_names_excludes_non_json():
    _list_provider_names.cache_clear()
    providers = _list_provider_names()
    # __init__.py should not appear
    assert "__init__" not in providers
    for p in providers:
        assert not p.endswith(".py")


def test_load_provider_returns_models(monkeypatch):
    monkeypatch.setenv("MLFLOW_MODEL_CATALOG_URI", "")
    _load_bundled_provider.cache_clear()
    models = _load_provider("openai")
    assert len(models) > 0
    assert "gpt-4o" in models
    info = models["gpt-4o"]
    assert info["mode"] == "chat"
    assert "input_cost_per_token" in info
    assert info["input_cost_per_token"] > 0


def test_load_provider_returns_empty_for_unknown(monkeypatch):
    monkeypatch.setenv("MLFLOW_MODEL_CATALOG_URI", "")
    _load_bundled_provider.cache_clear()
    assert _load_provider("nonexistent_provider_xyz") == {}


def test_load_provider_flattens_pricing(monkeypatch):
    monkeypatch.setenv("MLFLOW_MODEL_CATALOG_URI", "")
    _load_bundled_provider.cache_clear()
    models = _load_provider("anthropic")
    model = next(iter(models.values()))
    # Should have flat ModelInfo keys, not nested pricing/capabilities
    assert "input_cost_per_token" in model or "mode" in model
    assert "pricing" not in model
    assert "context_window" not in model


def _mock_catalog(provider_data):
    """Context manager that mocks the per-provider catalog with the given data.

    ``provider_data`` is a dict mapping provider names to ``{model_name: info}`` dicts.
    """
    return (
        mock.patch(
            "mlflow.utils.providers._load_provider",
            side_effect=lambda p: provider_data.get(p, {}),
        ),
        mock.patch(
            "mlflow.utils.providers._list_provider_names",
            return_value=list(provider_data.keys()),
        ),
    )


def test_get_all_providers_consolidates_vertex_ai_variants():
    data = {
        "openai": {"gpt-4o": {"mode": "chat"}},
        "anthropic": {"claude-3-5-sonnet": {"mode": "chat"}},
        "vertex_ai": {"gemini-1.5-pro": {"mode": "chat"}},
        "vertex_ai-llama_models": {"meta/llama-4-scout": {"mode": "chat"}},
        "vertex_ai-anthropic": {"claude-3-5-sonnet": {"mode": "chat"}},
    }
    with _mock_catalog(data)[0], _mock_catalog(data)[1]:
        providers = get_all_providers()

        assert "vertex_ai" in providers
        assert "vertex_ai-llama_models" not in providers
        assert "vertex_ai-anthropic" not in providers
        assert "openai" in providers
        assert "anthropic" in providers


def test_get_models_normalizes_vertex_ai_provider_and_strips_prefix():
    data = {
        "vertex_ai-llama_models": {
            "meta/llama-4-scout-17b-16e-instruct-maas": {
                "mode": "chat",
                "supports_function_calling": True,
            }
        },
        "vertex_ai-anthropic": {
            "claude-3-5-sonnet": {"mode": "chat", "supports_function_calling": True}
        },
        "vertex_ai": {"gemini-1.5-pro": {"mode": "chat", "supports_function_calling": True}},
    }
    with _mock_catalog(data)[0], _mock_catalog(data)[1]:
        models = get_models(provider="vertex_ai")

        assert len(models) == 3
        for model in models:
            assert model["provider"] == "vertex_ai"

        model_names = [m["model"] for m in models]
        assert "meta/llama-4-scout-17b-16e-instruct-maas" in model_names
        assert "claude-3-5-sonnet" in model_names
        assert "gemini-1.5-pro" in model_names


def test_get_models_filters_by_consolidated_provider():
    data = {
        "openai": {"gpt-4o": {"mode": "chat"}},
        "vertex_ai-llama_models": {"meta/llama-4-scout": {"mode": "chat"}},
    }
    with _mock_catalog(data)[0], _mock_catalog(data)[1]:
        vertex_models = get_models(provider="vertex_ai")
        assert len(vertex_models) == 1
        assert vertex_models[0]["model"] == "meta/llama-4-scout"

        openai_models = get_models(provider="openai")
        assert len(openai_models) == 1
        assert openai_models[0]["model"] == "gpt-4o"


def test_get_models_does_not_modify_other_providers():
    data = {
        "openai": {"gpt-4o": {"mode": "chat", "supports_function_calling": True}},
        "anthropic": {"claude-3-5-sonnet": {"mode": "chat", "supports_function_calling": True}},
    }
    with _mock_catalog(data)[0], _mock_catalog(data)[1]:
        models = get_models()

        openai_model = next(m for m in models if m["provider"] == "openai")
        assert openai_model["model"] == "gpt-4o"

        anthropic_model = next(m for m in models if m["provider"] == "anthropic")
        assert anthropic_model["model"] == "claude-3-5-sonnet"


def test_get_models_dedupes_models_after_normalization():
    data = {
        "vertex_ai": {
            "gemini-3-flash-preview": {"mode": "chat", "supports_function_calling": True}
        },
        "vertex_ai-chat-models": {
            "gemini-3-flash-preview": {"mode": "chat", "supports_function_calling": True}
        },
    }
    with _mock_catalog(data)[0], _mock_catalog(data)[1]:
        models = get_models(provider="vertex_ai")

        model_names = [m["model"] for m in models]
        assert model_names.count("gemini-3-flash-preview") == 1
        assert len(models) == 1


def test_get_all_providers_with_allowed_filter(monkeypatch):
    data = {
        "openai": {"gpt-4o": {"mode": "chat"}},
        "anthropic": {"claude-3-5-sonnet": {"mode": "chat"}},
        "gemini": {"gemini-1.5-pro": {"mode": "chat"}},
    }
    with _mock_catalog(data)[0], _mock_catalog(data)[1]:
        monkeypatch.setenv("MLFLOW_GATEWAY_ALLOWED_PROVIDERS", "openai,anthropic")
        providers = get_all_providers()
        assert "openai" in providers
        assert "anthropic" in providers
        assert "gemini" not in providers


def test_get_models_filters_with_allowed_providers(monkeypatch):
    data = {
        "openai": {"gpt-4o": {"mode": "chat", "supports_function_calling": True}},
        "anthropic": {"claude-3-5-sonnet": {"mode": "chat", "supports_function_calling": True}},
        "gemini": {"gemini-1.5-pro": {"mode": "chat", "supports_function_calling": True}},
    }
    with _mock_catalog(data)[0], _mock_catalog(data)[1]:
        monkeypatch.setenv("MLFLOW_GATEWAY_ALLOWED_PROVIDERS", "openai")
        models = get_models()
        providers_in_result = {m["provider"] for m in models}
        assert providers_in_result == {"openai"}


def test_get_provider_config_rejects_provider_not_in_allowed_list(monkeypatch):
    monkeypatch.setenv("MLFLOW_GATEWAY_ALLOWED_PROVIDERS", "anthropic")
    with pytest.raises(MlflowException, match="not allowed"):
        get_provider_config_response("openai")


def test_get_provider_config_bedrock_has_default_chain():
    config = get_provider_config_response("bedrock")
    modes = {m["mode"] for m in config["auth_modes"]}
    assert "default_chain" in modes

    default_chain = next(m for m in config["auth_modes"] if m["mode"] == "default_chain")
    assert default_chain["display_name"] == "Default Credential Chain"
    assert default_chain["secret_fields"] == []
    assert all(not f["required"] for f in default_chain["config_fields"])


def test_get_provider_config_sagemaker_has_default_chain():
    config = get_provider_config_response("sagemaker")
    modes = {m["mode"] for m in config["auth_modes"]}
    assert "default_chain" in modes


def test_get_provider_config_vertex_ai_has_default_chain():
    config = get_provider_config_response("vertex_ai")
    modes = {m["mode"] for m in config["auth_modes"]}
    assert "default_chain" in modes

    default_chain = next(m for m in config["auth_modes"] if m["mode"] == "default_chain")
    assert default_chain["display_name"] == "Application Default Credentials"
    assert default_chain["secret_fields"] == []
    project_field = next(f for f in default_chain["config_fields"] if f["name"] == "vertex_project")
    assert project_field["required"] is True


_MOCK_PROVIDER_DATA = {
    "test_provider": {
        "test-model": {
            "input_cost_per_token": 1e-6,
            "output_cost_per_token": 2e-6,
            "cache_read_input_token_cost": 5e-7,
            "cache_creation_input_token_cost": 3e-6,
        },
    },
    "openai": {
        "test-provider-model": {
            "input_cost_per_token": 1e-6,
            "output_cost_per_token": 2e-6,
        },
    },
}


def _mock_load_provider(provider):
    return _MOCK_PROVIDER_DATA.get(provider, {})


@pytest.fixture
def mock_model_cost():
    with (
        mock.patch(
            "mlflow.utils.providers._load_provider", side_effect=_mock_load_provider
        ) as m_load,
        mock.patch(
            "mlflow.utils.providers._load_bundled_provider", side_effect=_mock_load_provider
        ),
        mock.patch(
            "mlflow.utils.providers._list_provider_names",
            return_value=list(_MOCK_PROVIDER_DATA.keys()),
        ),
    ):
        yield m_load


def test_cost_per_token_basic(mock_model_cost):
    input_cost, output_cost = cost_per_token(
        model="test-model", prompt_tokens=1000, completion_tokens=500
    )
    # input: 1000 * 1e-6 = 0.001, output: 500 * 2e-6 = 0.001
    assert input_cost == pytest.approx(0.001)
    assert output_cost == pytest.approx(0.001)


def test_cost_per_token_with_provider_prefix(mock_model_cost):
    # "test-provider-model" only exists under "openai/" prefix, so provider lookup is exercised
    input_cost, output_cost = cost_per_token(
        model="test-provider-model",
        prompt_tokens=1000,
        completion_tokens=500,
        custom_llm_provider="openai",
    )
    assert input_cost == pytest.approx(0.001)
    assert output_cost == pytest.approx(0.001)


def test_cost_per_token_strips_provider_prefix(mock_model_cost):
    # "openai/test-model" should resolve to "test-model" by stripping the prefix
    input_cost, output_cost = cost_per_token(
        model="openai/test-model", prompt_tokens=1000, completion_tokens=500
    )
    assert input_cost == pytest.approx(0.001)
    assert output_cost == pytest.approx(0.001)


def test_cost_per_token_cache_read_tokens(mock_model_cost):
    input_cost, output_cost = cost_per_token(
        model="test-model",
        prompt_tokens=1000,
        completion_tokens=500,
        cache_read_input_tokens=200,
    )
    # regular: (1000-200) * 1e-6 = 0.0008
    # cache_read: 200 * 5e-7 = 0.0001
    assert input_cost == pytest.approx(0.0009)
    assert output_cost == pytest.approx(0.001)


def test_cost_per_token_cache_creation_tokens(mock_model_cost):
    input_cost, output_cost = cost_per_token(
        model="test-model",
        prompt_tokens=1000,
        completion_tokens=500,
        cache_creation_input_tokens=300,
    )
    # regular: (1000-300) * 1e-6 = 0.0007
    # cache_creation: 300 * 3e-6 = 0.0009
    assert input_cost == pytest.approx(0.0016)
    assert output_cost == pytest.approx(0.001)


def test_cost_per_token_zero_tokens(mock_model_cost):
    input_cost, output_cost = cost_per_token(
        model="test-model", prompt_tokens=0, completion_tokens=0
    )
    assert input_cost == 0.0
    assert output_cost == 0.0


def test_cost_per_token_unknown_model_returns_none(mock_model_cost):
    assert cost_per_token(model="totally-unknown-model", prompt_tokens=100) is None


def test_cost_per_token_unknown_model_with_provider_returns_none(mock_model_cost):
    assert (
        cost_per_token(
            model="totally-unknown-model",
            prompt_tokens=100,
            custom_llm_provider="unknown-provider",
        )
        is None
    )


def test_cost_per_token_no_cache_cost_falls_back_to_input_rate():
    no_cache_data = {
        "nocache_provider": {
            "test-model": {
                "input_cost_per_token": 1e-6,
                "output_cost_per_token": 2e-6,
            }
        }
    }
    with (
        mock.patch(
            "mlflow.utils.providers._load_provider",
            side_effect=lambda p: no_cache_data.get(p, {}),
        ),
        mock.patch(
            "mlflow.utils.providers._load_bundled_provider",
            side_effect=lambda p: no_cache_data.get(p, {}),
        ),
        mock.patch(
            "mlflow.utils.providers._list_provider_names",
            return_value=list(no_cache_data.keys()),
        ),
    ):
        input_cost, output_cost = cost_per_token(
            model="test-model",
            prompt_tokens=1000,
            completion_tokens=500,
            cache_read_input_tokens=200,
        )
        # No cache_read_input_token_cost, falls back to input_cost_per_token
        # regular: 800 * 1e-6 = 0.0008
        # cache_read: 200 * 1e-6 = 0.0002 (same rate as regular)
        assert input_cost == pytest.approx(0.001)
        assert output_cost == pytest.approx(0.001)


def test_flatten_catalog_entry():
    entry = {
        "mode": "chat",
        "context_window": {"max_input": 128000, "max_output": 16384},
        "pricing": {
            "input_per_million_tokens": 2.5,
            "output_per_million_tokens": 10.0,
            "cache_read_per_million_tokens": 1.25,
            "cache_write_per_million_tokens": 5.0,
        },
        "capabilities": {
            "function_calling": True,
            "vision": True,
            "reasoning": False,
            "prompt_caching": True,
            "response_schema": True,
        },
        "deprecation_date": "2026-01-01",
    }
    info = _flatten_catalog_entry(entry)
    assert info["mode"] == "chat"
    assert info["max_input_tokens"] == 128000
    assert info["max_output_tokens"] == 16384
    assert info["input_cost_per_token"] == pytest.approx(2.5e-6)
    assert info["output_cost_per_token"] == pytest.approx(1e-5)
    assert info["cache_read_input_token_cost"] == pytest.approx(1.25e-6)
    assert info["cache_creation_input_token_cost"] == pytest.approx(5e-6)
    assert info["supports_function_calling"] is True
    assert info["supports_vision"] is True
    assert info["supports_reasoning"] is False
    assert info["deprecation_date"] == "2026-01-01"


def test_flatten_catalog_entry_with_last_updated_at():
    entry = {
        "mode": "chat",
        "capabilities": {
            "function_calling": False,
            "vision": False,
            "reasoning": False,
            "prompt_caching": False,
            "response_schema": False,
        },
        "last_updated_at": "2025-01-15",
    }
    info = _flatten_catalog_entry(entry)
    assert info["last_updated_at"] == "2025-01-15"


def test_flatten_catalog_entry_without_last_updated_at():
    entry = {
        "mode": "chat",
        "capabilities": {
            "function_calling": False,
            "vision": False,
            "reasoning": False,
            "prompt_caching": False,
            "response_schema": False,
        },
    }
    info = _flatten_catalog_entry(entry)
    assert "last_updated_at" not in info


def test_flatten_catalog_entry_with_modality_pricing():
    entry = {
        "mode": "embedding",
        "context_window": {"max_input": 8172, "max_tokens": 8172},
        "pricing": {
            "input_per_million_tokens": 0.135,
            "output_per_million_tokens": 0.0,
            "modality": {
                "image": {"input_per_million_tokens": 60.0},
                "video": {"input_per_second": 0.0007},
                "audio": {"input_per_second": 0.00014},
            },
        },
        "capabilities": {
            "function_calling": False,
            "vision": True,
            "reasoning": False,
            "prompt_caching": False,
            "response_schema": False,
        },
    }
    info = _flatten_catalog_entry(entry)
    assert info["mode"] == "embedding"
    assert info["input_cost_per_token"] == pytest.approx(0.135e-6)
    assert info["modality"] == {
        "image": {"input_per_million_tokens": 60.0},
        "video": {"input_per_second": 0.0007},
        "audio": {"input_per_second": 0.00014},
    }
    assert info["supports_vision"] is True


def test_flatten_catalog_entry_without_modality_pricing():
    entry = {
        "mode": "chat",
        "pricing": {"input_per_million_tokens": 1.0, "output_per_million_tokens": 2.0},
        "capabilities": {
            "function_calling": False,
            "vision": False,
            "reasoning": False,
            "prompt_caching": False,
            "response_schema": False,
        },
    }
    info = _flatten_catalog_entry(entry)
    assert "modality" not in info


def test_load_provider_uses_remote_when_available():
    remote_data = {"test-model": {"mode": "chat", "input_cost_per_token": 1e-6}}
    with mock.patch(
        "mlflow.utils.providers._fetch_remote_provider", return_value=remote_data
    ) as mock_remote:
        result = _load_provider("openai")
        mock_remote.assert_called_once_with("openai")
        assert result is remote_data


def test_load_provider_falls_back_to_bundled_when_remote_fails():
    with mock.patch(
        "mlflow.utils.providers._fetch_remote_provider", return_value=None
    ) as mock_remote:
        result = _load_provider("openai")
        mock_remote.assert_called_once_with("openai")
        assert len(result) > 0
        assert "gpt-4o" in result


def test_fetch_remote_provider_disabled_when_url_empty(monkeypatch):
    monkeypatch.setenv("MLFLOW_MODEL_CATALOG_URI", "")
    assert _fetch_remote_provider("openai") is None


def test_fetch_remote_provider_supports_file_url(tmp_path, monkeypatch):
    catalog = {
        "schema_version": "1.0",
        "models": {"test-model": {"mode": "chat", "pricing": {"input_per_million_tokens": 1.0}}},
    }
    (tmp_path / "test_provider.json").write_text(json.dumps(catalog))
    monkeypatch.setenv("MLFLOW_MODEL_CATALOG_URI", tmp_path.as_uri())
    _get_remote_cache().clear()
    result = _fetch_remote_provider("test_provider")
    assert result is not None
    assert "test-model" in result