# Copyright DataStax, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import os
import sys
from collections.abc import Iterable
from typing import Any

from astrapy.info import EmbeddingProviderParameter, VectorServiceOptions
from astrapy.settings.defaults import (
    EMBEDDING_HEADER_API_KEY,
    EMBEDDING_HEADER_AWS_ACCESS_ID,
    EMBEDDING_HEADER_AWS_SECRET_ID,
)

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from live_provider_info import live_provider_info

alphanum = set("qwertyuiopasdfghjklzxcvbnmQWERTYUIOPASDFGHJKLZXCVBNM1234567890")

TEST_EXTENDED_VECTORIZE = bool(os.environ.get("TEST_EXTENDED_VECTORIZE"))

DEFAULT_TEST_ASSETS = {
    "samples": [
        (
            "car",
            "My car is fast and red, it's awesome.",
        ),
        (
            "landscape",
            "The landscape shows spiky, rocky mountaintops and a river nearby...",
        ),
        (
            "motorbike",
            "Your motorbike is certainly the best I have seen!",
        ),
    ],
    "probe": {
        "text": "Driving a red car is a joy.",
        "expected": ["car", "motorbike"],
    },
}

CODE_TEST_ASSETS = {
    "samples": [
        (
            "py_loop",
            "import os; for i in range(10): print(i + os.getenv('env'))",
        ),
        (
            "java",
            (
                "if (sortClause == null || sortClause.sortExpressions()."
                "isEmpty()){return Uni.createFrom().item(true);}"
            ),
        ),
        (
            "py_while",
            (
                "import json; val=json.load(open('limit.json')); "
                "v=0; while v<val: pprint(v); v = v + 1"
            ),
        ),
    ],
    "probe": {
        "text": "for(q in [10, 11, 12]): print(f'q={q}')",
        "expected": ["py_loop", "py_while"],
    },
}

TEST_ASSETS_MAP = {
    ("jinaAI", "jina-embeddings-v2-base-code"): CODE_TEST_ASSETS,
    ("voyageAI", "voyage-code-2"): CODE_TEST_ASSETS,
}

USE_INSERT_ONE_MAP: dict[tuple[str, str], bool] = {
    # ("upstageAI", "solar-1-mini-embedding"): True,
}

SECRET_NAME_ROOT_MAP = {
    "azureOpenAI": "AZURE_OPENAI",
    "bedrock": "BEDROCK",
    "cohere": "COHERE",
    "huggingface": "HUGGINGFACE",
    "huggingfaceDedicated": "HUGGINGFACEDED",
    "jinaAI": "JINAAI",
    "mistral": "MISTRAL",
    "nvidia": "NVIDIA",
    "openai": "OPENAI",
    "upstageAI": "UPSTAGE",
    "voyageAI": "VOYAGEAI",
}

# Any "model tag" listed here is skipped no matter what.
# (useful for models already from the findEmbProvs but not stable/working yet).
EXCLUDED_MODEL_QUADRUPLES = {
    # Cohere provider, while returned by FindEmbProv's, is not in scope for Vectorize and is skipped:
    ("cohere", "embed-english-v2.0", "HEADER", "f"),
    ("cohere", "embed-english-v3.0", "HEADER", "f"),
    ("cohere", "embed-multilingual-v3.0", "HEADER", "f"),
    # Pending a testable deployment, HuggingFace Dedicated testing is momentarily suspended:
    ("huggingfaceDedicated", "endpoint-defined-model", "HEADER", "f"),
    ("huggingfaceDedicated", "endpoint-defined-model", "SHARED_SECRET", "f"),
}

# this is a way to suppress/limit certain combinations of
# "full param" testing from the start. If even one of the optional params
# in a test model is PARAM_SKIP_MARKER, the combination is not emitted at all.
PARAM_SKIP_MARKER = "__SKIP_ME__"

PARAMETER_VALUE_MAP: dict[tuple[str, str, str], Any]
FORCE_DIMENSION_MAP: dict[tuple[str, str], Any]
if TEST_EXTENDED_VECTORIZE:
    PARAMETER_VALUE_MAP = {
        ("azureOpenAI", "text-embedding-3-large", "deploymentId"): os.environ[
            "AZURE_OPENAI_DEPLOY_ID_EMB3LARGE"
        ],
        ("azureOpenAI", "text-embedding-3-large", "resourceName"): os.environ[
            "AZURE_OPENAI_RESNAME_EMB3LARGE"
        ],
        ("azureOpenAI", "text-embedding-3-small", "deploymentId"): os.environ[
            "AZURE_OPENAI_DEPLOY_ID_EMB3SMALL"
        ],
        ("azureOpenAI", "text-embedding-3-small", "resourceName"): os.environ[
            "AZURE_OPENAI_RESNAME_EMB3SMALL"
        ],
        ("azureOpenAI", "text-embedding-ada-002", "deploymentId"): os.environ[
            "AZURE_OPENAI_DEPLOY_ID_ADA2"
        ],
        ("azureOpenAI", "text-embedding-ada-002", "resourceName"): os.environ[
            "AZURE_OPENAI_RESNAME_ADA2"
        ],
        ("jinaAI", "jina-embeddings-v3", "late_chunking"): True,
        ("jinaAI", "jina-embeddings-v3", "task"): "text-matching",
        ("voyageAI", "voyage-2", "autoTruncate"): True,
        ("voyageAI", "voyage-code-2", "autoTruncate"): True,
        ("voyageAI", "voyage-finance-2", "autoTruncate"): True,
        ("voyageAI", "voyage-large-2", "autoTruncate"): True,
        ("voyageAI", "voyage-large-2-instruct", "autoTruncate"): True,
        ("voyageAI", "voyage-law-2", "autoTruncate"): True,
        ("voyageAI", "voyage-multilingual-2", "autoTruncate"): True,
        #
        ("huggingfaceDedicated", "endpoint-defined-model", "endpointName"): os.environ[
            "HUGGINGFACEDED_ENDPOINTNAME"
        ],
        ("huggingfaceDedicated", "endpoint-defined-model", "regionName"): os.environ[
            "HUGGINGFACEDED_REGIONNAME"
        ],
        ("huggingfaceDedicated", "endpoint-defined-model", "cloudName"): os.environ[
            "HUGGINGFACEDED_CLOUDNAME"
        ],
        #
        ("openai", "text-embedding-3-large", "organizationId"): os.environ[
            "OPENAI_ORGANIZATION_ID"
        ],
        ("openai", "text-embedding-3-large", "projectId"): os.environ[
            "OPENAI_PROJECT_ID"
        ],
        ("openai", "text-embedding-3-small", "organizationId"): os.environ[
            "OPENAI_ORGANIZATION_ID"
        ],
        ("openai", "text-embedding-3-small", "projectId"): os.environ[
            "OPENAI_PROJECT_ID"
        ],
        ("openai", "text-embedding-ada-002", "organizationId"): os.environ[
            "OPENAI_ORGANIZATION_ID"
        ],
        ("openai", "text-embedding-ada-002", "projectId"): os.environ[
            "OPENAI_PROJECT_ID"
        ],
        #
        ("bedrock", "amazon.titan-embed-text-v1", "region"): os.environ[
            "BEDROCK_REGION"
        ],
        ("bedrock", "amazon.titan-embed-text-v2:0", "region"): os.environ[
            "BEDROCK_REGION"
        ],
    }

    # this is ad-hoc for HF dedicated. Models here, though "optional" dimension,
    # do not undergo the f/0 optional dimension because of that, rather have
    # a forced fixed, provided dimension.
    FORCE_DIMENSION_MAP = {
        ("huggingfaceDedicated", "endpoint-defined-model"): int(
            os.environ["HUGGINGFACEDED_DIMENSION"]
        ),
    }
else:
    PARAMETER_VALUE_MAP = {}
    FORCE_DIMENSION_MAP = {}


def live_test_models() -> Iterable[dict[str, Any]]:
    def _from_validation(pspec: EmbeddingProviderParameter) -> int:
        assert pspec.parameter_type.upper() == "NUMBER"
        if "numericRange" in pspec.validation:
            m0: int = pspec.validation["numericRange"][0]
            m1: int = pspec.validation["numericRange"][1]
            return (m0 + m1) // 2
        elif "options" in pspec.validation:
            options: list[int] = pspec.validation["options"]
            if len(options) > 1:
                return options[1]
            else:
                return options[0]
        else:
            raise ValueError("unsupported pspec")

    def _collapse(longt: str) -> str:
        if len(longt) <= 40:
            return longt
        else:
            return f"{longt[:30]}_{longt[-5:]}"

    if TEST_EXTENDED_VECTORIZE:
        # generate the full list of models based on the live provider endpoint
        provider_info = live_provider_info()
        for provider_name, provider_desc in sorted(
            provider_info.embedding_providers.items()
        ):
            for model in provider_desc.models:
                for auth_type_name, auth_type_desc in sorted(
                    provider_desc.supported_authentication.items()
                ):
                    if auth_type_desc.enabled:
                        # test assumptions on auth type
                        if auth_type_name == "NONE":
                            assert auth_type_desc.tokens == []
                        elif auth_type_name == "HEADER":
                            header_names_lower = tuple(
                                sorted(
                                    t.accepted.lower() for t in auth_type_desc.tokens
                                )
                            )
                            assert header_names_lower in {
                                (EMBEDDING_HEADER_API_KEY.lower(),),
                                (
                                    EMBEDDING_HEADER_AWS_ACCESS_ID.lower(),
                                    EMBEDDING_HEADER_AWS_SECRET_ID.lower(),
                                ),
                            }
                        elif auth_type_name == "SHARED_SECRET":
                            authkey_names = tuple(
                                sorted(t.accepted for t in auth_type_desc.tokens)
                            )
                            assert authkey_names in {
                                ("providerKey",),
                                ("accessId", "secretKey"),
                            }
                        else:
                            raise ValueError("Unknown auth type")

                        # params
                        collated_params = provider_desc.parameters + model.parameters
                        all_nond_params = [
                            param
                            for param in collated_params
                            if param.name != "vectorDimension"
                        ]
                        required_nond_params = {
                            param.name for param in all_nond_params if param.required
                        }
                        optional_nond_params = {
                            param.name
                            for param in all_nond_params
                            if not param.required
                        }
                        #
                        d_params = [
                            param
                            for param in collated_params
                            if param.name == "vectorDimension"
                        ]
                        if d_params:
                            d_param = d_params[0]
                            if (provider_name, model.name) in FORCE_DIMENSION_MAP:
                                optional_dimension = False
                                dimension = FORCE_DIMENSION_MAP[
                                    (provider_name, model.name)
                                ]
                            elif d_param.default_value is not None:
                                optional_dimension = True
                                assert model.vector_dimension is None
                                dimension = _from_validation(d_param)
                            else:
                                optional_dimension = False
                                assert model.vector_dimension is None
                                dimension = _from_validation(d_param)
                        else:
                            optional_dimension = False
                            assert model.vector_dimension is not None
                            assert model.vector_dimension > 0
                            dimension = model.vector_dimension

                        model_parameters = {
                            param_name: PARAMETER_VALUE_MAP[
                                (provider_name, model.name, param_name)
                            ]
                            for param_name in required_nond_params
                        }
                        optional_model_parameters = {
                            param_name: PARAMETER_VALUE_MAP[
                                (provider_name, model.name, param_name)
                            ]
                            for param_name in optional_nond_params
                        }

                        minimal_model_quadruple = (
                            provider_name,
                            model.name,
                            auth_type_name,
                            "0",
                        )
                        if minimal_model_quadruple not in EXCLUDED_MODEL_QUADRUPLES:
                            if optional_dimension or optional_nond_params != set():
                                # we issue a minimal-params version
                                model_tag_0 = (
                                    f"{provider_name}/{model.name}/{auth_type_name}/0"
                                )
                                this_minimal_model = {
                                    "model_tag": model_tag_0,
                                    "simple_tag": _collapse(
                                        "".join(c for c in model_tag_0 if c in alphanum)
                                    ),
                                    "auth_type_name": auth_type_name,
                                    "auth_type_tokens": auth_type_desc.tokens,
                                    "secret_tag": SECRET_NAME_ROOT_MAP[provider_name],
                                    "test_assets": TEST_ASSETS_MAP.get(
                                        (provider_name, model.name), DEFAULT_TEST_ASSETS
                                    ),
                                    "use_insert_one": USE_INSERT_ONE_MAP.get(
                                        (provider_name, model.name), False
                                    ),
                                    "service_options": VectorServiceOptions(
                                        provider=provider_name,
                                        model_name=model.name,
                                        parameters=model_parameters,
                                    ),
                                }
                                yield this_minimal_model

                        # and in any case we issue a 'full-spec' one ...
                        # ... unless explicitly marked as skipped
                        if all(
                            v != PARAM_SKIP_MARKER
                            for v in optional_model_parameters.values()
                        ):
                            root_model = {
                                "auth_type_name": auth_type_name,
                                "auth_type_tokens": auth_type_desc.tokens,
                                "dimension": dimension,
                                "secret_tag": SECRET_NAME_ROOT_MAP[provider_name],
                                "test_assets": TEST_ASSETS_MAP.get(
                                    (provider_name, model.name), DEFAULT_TEST_ASSETS
                                ),
                                "use_insert_one": USE_INSERT_ONE_MAP.get(
                                    (provider_name, model.name), False
                                ),
                            }

                            model_tag_f = (
                                f"{provider_name}/{model.name}/{auth_type_name}/f"
                            )
                            full_model_quadruple = (
                                provider_name,
                                model.name,
                                auth_type_name,
                                "f",
                            )
                            if full_model_quadruple not in EXCLUDED_MODEL_QUADRUPLES:
                                this_model = {
                                    "model_tag": model_tag_f,
                                    "simple_tag": _collapse(
                                        "".join(c for c in model_tag_f if c in alphanum)
                                    ),
                                    "service_options": VectorServiceOptions(
                                        provider=provider_name,
                                        model_name=model.name,
                                        parameters={
                                            **model_parameters,
                                            **optional_model_parameters,
                                        },
                                    ),
                                    **root_model,
                                }
                                yield this_model
