# -*- 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

from unittest import mock

from absl.testing import absltest

from google.api_core import exceptions
from google.generativeai import text
from google.generativeai.notebook import text_model
from google.generativeai.notebook.lib import model as model_lib


def _fake_generator(
    prompt: str,
    model: str | None = None,
    temperature: float | None = None,
    candidate_count: int | None = None,
) -> text.Completion:
    return text.Completion(
        prompt=prompt,
        model=model,
        temperature=temperature,
        candidate_count=candidate_count,
        # Smuggle the parameters as text output, so we can make assertions.
        candidates=[
            {"output": f"{prompt}_1"},
            {"output": model},
            {"output": temperature},
            {"output": candidate_count},
        ],
    )


class TestModel(text_model.TextModel):
    """A TextModel, but with _generate_text stubbed out."""

    def _generate_text(
        self,
        prompt: str,
        model: str | None = None,
        temperature: float | None = None,
        candidate_count: int | None = None,
        **kwargs,
    ) -> text.Completion:
        return _fake_generator(
            prompt=prompt,
            model=model,
            temperature=temperature,
            candidate_count=candidate_count,
        )


class TextModelTestCase(absltest.TestCase):
    def test_generate_text(self):
        model = TestModel()

        result = model.call_model("prompt goes in")
        self.assertEqual(result.text_results[0], "prompt goes in_1")
        self.assertIsNone(result.text_results[1])
        self.assertIsNone(result.text_results[2])
        self.assertIsNone(result.text_results[3])

        args = model_lib.ModelArguments(model="model_name", temperature=0.42, candidate_count=5)
        result = model.call_model("prompt goes in", args)
        self.assertEqual(result.text_results[0], "prompt goes in_1")
        self.assertEqual(result.text_results[1], "model_name")
        self.assertEqual(result.text_results[2], 0.42)
        self.assertEqual(result.text_results[3], 5)

    def test_retry(self):
        model = TestModel()

        with mock.patch.object(model, "_generate_text") as erroneous_generator:
            erroneous_generator.side_effect = [
                exceptions.ResourceExhausted("Over quota"),
                mock.DEFAULT,
            ]

            _ = model.call_model("phew it worked")

        self.assertEqual(erroneous_generator.call_count, 2)


if __name__ == "__main__":
    absltest.main()