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

Common methods for Commands such as RunCommand and CompileCommand.
"""
from __future__ import annotations

from typing import AbstractSet, Any, Callable, Sequence

from google.generativeai.notebook import ipython_env
from google.generativeai.notebook import model_registry
from google.generativeai.notebook import parsed_args_lib
from google.generativeai.notebook import post_process_utils
from google.generativeai.notebook.lib import llm_function
from google.generativeai.notebook.lib import llmfn_input_utils
from google.generativeai.notebook.lib import llmfn_output_row
from google.generativeai.notebook.lib import llmfn_outputs
from google.generativeai.notebook.lib import unique_fn


class _GroundTruthLLMFunction(llm_function.LLMFunction):
  """LLMFunction that returns pre-generated ground truth data."""

  def __init__(self, data: Sequence[str]):
    super().__init__(outputs_ipython_display_fn=None)
    self._data = data

  def get_placeholders(self) -> AbstractSet[str]:
    # Ground truth is fixed and thus has no placeholders.
    return frozenset({})

  def _call_impl(
      self, inputs: llmfn_input_utils.LLMFunctionInputs | None
  ) -> Sequence[llmfn_outputs.LLMFnOutputEntry]:
    normalized_inputs = llmfn_input_utils.to_normalized_inputs(inputs)
    if len(self._data) != len(normalized_inputs):
      raise RuntimeError(
          "Ground truth should have same number of entries as inputs: {} vs {}"
          .format(len(self._data), len(normalized_inputs))
      )

    outputs: list[llmfn_outputs.LLMFnOutputEntry] = []
    for idx, (value, prompt_vars) in enumerate(
        zip(self._data, normalized_inputs)
    ):
      output_row = llmfn_output_row.LLMFnOutputRow(
          data={
              llmfn_outputs.ColumnNames.RESULT_NUM: 0,
              llmfn_outputs.ColumnNames.TEXT_RESULT: value,
          },
          result_type=str,
      )
      outputs.append(
          llmfn_outputs.LLMFnOutputEntry(
              prompt_num=0,
              input_num=idx,
              prompt_vars=prompt_vars,
              output_rows=[output_row],
          )
      )
    return outputs


def _get_ipython_display_fn(
    env: ipython_env.IPythonEnv,
) -> Callable[[llmfn_outputs.LLMFnOutputs], None]:
  return lambda x: env.display(x.as_pandas_dataframe())


def create_llm_function(
    models: model_registry.ModelRegistry,
    env: ipython_env.IPythonEnv | None,
    parsed_args: parsed_args_lib.ParsedArgs,
    cell_content: str,
    post_processing_fns: Sequence[post_process_utils.ParsedPostProcessExpr],
) -> llm_function.LLMFunction:
  """Creates an LLMFunction from Command.execute() arguments."""
  prompts: list[str] = [cell_content]

  llmfn_outputs_display_fn = _get_ipython_display_fn(env) if env else None

  llm_fn = llm_function.LLMFunctionImpl(
      model=models.get_model(parsed_args.model_type),
      model_args=parsed_args.model_args,
      prompts=prompts,
      outputs_ipython_display_fn=llmfn_outputs_display_fn,
  )
  if parsed_args.unique:
    llm_fn = llm_fn.add_post_process_reorder_fn(
        name="unique", fn=unique_fn.unique_fn
    )
  for fn in post_processing_fns:
    llm_fn = fn.add_to_llm_function(llm_fn)

  return llm_fn


def _convert_simple_compare_fn(
    name_and_simple_fn: tuple[str, Callable[[str, str], Any]]
) -> tuple[str, llm_function.CompareFn]:
  simple_fn = name_and_simple_fn[1]
  new_fn = lambda x, y: simple_fn(x.result_value(), y.result_value())
  return name_and_simple_fn[0], new_fn


def create_llm_compare_function(
    env: ipython_env.IPythonEnv | None,
    parsed_args: parsed_args_lib.ParsedArgs,
    post_processing_fns: Sequence[post_process_utils.ParsedPostProcessExpr],
) -> llm_function.LLMFunction:
  """Creates an LLMCompareFunction from Command.execute() arguments."""
  llmfn_outputs_display_fn = _get_ipython_display_fn(env) if env else None

  llm_cmp_fn = llm_function.LLMCompareFunction(
      lhs_name_and_fn=parsed_args.lhs_name_and_fn,
      rhs_name_and_fn=parsed_args.rhs_name_and_fn,
      compare_name_and_fns=[
          _convert_simple_compare_fn(x) for x in parsed_args.compare_fn
      ],
      outputs_ipython_display_fn=llmfn_outputs_display_fn,
  )
  for fn in post_processing_fns:
    llm_cmp_fn = fn.add_to_llm_function(llm_cmp_fn)

  return llm_cmp_fn


def create_llm_eval_function(
    models: model_registry.ModelRegistry,
    env: ipython_env.IPythonEnv | None,
    parsed_args: parsed_args_lib.ParsedArgs,
    cell_content: str,
    post_processing_fns: Sequence[post_process_utils.ParsedPostProcessExpr],
) -> llm_function.LLMFunction:
  """Creates an LLMCompareFunction from Command.execute() arguments."""
  llmfn_outputs_display_fn = _get_ipython_display_fn(env) if env else None

  # First construct a regular LLMFunction from the cell contents.
  llm_fn = create_llm_function(
      models=models,
      env=env,
      parsed_args=parsed_args,
      cell_content=cell_content,
      post_processing_fns=post_processing_fns,
  )

  # Next create a LLMCompareFunction.
  ground_truth_fn = _GroundTruthLLMFunction(data=parsed_args.ground_truth)
  llm_cmp_fn = llm_function.LLMCompareFunction(
      lhs_name_and_fn=("actual", llm_fn),
      rhs_name_and_fn=("ground_truth", ground_truth_fn),
      compare_name_and_fns=[
          _convert_simple_compare_fn(x) for x in parsed_args.compare_fn
      ],
      outputs_ipython_display_fn=llmfn_outputs_display_fn,
  )

  return llm_cmp_fn