# -*- 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.
"""LLMFnOutputRow."""
from __future__ import annotations

import abc
from typing import Any, Iterator, Mapping


# The type of value stored in a cell.
_CELLVALUETYPE = Any


def _get_name_of_type(x: type[Any]) -> str:
  if hasattr(x, "__name__"):
    return x.__name__
  return str(x)


def _validate_is_result_type(value: Any, result_type: type[Any]) -> None:
  if result_type == Any:
    return
  if not isinstance(value, result_type):
    raise ValueError(
        'Value of last entry must be of type "{}", got "{}"'.format(
            _get_name_of_type(result_type),
            _get_name_of_type(type(value)),
        )
    )


class LLMFnOutputRowView(Mapping[str, _CELLVALUETYPE], metaclass=abc.ABCMeta):
  """Immutable view of LLMFnOutputRow."""

  # Additional methods (not required by Mapping[str, _CELLVALUETYPE])
  @abc.abstractmethod
  def __contains__(self, k: str) -> bool:
    """For expressions like: x in this_instance."""

  @abc.abstractmethod
  def __str__(self) -> str:
    """For expressions like: str(this_instance)."""

  # Own methods.
  @abc.abstractmethod
  def result_type(self) -> type[Any]:
    """Returns the type enforced for the result cell."""

  @abc.abstractmethod
  def result_value(self) -> Any:
    """Get the value of the result cell."""

  @abc.abstractmethod
  def result_key(self) -> str:
    """Get the key of the result cell."""


class LLMFnOutputRow(LLMFnOutputRowView):
  """Container that represents a single row in a table of outputs.

  We represent outputs as a table. This class represents a single row in the
  table like a dictionary, where the key is the column name and the value is the
  cell value.

  A single cell is designated the "result". This contains the output of the LLM
  model after running any post-processing functions specified by the user.

  In addition to behaving like a dictionary, this class provides additional
  methods, including:
  - Getting the value of the "result" cell
  - Setting the value (and optionally the key) of the "result" cell.
  - Add a new non-result cell

  Notes: As an implementation detail, the result-cell is always kept as the
  rightmost cell.
  """

  def __init__(
      self, data: Mapping[str, _CELLVALUETYPE], result_type: type[Any]
  ):
    """Constructor.

    Args:
      data: The initial value of the row. The last entry will be treated as the
        result. Cannot be empty. The value of the last entry must be `str`.
      result_type: The type of the result cell. This will be enforced at
        runtime.
    """
    self._data: dict[str, _CELLVALUETYPE] = dict(data)
    if not self._data:
      raise ValueError("Must provide non-empty data")

    self._result_type = result_type
    result_value = list(self._data.values())[-1]
    _validate_is_result_type(result_value, self._result_type)

  # Methods needed for Mapping[str, _CELLVALUETYPE]:
  def __iter__(self) -> Iterator[str]:
    return self._data.__iter__()

  def __len__(self) -> int:
    return self._data.__len__()

  def __getitem__(self, k: str) -> _CELLVALUETYPE:
    return self._data.__getitem__(k)

  # Additional methods for LLMFnOutputRowView.
  def __contains__(self, k: str) -> bool:
    return self._data.__contains__(k)

  def __str__(self) -> str:
    return "LLMFnOutputRow: {}".format(self._data.__str__())

  def result_type(self) -> type[Any]:
    return self._result_type

  def result_value(self) -> Any:
    return self._data[self.result_key()]

  def result_key(self) -> str:
    # Our invariant is that the result-cell is always the rightmost cell.
    return list(self._data.keys())[-1]

  # Mutable methods.
  def set_result_value(self, value: Any, key: str | None = None) -> None:
    """Set the value of the result cell.

    Sets the value (and optionally the key) of the result cell.

    Args:
      value: The value to set the result cell today.
      key: Optionally change the key as well.
    """
    _validate_is_result_type(value, self._result_type)

    current_key = self.result_key()
    if key is None or key == current_key:
      self._data[current_key] = value
      return

    del self._data[current_key]
    self._data[key] = value

  def add(self, key: str, value: _CELLVALUETYPE) -> None:
    """Add a non-result cell.

    Adds a new non-result cell. This does not affect the result cell.

    Args:
      key: The key of the new cell to add.
      value: The value of the new cell to add.
    """
    # Handle collisions with `key`.
    if key in self._data:
      idx = 1
      candidate_key = key
      while candidate_key in self._data:
        candidate_key = "{}_{}".format(key, idx)
        idx = idx + 1
      key = candidate_key

    # Insert the new key/value into the second rightmost position to keep
    # the result cell as the rightmost cell.
    result_key = self.result_key()
    result_value = self._data.pop(result_key)
    self._data[key] = value
    self._data[result_key] = result_value