import pytest  # noqa F401
from dotenv import load_dotenv

from neollm import MyLLM
from neollm.types import Messages, Response
from neollm.utils.inference import execute_parallel

env_file_path = "project/.env"
assert load_dotenv(env_file_path, override=True)


class SampleMyLLM(MyLLM[str, str]):
    def _preprocess(self, inputs: str) -> Messages:
        messages: Messages = [
            {"role": "system", "content": ("1+1={int}")},
            {"role": "user", "content": str(inputs) + "\n{int}"},
        ]
        return messages

    def _postprocess(self, response: Response) -> str:
        return response.choices[0].message.content or ""


class TestMyLLM:
    def call_myllm(self, model: str, platform: str) -> None:
        myllm = SampleMyLLM(model=model, platform=platform, llm_settings={"max_tokens": 1})
        myllm("1+1")

    def call_stream_myllm(self, model: str, platform: str) -> None:
        myllm = SampleMyLLM(model=model, platform=platform)
        myllm.call_stream("1+1")

    def test_azure(self) -> None:
        kwargs = [
            {"model": model, "platform": "azure"}
            for model in [
                # "gpt-4o-2024-05-13", # 2024/05/15時点、AOAIで使用不可（リージョン不明）
                # "gpt-4-turbo-2024-04-09", # スウェーデンにない
                # "gpt-3.5-turbo-0125",# スウェーデンにない
                # "gpt-4-turbo-0125",# スウェーデンにない
                "gpt-3.5-turbo-1106",
                "gpt-4-turbo-1106",
                # "gpt-4v-turbo-1106", # スウェーデンにない
                "gpt-3.5-turbo-0613",
                "gpt-3.5-turbo-16k-0613",
                "gpt-4-0613",
                "gpt-4-32k-0613",
                # w/o date
                "gpt-4-turbo",
                "gpt-3.5-turbo",
                # "gpt-4v-turbo", # スウェーデンにない
                "gpt-3.5-turbo-16k",
                "gpt-4",
                "gpt-4-32k",
            ]
        ]
        execute_parallel(self.call_myllm, kwargs, max_workers=10)
        execute_parallel(self.call_stream_myllm, kwargs, max_workers=10)

    def test_openai(self) -> None:
        kwargs = [
            {"model": model, "platform": "openai"}
            for model in [
                "gpt-4o-2024-05-13",
                "gpt-4-turbo-2024-04-09",
                "gpt-3.5-turbo-0125",
                "gpt-4-turbo-0125",
                "gpt-3.5-turbo-1106",
                "gpt-4-turbo-1106",
                "gpt-4v-turbo-1106",
                "gpt-3.5-turbo-0613",
                "gpt-3.5-turbo-16k-0613",
                "gpt-4-0613",
                # "gpt-4-32k-0613", #　使えなくなってる
                # w/o date
                "gpt-4-turbo",
                "gpt-3.5-turbo",
                "gpt-4v-turbo",
                "gpt-3.5-turbo-16k",
                "gpt-4",
                # "gpt-4-32k", #　使えなくなってる
            ]
        ]
        execute_parallel(self.call_myllm, kwargs, max_workers=10)
        execute_parallel(self.call_stream_myllm, kwargs, max_workers=10)

    def test_anthropic(self) -> None:
        kwargs = [
            {"model": model, "platform": "anthropic"}
            for model in [
                "claude-3-opus",
                "claude-3-sonnet",
                "claude-3-haiku",
                "claude-3-opus-20240229",
                "claude-3-sonnet-20240229",
                "claude-3-haiku-20240307",
            ]
        ]
        execute_parallel(self.call_myllm, kwargs, max_workers=10)
        execute_parallel(self.call_stream_myllm, kwargs, max_workers=10)

    # TODO! google
    def _test_gcp(self) -> None:
        kwargs = [
            {"model": model, "platform": "gcp"}
            for model in [
                "claude-3-opus",
                "claude-3-sonnet",
                "claude-3-haiku",
                "claude-3-opus@20240229",
                "claude-3-sonnet@20240229",
                "claude-3-haiku@20240307",
            ]
        ]
        execute_parallel(self.call_myllm, kwargs, max_workers=10)
        execute_parallel(self.call_stream_myllm, kwargs, max_workers=10)