from pathlib import Path
from typing import Callable, Union

import pytest

from weaviate.collections.classes.generative import (
    GenerativeConfig,
    GenerativeParameters,
    _GenerativeConfigRuntimeOptions,
    _GroupedTask,
    _SinglePrompt,
)
from weaviate.proto.v1 import base_pb2
from weaviate.proto.v1 import generative_pb2
from weaviate.types import BLOB_INPUT

LOGO = "test/collection/weaviate-logo.png"
LOGO_ENCODED = " "


@pytest.mark.parametrize(
    "provider",
    [
        lambda image: GenerativeParameters.grouped_task("whatever", images=[image]),
        lambda image: GenerativeParameters.single_prompt("whatever", images=[image]),
        lambda image: GenerativeParameters.grouped_task("whatever", images=image),
        lambda image: GenerativeParameters.single_prompt("whatever", images=image),
    ],
)
def test_generative_parameters_images_parsing(
    provider: Union[
        Callable[[BLOB_INPUT], _GroupedTask],
        Callable[[BLOB_INPUT], _SinglePrompt],
    ],
) -> None:
    # tests image as a path string
    p = provider(LOGO)
    assert p.images is not None
    assert list(p.images) == [LOGO_ENCODED]

    # tests image as a path object
    p = provider(Path(LOGO))
    assert p.images is not None
    assert list(p.images) == [LOGO_ENCODED]

    # tests image as an open file
    with open(LOGO, "rb") as f:
        p = provider(f)
        assert p.images is not None
        assert list(p.images) == [LOGO_ENCODED]


@pytest.mark.parametrize(
    "actual,expected",
    [
        (
            GenerativeConfig.anthropic(
                base_url="http://localhost:8080",
                model="text-to-image",
                max_tokens=100,
                stop_sequences=["\n"],
                temperature=0.5,
                top_k=50,
                top_p=0.9,
            )._to_grpc(
                _GenerativeConfigRuntimeOptions(
                    return_metadata=True, images=[LOGO_ENCODED], image_properties=["image"]
                )
            ),
            generative_pb2.GenerativeProvider(
                return_metadata=True,
                anthropic=generative_pb2.GenerativeAnthropic(
                    base_url="http://localhost:8080",
                    model="text-to-image",
                    max_tokens=100,
                    stop_sequences=base_pb2.TextArray(values=["\n"]),
                    temperature=0.5,
                    top_k=50,
                    top_p=0.9,
                    images=base_pb2.TextArray(values=[LOGO_ENCODED]),
                    image_properties=base_pb2.TextArray(values=["image"]),
                ),
            ),
        ),
        (
            GenerativeConfig.anyscale(
                base_url="http://localhost:8080",
                model="text-to-image",
                temperature=0.5,
            )._to_grpc(_GenerativeConfigRuntimeOptions(return_metadata=True)),
            generative_pb2.GenerativeProvider(
                return_metadata=True,
                anyscale=generative_pb2.GenerativeAnyscale(
                    base_url="http://localhost:8080",
                    model="text-to-image",
                    temperature=0.5,
                ),
            ),
        ),
        (
            GenerativeConfig.aws(
                endpoint="http://localhost:8080",
                model="text-to-image",
                region="us-west-2",
                service="sagemaker",
                target_model="arn:aws:sagemaker:us-west-2:123456789012:model/text-to-image",
                target_variant="variant-1",
                temperature=0.5,
            )._to_grpc(
                _GenerativeConfigRuntimeOptions(
                    return_metadata=True, images=[LOGO_ENCODED], image_properties=["image"]
                )
            ),
            generative_pb2.GenerativeProvider(
                return_metadata=True,
                aws=generative_pb2.GenerativeAWS(
                    endpoint="http://localhost:8080",
                    model="text-to-image",
                    region="us-west-2",
                    service="sagemaker",
                    target_model="arn:aws:sagemaker:us-west-2:123456789012:model/text-to-image",
                    target_variant="variant-1",
                    temperature=0.5,
                    images=base_pb2.TextArray(values=[LOGO_ENCODED]),
                    image_properties=base_pb2.TextArray(values=["image"]),
                ),
            ),
        ),
        (
            GenerativeConfig.cohere(
                base_url="http://localhost:8080",
                k=5,
                max_tokens=100,
                model="text-to-image",
                p=0.9,
                presence_penalty=0.5,
                stop_sequences=["\n"],
                temperature=0.5,
            )._to_grpc(_GenerativeConfigRuntimeOptions(return_metadata=True)),
            generative_pb2.GenerativeProvider(
                return_metadata=True,
                cohere=generative_pb2.GenerativeCohere(
                    base_url="http://localhost:8080",
                    k=5,
                    max_tokens=100,
                    model="text-to-image",
                    p=0.9,
                    presence_penalty=0.5,
                    stop_sequences=base_pb2.TextArray(values=["\n"]),
                    temperature=0.5,
                ),
            ),
        ),
        (
            GenerativeConfig.databricks(
                endpoint="http://localhost:8080",
                frequency_penalty=0.5,
                log_probs=True,
                max_tokens=100,
                model="text-to-image",
                n=5,
                presence_penalty=0.5,
                stop=["\n"],
                temperature=0.5,
                top_log_probs=5,
                top_p=0.9,
            )._to_grpc(_GenerativeConfigRuntimeOptions(return_metadata=True)),
            generative_pb2.GenerativeProvider(
                return_metadata=True,
                databricks=generative_pb2.GenerativeDatabricks(
                    endpoint="http://localhost:8080",
                    frequency_penalty=0.5,
                    log_probs=True,
                    max_tokens=100,
                    model="text-to-image",
                    n=5,
                    presence_penalty=0.5,
                    stop=base_pb2.TextArray(values=["\n"]),
                    temperature=0.5,
                    top_log_probs=5,
                    top_p=0.9,
                ),
            ),
        ),
        (
            GenerativeConfig.dummy()._to_grpc(
                _GenerativeConfigRuntimeOptions(return_metadata=True)
            ),
            generative_pb2.GenerativeProvider(
                return_metadata=True, dummy=generative_pb2.GenerativeDummy()
            ),
        ),
        (
            GenerativeConfig.friendliai(
                base_url="http://localhost:8080",
                max_tokens=100,
                model="text-to-image",
                n=5,
                temperature=0.5,
                top_p=0.9,
            )._to_grpc(_GenerativeConfigRuntimeOptions(return_metadata=True)),
            generative_pb2.GenerativeProvider(
                return_metadata=True,
                friendliai=generative_pb2.GenerativeFriendliAI(
                    base_url="http://localhost:8080",
                    max_tokens=100,
                    model="text-to-image",
                    n=5,
                    temperature=0.5,
                    top_p=0.9,
                ),
            ),
        ),
        (
            GenerativeConfig.google(
                api_endpoint="http://localhost:8080",
                endpoint_id="12345678901234567890123456789012",
                frequency_penalty=0.5,
                max_tokens=100,
                model="text-to-image",
                presence_penalty=0.5,
                project_id="my-project",
                region="us-west1",
                stop_sequences=["\n"],
                temperature=0.5,
                top_k=50,
                top_p=0.9,
            )._to_grpc(
                _GenerativeConfigRuntimeOptions(
                    return_metadata=True, images=[LOGO_ENCODED], image_properties=["image"]
                )
            ),
            generative_pb2.GenerativeProvider(
                return_metadata=True,
                google=generative_pb2.GenerativeGoogle(
                    api_endpoint="localhost:8080",
                    endpoint_id="12345678901234567890123456789012",
                    frequency_penalty=0.5,
                    max_tokens=100,
                    model="text-to-image",
                    presence_penalty=0.5,
                    project_id="my-project",
                    region="us-west1",
                    stop_sequences=base_pb2.TextArray(values=["\n"]),
                    temperature=0.5,
                    top_k=50,
                    top_p=0.9,
                    images=base_pb2.TextArray(values=[LOGO_ENCODED]),
                    image_properties=base_pb2.TextArray(values=["image"]),
                ),
            ),
        ),
        (
            GenerativeConfig.mistral(
                base_url="http://localhost:8080",
                max_tokens=100,
                model="text-to-image",
                temperature=0.5,
                top_p=0.9,
            )._to_grpc(_GenerativeConfigRuntimeOptions(return_metadata=True)),
            generative_pb2.GenerativeProvider(
                return_metadata=True,
                mistral=generative_pb2.GenerativeMistral(
                    base_url="http://localhost:8080",
                    max_tokens=100,
                    model="text-to-image",
                    temperature=0.5,
                    top_p=0.9,
                ),
            ),
        ),
        (
            GenerativeConfig.nvidia(
                base_url="http://localhost:8080",
                max_tokens=100,
                model="text-to-image",
                temperature=0.5,
                top_p=0.9,
            )._to_grpc(_GenerativeConfigRuntimeOptions(return_metadata=True)),
            generative_pb2.GenerativeProvider(
                return_metadata=True,
                nvidia=generative_pb2.GenerativeNvidia(
                    base_url="http://localhost:8080",
                    max_tokens=100,
                    model="text-to-image",
                    temperature=0.5,
                    top_p=0.9,
                ),
            ),
        ),
        (
            GenerativeConfig.ollama(
                api_endpoint="http://localhost:8080",
                model="text-to-image",
                temperature=0.5,
            )._to_grpc(
                _GenerativeConfigRuntimeOptions(
                    return_metadata=True, images=[LOGO_ENCODED], image_properties=["image"]
                )
            ),
            generative_pb2.GenerativeProvider(
                return_metadata=True,
                ollama=generative_pb2.GenerativeOllama(
                    api_endpoint="http://localhost:8080",
                    model="text-to-image",
                    temperature=0.5,
                    images=base_pb2.TextArray(values=[LOGO_ENCODED]),
                    image_properties=base_pb2.TextArray(values=["image"]),
                ),
            ),
        ),
        (
            GenerativeConfig.openai(
                api_version="2020-05-03",
                base_url="http://localhost:8080",
                deployment_id="12345678901234567890123456789012",
                frequency_penalty=0.5,
                max_tokens=100,
                model="text-to-image",
                presence_penalty=0.5,
                resource_name="projects/my-project/locations/us-west1/endpoints/12345678901234567890123456789012",
                stop=["\n"],
                temperature=0.5,
                top_p=50,
                reasoning_effort="high",
                verbosity="low",
            )._to_grpc(
                _GenerativeConfigRuntimeOptions(
                    return_metadata=True, images=[LOGO_ENCODED], image_properties=["image"]
                )
            ),
            generative_pb2.GenerativeProvider(
                return_metadata=True,
                openai=generative_pb2.GenerativeOpenAI(
                    api_version="2020-05-03",
                    base_url="http://localhost:8080",
                    deployment_id="12345678901234567890123456789012",
                    frequency_penalty=0.5,
                    max_tokens=100,
                    model="text-to-image",
                    presence_penalty=0.5,
                    resource_name="projects/my-project/locations/us-west1/endpoints/12345678901234567890123456789012",
                    stop=base_pb2.TextArray(values=["\n"]),
                    temperature=0.5,
                    top_p=50,
                    images=base_pb2.TextArray(values=[LOGO_ENCODED]),
                    image_properties=base_pb2.TextArray(values=["image"]),
                    is_azure=False,
                    reasoning_effort=generative_pb2.GenerativeOpenAI.ReasoningEffort.REASONING_EFFORT_HIGH,
                    verbosity=generative_pb2.GenerativeOpenAI.Verbosity.VERBOSITY_LOW,
                ),
            ),
        ),
        (
            GenerativeConfig.azure_openai(
                api_version="2020-05-03",
                base_url="http://localhost:8080",
                deployment_id="12345678901234567890123456789012",
                frequency_penalty=0.5,
                max_tokens=100,
                model="text-to-image",
                presence_penalty=0.5,
                resource_name="projects/my-project/locations/us-west1/endpoints/12345678901234567890123456789012",
                stop=["\n"],
                temperature=0.5,
                top_p=50,
            )._to_grpc(
                _GenerativeConfigRuntimeOptions(
                    return_metadata=True, images=[LOGO_ENCODED], image_properties=["image"]
                )
            ),
            generative_pb2.GenerativeProvider(
                return_metadata=True,
                openai=generative_pb2.GenerativeOpenAI(
                    api_version="2020-05-03",
                    base_url="http://localhost:8080",
                    deployment_id="12345678901234567890123456789012",
                    frequency_penalty=0.5,
                    max_tokens=100,
                    model="text-to-image",
                    presence_penalty=0.5,
                    resource_name="projects/my-project/locations/us-west1/endpoints/12345678901234567890123456789012",
                    stop=base_pb2.TextArray(values=["\n"]),
                    temperature=0.5,
                    top_p=50,
                    images=base_pb2.TextArray(values=[LOGO_ENCODED]),
                    image_properties=base_pb2.TextArray(values=["image"]),
                    is_azure=True,
                ),
            ),
        ),
        (
            GenerativeConfig.xai(
                base_url="http://localhost:8080",
                max_tokens=100,
                model="text-to-image",
                temperature=0.5,
                top_p=50,
            )._to_grpc(
                _GenerativeConfigRuntimeOptions(
                    return_metadata=True, images=[LOGO_ENCODED], image_properties=["image"]
                )
            ),
            generative_pb2.GenerativeProvider(
                return_metadata=True,
                xai=generative_pb2.GenerativeXAI(
                    base_url="http://localhost:8080",
                    max_tokens=100,
                    model="text-to-image",
                    temperature=0.5,
                    top_p=50,
                    images=base_pb2.TextArray(values=[LOGO_ENCODED]),
                    image_properties=base_pb2.TextArray(values=["image"]),
                ),
            ),
        ),
    ],
)
def test_generative_provider_to_grpc(
    actual: generative_pb2.GenerativeProvider, expected: generative_pb2.GenerativeProvider
) -> None:
    assert expected == actual
