# -*- 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.
"""Utilities for handling input variables."""
from __future__ import annotations

from typing import Any, Mapping, Sequence, Union

from google.generativeai.notebook.lib import llmfn_inputs_source


_NormalizedInputsList = llmfn_inputs_source.NormalizedInputsList

_ColumnOrderValuesList = Mapping[str, Sequence[str]]

LLMFunctionInputs = Union[
    _ColumnOrderValuesList,
    llmfn_inputs_source.LLMFnInputsSource,
]


def _is_column_order_values_list(inputs: Any) -> bool:
  """See if inputs is of the form: {"key1": ["val1", "val2", ...]}.

  This is similar to the format produced by:
    pandas.DataFrame.to_dict(orient="list")

  Args:
    inputs: The inputs passed into an LLMFunction.

  Returns:
    Whether `inputs` is a column-ordered list of values.
  """
  if not isinstance(inputs, Mapping):
    return False
  for x in inputs.values():
    if not isinstance(x, Sequence):
      return False
    # Strings and bytes are also considered Sequences but we disallow them
    # here because the values contained in their Sequences are single
    # characters rather than words.
    if isinstance(x, str) or isinstance(x, bytes):
      return False
  return True


# TODO(b/273688393): Perform stricter validation on `inputs`.
def _normalize_column_order_values_list(
    inputs: _ColumnOrderValuesList,
) -> _NormalizedInputsList:
  """Transforms prompt inputs into a list of dictionaries."""
  return_list: list[dict[str, str]] = []
  keys = list(inputs.keys())
  if keys:
    first_key = keys[0]
    for row_num in range(len(inputs[first_key])):
      row_dict = {}
      return_list.append(row_dict)
      for key in keys:
        row_dict[key] = inputs[key][row_num]
  return return_list


def to_normalized_inputs(inputs: LLMFunctionInputs) -> _NormalizedInputsList:
  """Handles the different types of `inputs` and returns a normalized form."""
  normalized_inputs: list[Mapping[str, str]] = []
  if isinstance(inputs, llmfn_inputs_source.LLMFnInputsSource):
    normalized_inputs.extend(inputs.to_normalized_inputs())
  elif _is_column_order_values_list(inputs):
    normalized_inputs.extend(_normalize_column_order_values_list(inputs))
  else:
    raise ValueError("Unsupported input type {!r}".format(inputs))
  return normalized_inputs