from neollm.llm.gpt.azure_llm import (
    AzureGPT4_0613,
    AzureGPT4O_20240513,
    AzureGPT4T_0125,
    AzureGPT4T_1106,
    AzureGPT4T_20240409,
    AzureGPT4VT_1106,
    AzureGPT35FT,
    AzureGPT35T16k_0613,
    AzureGPT35T_0125,
    AzureGPT35T_0613,
    AzureGPT35T_1106,
    AzureGPT432k_0613,
    get_azure_llm,
)
from neollm.types.info import APIPricing


def test_get_azure_llm() -> None:

    # no date
    assert get_azure_llm("gpt-3.5-turbo", {}).__class__ == AzureGPT35T_0613
    assert get_azure_llm("gpt-35-turbo", {}).__class__ == AzureGPT35T_0613
    assert get_azure_llm("gpt-3.5-turbo-16k", {}).__class__ == AzureGPT35T16k_0613
    assert get_azure_llm("gpt-35-turbo-16k", {}).__class__ == AzureGPT35T16k_0613
    assert get_azure_llm("gpt-4", {}).__class__ == AzureGPT4_0613
    assert get_azure_llm("gpt-4-32k", {}).__class__ == AzureGPT432k_0613
    assert get_azure_llm("gpt-4-turbo", {}).__class__ == AzureGPT4T_1106
    assert get_azure_llm("gpt-4v-turbo", {}).__class__ == AzureGPT4VT_1106
    assert get_azure_llm("gpt-4o", {}).__class__ == AzureGPT4O_20240513
    # with date
    assert get_azure_llm("gpt-4o-2024-05-13", {}).__class__ == AzureGPT4O_20240513
    assert get_azure_llm("gpt-4-turbo-2024-04-09", {}).__class__ == AzureGPT4T_20240409
    assert get_azure_llm("gpt-3.5-turbo-0125", {}).__class__ == AzureGPT35T_0125
    assert get_azure_llm("gpt-35-turbo-0125", {}).__class__ == AzureGPT35T_0125
    assert get_azure_llm("gpt-4-turbo-0125", {}).__class__ == AzureGPT4T_0125
    assert get_azure_llm("gpt-3.5-turbo-1106", {}).__class__ == AzureGPT35T_1106
    assert get_azure_llm("gpt-35-turbo-1106", {}).__class__ == AzureGPT35T_1106
    assert get_azure_llm("gpt-4-turbo-1106", {}).__class__ == AzureGPT4T_1106
    assert get_azure_llm("gpt-4v-turbo-1106", {}).__class__ == AzureGPT4VT_1106
    assert get_azure_llm("gpt-3.5-turbo-0613", {}).__class__ == AzureGPT35T_0613
    assert get_azure_llm("gpt-35-turbo-0613", {}).__class__ == AzureGPT35T_0613
    assert get_azure_llm("gpt-3.5-turbo-16k-0613", {}).__class__ == AzureGPT35T16k_0613
    assert get_azure_llm("gpt-35-turbo-16k-0613", {}).__class__ == AzureGPT35T16k_0613
    assert get_azure_llm("gpt-4-0613", {}).__class__ == AzureGPT4_0613
    assert get_azure_llm("gpt-4-32k-0613", {}).__class__ == AzureGPT432k_0613
    # ft
    assert get_azure_llm("ft:gpt-3.5-turbo-1106-XXXX", {}).__class__ == AzureGPT35FT


def test_check_price() -> None:
    # https://azure.microsoft.com/ja-jp/pricing/details/cognitive-services/openai-service/

    # これからのモデル --------------------------------------------------------
    assert AzureGPT4T_20240409.dollar_per_ktoken == APIPricing(input=0.01, output=0.03)
    # Updated --------------------------------------------------------
    # GPT3.5T
    assert AzureGPT35T_0125.dollar_per_ktoken == APIPricing(input=0.0005, output=0.0015)
    # GPT4
    assert AzureGPT4O_20240513.dollar_per_ktoken == APIPricing(input=0.005, output=0.015)
    assert AzureGPT4T_0125.dollar_per_ktoken == APIPricing(input=0.01, output=0.03)
    assert AzureGPT4VT_1106.dollar_per_ktoken == APIPricing(input=0.01, output=0.03)
    assert AzureGPT4T_1106.dollar_per_ktoken == APIPricing(input=0.01, output=0.03)
    assert AzureGPT4_0613.dollar_per_ktoken == APIPricing(input=0.03, output=0.06)
    assert AzureGPT432k_0613.dollar_per_ktoken == APIPricing(input=0.06, output=0.12)
    # FT
    assert AzureGPT35FT.dollar_per_ktoken == APIPricing(input=0.0005, output=0.0015)
    # Legacy ---------------------------------------------------------
    # AzureGPT35T_0301 なし
    assert AzureGPT35T_0613.dollar_per_ktoken == APIPricing(input=0.0015, output=0.002)
    assert AzureGPT35T16k_0613.dollar_per_ktoken == APIPricing(input=0.003, output=0.004)
    assert AzureGPT35T_1106.dollar_per_ktoken == APIPricing(input=0.001, output=0.002)


def test_check_context_window() -> None:
    # https://learn.microsoft.com/ja-jp/azure/ai-services/openai/concepts/models#gpt-4-and-gpt-4-turbo-preview-models
    assert AzureGPT4T_20240409.context_window == 128_000

    assert AzureGPT4T_0125.context_window == 128_000
    assert AzureGPT35T_0125.context_window == 16_385

    assert AzureGPT4O_20240513.context_window == 128_000
    assert AzureGPT4T_1106.context_window == 128_000
    assert AzureGPT4VT_1106.context_window == 128_000
    assert AzureGPT35T_1106.context_window == 16_385

    assert AzureGPT35T_0613.context_window == 4_096
    assert AzureGPT4_0613.context_window == 8_192
    assert AzureGPT35T16k_0613.context_window == 16_385
    assert AzureGPT432k_0613.context_window == 32_768

    assert AzureGPT35FT.context_window == 4_096