# -*- 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.
"""Unittest for Argument Definition."""
from __future__ import annotations

import argparse
import enum
import math
from absl.testing import absltest
from google.generativeai.notebook import argument_parser
from google.generativeai.notebook import flag_def


def _new_parser(flag: flag_def.FlagDef) -> argparse.ArgumentParser:
  """Returns a new ArgumentParser with `flag` added as an argument."""
  parser = argument_parser.ArgumentParser()
  flag.add_argument_to_parser(parser)
  return parser


class SingleValueFlagDefTest(absltest.TestCase):

  def test_short_name(self):
    flag = flag_def.SingleValueFlagDef(
        name="value", short_name="v", parse_type=str, required=True
    )

    results = _new_parser(flag).parse_args(["--value=forty-one"])
    self.assertEqual("forty-one", results.value)
    results = _new_parser(flag).parse_args(["-v", "forty-two"])
    self.assertEqual("forty-two", results.value)

  def test_cardinality(self):
    flag = flag_def.SingleValueFlagDef(
        name="value", parse_type=str, required=True
    )

    # Parser should not accept flag without no values.
    with self.assertRaisesRegex(
        argument_parser.ParserError, "expected 1 argument"
    ):
      _new_parser(flag).parse_args(["--value"])

    # Parser should not accept flag with more-than-one value.
    with self.assertRaisesRegex(
        argument_parser.ParserError, "unrecognized arguments: forty-two"
    ):
      _new_parser(flag).parse_args(["--value", "forty-one", "forty-two"])

    # Parser should not accept a command-line with the flag specified
    # more-than-once.
    with self.assertRaisesRegex(
        argument_parser.ParserError, "Cannot set --value more than once"
    ):
      _new_parser(flag).parse_args(
          ["--value", "forty-one", "--value", "forty-two"]
      )

    results = _new_parser(flag).parse_args(["--value", "forty-one"])
    self.assertEqual("forty-one", results.value)

  def test_required(self):
    req_flag = flag_def.SingleValueFlagDef(
        name="value",
        parse_type=str,
        required=True,
    )

    # Parser is able to parse a command line containing the required argument.
    results = _new_parser(req_flag).parse_args(["--value=forty-two"])
    self.assertEqual("forty-two", results.value)

    # Parser should raise an Exception if the commandline does not contain the
    # required argument.
    with self.assertRaisesRegex(
        argument_parser.ParserError, "the following arguments are required"
    ):
      _new_parser(req_flag).parse_args([])

  def test_optional(self):
    # Optional flags should have a default value
    with self.assertRaisesRegex(
        ValueError, "Optional flags must have a default value"
    ):
      flag_def.SingleValueFlagDef(
          name="value",
          parse_type=str,
          required=False,
      )

    opt_flag = flag_def.SingleValueFlagDef(
        name="value",
        parse_type=str,
        required=False,
        default_value="zero",
    )

    # Parser is able to parse a command line containing the optional argument.
    results = _new_parser(opt_flag).parse_args(["--value=forty-two"])
    self.assertEqual("forty-two", results.value)

    # Parser should return the default value if the command line does not
    # contain the optional argument.
    results = _new_parser(opt_flag).parse_args([])
    self.assertEqual("zero", results.value)

    # Parser should not accept flag without a value.
    with self.assertRaisesRegex(
        argument_parser.ParserError, "expected 1 argument"
    ):
      _new_parser(opt_flag).parse_args(["--value"])

  def test_default_is_none(self):
    """Make sure None can be accepted as a default value."""
    opt_flag = flag_def.SingleValueFlagDef(
        name="value",
        parse_type=str,
        required=False,
        default_value=None,
    )

    # Parser is able to parse a command line containing the optional argument.
    results = _new_parser(opt_flag).parse_args(["--value=forty-two"])
    self.assertEqual("forty-two", results.value)

    # Parser should return the default value if the command line does not
    # contain the optional argument.
    results = _new_parser(opt_flag).parse_args([])
    self.assertIsNone(results.value)

    # Parser should not accept flag without a value.
    with self.assertRaisesRegex(
        argument_parser.ParserError, "expected 1 argument"
    ):
      _new_parser(opt_flag).parse_args(["--value"])

  def test_type_conversion(self):
    # Default value must be of the same type as destination type.
    with self.assertRaisesRegex(
        ValueError,
        "Default value must be of the same type as the destination type",
    ):
      flag_def.SingleValueFlagDef(
          name="value",
          parse_type=int,
          required=False,
          default_value="zero",
      )

    int_flag_def = flag_def.SingleValueFlagDef(
        name="value",
        parse_type=int,
        required=False,
        default_value=0,
    )

    # Parser should not accept a value of the wrong type.
    with self.assertRaisesRegex(
        argument_parser.ParserError, "invalid int value: 'forty-two'"
    ):
      _new_parser(int_flag_def).parse_args(["--value", "forty-two"])

    results = _new_parser(int_flag_def).parse_args(["--value", "42"])
    self.assertEqual(42, results.value)

  def test_validation(self):
    def _check_is_not_nan(x: float) -> float:
      if math.isnan(x):
        raise ValueError("Must not be NAN")
      return x

    float_flag_def = flag_def.SingleValueFlagDef(
        name="value",
        parse_type=float,
        parse_to_dest_type_fn=_check_is_not_nan,
        required=True,
    )

    results = _new_parser(float_flag_def).parse_args(["--value", "0.25"])
    self.assertEqual(0.25, results.value)

    with self.assertRaisesRegex(argument_parser.ParserError, "Must not be NAN"):
      _new_parser(float_flag_def).parse_args(["--value", "nan"])


class ColorsEnum(enum.Enum):
  RED = "red"
  GREEN = "green"
  BLUE = "blue"


class EnumFlagDefTest(absltest.TestCase):

  def test_construction(self):
    # "enum_type" must be provided.
    with self.assertRaisesRegex(
        TypeError, "missing 1 required keyword-only argument"
    ):
      # pylint: disable-next=missing-kwoa
      flag_def.EnumFlagDef(name="color", required=True)  # type: ignore

    # "parse_type" cannot be provided.
    with self.assertRaisesRegex(
        ValueError, 'Cannot set "parse_type" for EnumFlagDef'
    ):
      flag_def.EnumFlagDef(
          name="color", required=True, enum_type=ColorsEnum, parse_type=int
      )
    # "dest_type" cannot be provided.
    with self.assertRaisesRegex(
        ValueError, 'Cannot set "dest_type" for EnumFlagDef'
    ):
      flag_def.EnumFlagDef(
          name="color", required=True, enum_type=ColorsEnum, dest_type=str
      )
    # This should succeed.
    flag_def.EnumFlagDef(name="color", required=True, enum_type=ColorsEnum)

  def test_parsing(self):
    flag = flag_def.EnumFlagDef(
        name="color",
        required=False,
        enum_type=ColorsEnum,
        default_value=ColorsEnum.RED,
    )

    # "teal" is not one of the enum values.
    with self.assertRaisesRegex(
        argument_parser.ParserError, "invalid choice: 'teal'"
    ):
      _new_parser(flag).parse_args(["--color=teal"])

    results = _new_parser(flag).parse_args(["--color=red"])
    self.assertEqual(ColorsEnum.RED, results.color)
    results = _new_parser(flag).parse_args(["--color=green"])
    self.assertEqual(ColorsEnum.GREEN, results.color)
    results = _new_parser(flag).parse_args(["--color=blue"])
    self.assertEqual(ColorsEnum.BLUE, results.color)

  def test_choices(self):
    # If `choices` is provided, all values must be valid.
    with self.assertRaisesRegex(ValueError, 'Invalid value in "choices"'):
      flag_def.EnumFlagDef(
          name="color",
          required=True,
          enum_type=ColorsEnum,
          choices=["red", "green", "teal"],
      )

    # Exclude "blue".
    flag = flag_def.EnumFlagDef(
        name="color",
        required=True,
        enum_type=ColorsEnum,
        choices=["red", "green"],
    )

    # "blue" is no longer one of the choices.
    with self.assertRaisesRegex(
        argument_parser.ParserError, "invalid choice: 'blue'"
    ):
      _new_parser(flag).parse_args(["--color=blue"])

    results = _new_parser(flag).parse_args(["--color=red"])
    self.assertEqual(ColorsEnum.RED, results.color)
    results = _new_parser(flag).parse_args(["--color=green"])
    self.assertEqual(ColorsEnum.GREEN, results.color)


class MultiValuesFlagDefTest(absltest.TestCase):

  def test_basic(self):
    # Default value is not needed even if optional; the value would just be the
    # empty list.
    flag = flag_def.MultiValuesFlagDef(
        name="colors", parse_type=str, required=False
    )

    # Default value is the empty list.
    results = _new_parser(flag).parse_args([])
    self.assertEmpty(results.colors)

    results = _new_parser(flag).parse_args(["--colors", "red"])
    self.assertEqual(["red"], results.colors)

    results = _new_parser(flag).parse_args(["--colors", "red", "green"])
    self.assertEqual(["red", "green"], results.colors)

  def test_required(self):
    flag = flag_def.MultiValuesFlagDef(
        name="colors",
        parse_type=str,
        required=True,
    )

    # Parser is able to parse a command line containing the required argument.
    results = _new_parser(flag).parse_args(["--colors", "red"])
    self.assertEqual(["red"], results.colors)

    # Parser should raise an Exception if the commandline does not contain the
    # required argument.
    with self.assertRaisesRegex(
        argument_parser.ParserError, "the following arguments are required"
    ):
      _new_parser(flag).parse_args([])

  def test_cannot_set_default_value(self):
    # `default_value` is not a field for MultiValueFlagsDef.
    with self.assertRaisesRegex(
        TypeError, "got an unexpected keyword argument"
    ):
      # pylint: disable-next=unexpected-keyword-arg
      flag_def.MultiValuesFlagDef(  # type: ignore
          name="colors",
          parse_type=str,
          required=False,
          default_value="fuschia",
      )

  def test_values_must_be_unique(self):
    flag = flag_def.MultiValuesFlagDef(name="colors")

    # Cannot specify "red" more than once.
    with self.assertRaisesRegex(
        argument_parser.ParserError, 'Duplicate values "red"'
    ):
      _new_parser(flag).parse_args(["--colors", "red", "green", "red"])

  def test_cardinality(self):
    flag = flag_def.MultiValuesFlagDef(
        name="colors",
        parse_type=str,
        required=False,
    )

    # Must have at least one argument.
    with self.assertRaisesRegex(
        argument_parser.ParserError, "expected at least one argument"
    ):
      _new_parser(flag).parse_args(["--colors"])

    # Cannot specify "--colors" more than once.
    with self.assertRaisesRegex(
        argument_parser.ParserError, "Cannot set --colors more than once"
    ):
      _new_parser(flag).parse_args(["--colors", "red", "--colors", "blue"])

  def test_dest_type_conversion(self):
    flag = flag_def.MultiValuesFlagDef(
        name="colors",
        parse_type=str,
        dest_type=ColorsEnum,
        required=False,
        choices=[x.value for x in ColorsEnum],
    )

    # "fuschia" is not a valid value for enum.
    with self.assertRaisesRegex(
        argument_parser.ParserError, "invalid choice: 'fuschia'"
    ):
      _new_parser(flag).parse_args(["--colors", "fuschia"])

    # Results are converted to a list of enums.
    results = _new_parser(flag).parse_args(["--colors", "red", "green"])
    self.assertEqual([ColorsEnum.RED, ColorsEnum.GREEN], results.colors)

  def test_validation(self):
    def _check_is_not_nan(x: float) -> float:
      if math.isnan(x):
        raise ValueError("Must not be NAN")
      return x

    flag = flag_def.MultiValuesFlagDef(
        name="values",
        parse_type=float,
        parse_to_dest_type_fn=_check_is_not_nan,
    )

    results = _new_parser(flag).parse_args(["--value", "0.25", "0.5"])
    self.assertEqual([0.25, 0.5], results.values)

    with self.assertRaisesRegex(argument_parser.ParserError, "Must not be NAN"):
      _new_parser(flag).parse_args(["--value", "0.25", "nan"])


class BooleanFlagDefTest(absltest.TestCase):

  def test_basic(self):
    flag = flag_def.BooleanFlagDef(name="unique")

    results = _new_parser(flag).parse_args([])
    self.assertFalse(results.unique)

    results = _new_parser(flag).parse_args(["--unique"])
    self.assertTrue(results.unique)

  def test_constructor(self):
    """Check that invalid constructor arguments are rejected."""
    with self.assertRaisesRegex(
        ValueError, "dest_type cannot be set for BooleanFlagDef"
    ):
      flag_def.BooleanFlagDef(name="unique", dest_type=bool)
    with self.assertRaisesRegex(
        ValueError, "parse_to_dest_type_fn cannot be set for BooleanFlagDef"
    ):
      flag_def.BooleanFlagDef(
          name="unique", parse_to_dest_type_fn=lambda x: True
      )
    with self.assertRaisesRegex(
        ValueError, "choices cannot be set for BooleanFlagDef"
    ):
      flag_def.BooleanFlagDef(name="unique", choices=[True])

  def test_cardinality(self):
    flag = flag_def.BooleanFlagDef(name="unique")

    with self.assertRaisesRegex(
        argument_parser.ParserError, "error: unrecognized arguments: True"
    ):
      _new_parser(flag).parse_args(["--unique", "True"])

    with self.assertRaisesRegex(
        argument_parser.ParserError, "Cannot set --unique more than once"
    ):
      _new_parser(flag).parse_args(["--unique", "--unique"])


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