# -*- coding: utf-8 -*-
# Copyright 2023 Google LLC
#
# 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 dataclasses
from collections.abc import Iterable, Sequence
import itertools
from typing import Any, Iterable, overload, TypeVar

import google.ai.generativelanguage as glm

from google.generativeai import protos

from google.generativeai.client import get_default_text_client
from google.generativeai import string_utils
from google.generativeai.types import helper_types
from google.generativeai.types import text_types
from google.generativeai.types import model_types
from google.generativeai import models
from google.generativeai.types import palm_safety_types

DEFAULT_TEXT_MODEL = "models/text-bison-001"
EMBEDDING_MAX_BATCH_SIZE = 100

try:
    # python 3.12+
    _batched = itertools.batched  # type: ignore
except AttributeError:
    T = TypeVar("T")

    def _batched(iterable: Iterable[T], n: int) -> Iterable[list[T]]:
        if n < 1:
            raise ValueError(f"Batch size `n` must be >1, got: {n}")
        batch = []
        for item in iterable:
            batch.append(item)
            if len(batch) == n:
                yield batch
                batch = []

        if batch:
            yield batch


def _make_text_prompt(prompt: str | dict[str, str]) -> protos.TextPrompt:
    """
    Creates a `protos.TextPrompt` object based on the provided prompt input.

    Args:
        prompt: The prompt input, either a string or a dictionary.

    Returns:
        protos.TextPrompt: A TextPrompt object containing the prompt text.

    Raises:
        TypeError: If the provided prompt is neither a string nor a dictionary.
    """
    if isinstance(prompt, str):
        return protos.TextPrompt(text=prompt)
    elif isinstance(prompt, dict):
        return protos.TextPrompt(prompt)
    else:
        raise TypeError(
            "Invalid argument type: Expected a string or dictionary for the text prompt."
        )


def _make_generate_text_request(
    *,
    model: model_types.AnyModelNameOptions = DEFAULT_TEXT_MODEL,
    prompt: str | None = None,
    temperature: float | None = None,
    candidate_count: int | None = None,
    max_output_tokens: int | None = None,
    top_p: int | None = None,
    top_k: int | None = None,
    safety_settings: palm_safety_types.SafetySettingOptions | None = None,
    stop_sequences: str | Iterable[str] | None = None,
) -> protos.GenerateTextRequest:
    """
    Creates a `protos.GenerateTextRequest` object based on the provided parameters.

    This function generates a `protos.GenerateTextRequest` object with the specified
    parameters. It prepares the input parameters and creates a request that can be
    used for generating text using the chosen model.

    Args:
        model: The model to use for text generation.
        prompt: The prompt for text generation. Defaults to None.
        temperature: The temperature for randomness in generation. Defaults to None.
        candidate_count: The number of candidates to consider. Defaults to None.
        max_output_tokens: The maximum number of output tokens. Defaults to None.
        top_p: The nucleus sampling probability threshold. Defaults to None.
        top_k: The top-k sampling parameter. Defaults to None.
        safety_settings: Safety settings for generated text. Defaults to None.
        stop_sequences: Stop sequences to halt text generation. Can be a string
             or iterable of strings. Defaults to None.

    Returns:
        `protos.GenerateTextRequest`: A `GenerateTextRequest` object configured with the specified parameters.
    """
    model = model_types.make_model_name(model)
    prompt = _make_text_prompt(prompt=prompt)
    safety_settings = palm_safety_types.normalize_safety_settings(safety_settings)
    if isinstance(stop_sequences, str):
        stop_sequences = [stop_sequences]
    if stop_sequences:
        stop_sequences = list(stop_sequences)

    return protos.GenerateTextRequest(
        model=model,
        prompt=prompt,
        temperature=temperature,
        candidate_count=candidate_count,
        max_output_tokens=max_output_tokens,
        top_p=top_p,
        top_k=top_k,
        safety_settings=safety_settings,
        stop_sequences=stop_sequences,
    )


def generate_text(
    *,
    model: model_types.AnyModelNameOptions = DEFAULT_TEXT_MODEL,
    prompt: str,
    temperature: float | None = None,
    candidate_count: int | None = None,
    max_output_tokens: int | None = None,
    top_p: float | None = None,
    top_k: float | None = None,
    safety_settings: palm_safety_types.SafetySettingOptions | None = None,
    stop_sequences: str | Iterable[str] | None = None,
    client: glm.TextServiceClient | None = None,
    request_options: helper_types.RequestOptionsType | None = None,
) -> text_types.Completion:
    """Calls the API to generate text based on the provided prompt.

    Args:
        model: Which model to call, as a string or a `types.Model`.
        prompt: Free-form input text given to the model. Given a prompt, the model will
                generate text that completes the input text.
        temperature: Controls the randomness of the output. Must be positive.
            Typical values are in the range: `[0.0,1.0]`. Higher values produce a
            more random and varied response. A temperature of zero will be deterministic.
        candidate_count: The **maximum** number of generated response messages to return.
            This value must be between `[1, 8]`, inclusive. If unset, this
            will default to `1`.

            Note: Only unique candidates are returned. Higher temperatures are more
            likely to produce unique candidates. Setting `temperature=0.0` will always
            return 1 candidate regardless of the `candidate_count`.
        max_output_tokens: Maximum number of tokens to include in a candidate. Must be greater
                           than zero. If unset, will default to 64.
        top_k: The API uses combined [nucleus](https://arxiv.org/abs/1904.09751) and top-k sampling.
            `top_k` sets the maximum number of tokens to sample from on each step.
        top_p: The API uses combined [nucleus](https://arxiv.org/abs/1904.09751) and top-k sampling.
            `top_p` configures the nucleus sampling. It sets the maximum cumulative
            probability of tokens to sample from.
            For example, if the sorted probabilities are
            `[0.5, 0.2, 0.1, 0.1, 0.05, 0.05]` a `top_p` of `0.8` will sample
            as `[0.625, 0.25, 0.125, 0, 0, 0]`.
        safety_settings: A list of unique `types.SafetySetting` instances for blocking unsafe content.
           These will be enforced on the `prompt` and
           `candidates`. There should not be more than one
           setting for each `types.SafetyCategory` type. The API will block any prompts and
           responses that fail to meet the thresholds set by these settings. This list
           overrides the default settings for each `SafetyCategory` specified in the
           safety_settings. If there is no `types.SafetySetting` for a given
           `SafetyCategory` provided in the list, the API will use the default safety
           setting for that category.
        stop_sequences: A set of up to 5 character sequences that will stop output generation.
          If specified, the API will stop at the first appearance of a stop
          sequence. The stop sequence will not be included as part of the response.
        client: If you're not relying on a default client, you pass a `glm.TextServiceClient` instead.
        request_options: Options for the request.

    Returns:
        A `types.Completion` containing the model's text completion response.
    """
    request = _make_generate_text_request(
        model=model,
        prompt=prompt,
        temperature=temperature,
        candidate_count=candidate_count,
        max_output_tokens=max_output_tokens,
        top_p=top_p,
        top_k=top_k,
        safety_settings=safety_settings,
        stop_sequences=stop_sequences,
    )

    return _generate_response(client=client, request=request, request_options=request_options)


@string_utils.prettyprint
@dataclasses.dataclass(init=False)
class Completion(text_types.Completion):
    def __init__(self, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)

        self.result = None
        if self.candidates:
            self.result = self.candidates[0]["output"]


def _generate_response(
    request: protos.GenerateTextRequest,
    client: glm.TextServiceClient = None,
    request_options: helper_types.RequestOptionsType | None = None,
) -> Completion:
    """
    Generates a response using the provided `protos.GenerateTextRequest` and client.

    Args:
        request: The text generation request.
        client: The client to use for text generation. Defaults to None, in which
            case the default text client is used.
        request_options: Options for the request.

    Returns:
        `Completion`: A `Completion` object with the generated text and response information.
    """
    if request_options is None:
        request_options = {}

    if client is None:
        client = get_default_text_client()

    response = client.generate_text(request, **request_options)
    response = type(response).to_dict(response)

    response["filters"] = palm_safety_types.convert_filters_to_enums(response["filters"])
    response["safety_feedback"] = palm_safety_types.convert_safety_feedback_to_enums(
        response["safety_feedback"]
    )
    response["candidates"] = palm_safety_types.convert_candidate_enums(response["candidates"])

    return Completion(_client=client, **response)


def count_text_tokens(
    model: model_types.AnyModelNameOptions,
    prompt: str,
    client: glm.TextServiceClient | None = None,
    request_options: helper_types.RequestOptionsType | None = None,
) -> text_types.TokenCount:
    """Calls the API to count the number of tokens in the text prompt."""

    base_model = models.get_base_model_name(model)

    if request_options is None:
        request_options = {}

    if client is None:
        client = get_default_text_client()

    result = client.count_text_tokens(
        protos.CountTextTokensRequest(model=base_model, prompt={"text": prompt}),
        **request_options,
    )

    return type(result).to_dict(result)


@overload
def generate_embeddings(
    model: model_types.BaseModelNameOptions,
    text: str,
    client: glm.TextServiceClient = None,
    request_options: helper_types.RequestOptionsType | None = None,
) -> text_types.EmbeddingDict: ...


@overload
def generate_embeddings(
    model: model_types.BaseModelNameOptions,
    text: Sequence[str],
    client: glm.TextServiceClient = None,
    request_options: helper_types.RequestOptionsType | None = None,
) -> text_types.BatchEmbeddingDict: ...


def generate_embeddings(
    model: model_types.BaseModelNameOptions,
    text: str | Sequence[str],
    client: glm.TextServiceClient = None,
    request_options: helper_types.RequestOptionsType | None = None,
) -> text_types.EmbeddingDict | text_types.BatchEmbeddingDict:
    """Calls the API to create an embedding for the text passed in.

    Args:
        model: Which model to call, as a string or a `types.Model`.

        text: Free-form input text given to the model. Given a string, the model will
              generate an embedding based on the input text.

        client: If you're not relying on a default client, you pass a `glm.TextServiceClient` instead.

        request_options: Options for the request.

    Returns:
        Dictionary containing the embedding (list of float values) for the input text.
    """
    model = model_types.make_model_name(model)

    if request_options is None:
        request_options = {}

    if client is None:
        client = get_default_text_client()

    if isinstance(text, str):
        embedding_request = protos.EmbedTextRequest(model=model, text=text)
        embedding_response = client.embed_text(
            embedding_request,
            **request_options,
        )
        embedding_dict = type(embedding_response).to_dict(embedding_response)
        embedding_dict["embedding"] = embedding_dict["embedding"]["value"]
    else:
        result = {"embedding": []}
        for batch in _batched(text, EMBEDDING_MAX_BATCH_SIZE):
            # TODO(markdaoust): This could use an option for returning an iterator or wait-bar.
            embedding_request = protos.BatchEmbedTextRequest(model=model, texts=batch)
            embedding_response = client.batch_embed_text(
                embedding_request,
                **request_options,
            )
            embedding_dict = type(embedding_response).to_dict(embedding_response)
            result["embedding"].extend(e["value"] for e in embedding_dict["embeddings"])
        return result

    return embedding_dict