# Copyright 2025 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

import logging
import re
from typing import ClassVar
from typing import Optional

from google.genai import types as genai_types
from pydantic import ValidationError
from typing_extensions import override

from ...models.base_llm import BaseLlm
from ...models.llm_request import LlmRequest
from ...models.llm_response import LlmResponse
from ...models.registry import LLMRegistry
from ...utils.context_utils import Aclosing
from ...utils.feature_decorator import experimental
from .._retry_options_utils import add_default_retry_options_if_not_present
from ..eval_case import ConversationScenario
from ..eval_case import Invocation
from ..eval_metrics import BaseCriterion
from ..eval_metrics import EvalMetric
from ..eval_metrics import EvalStatus
from ..eval_metrics import LlmBackedUserSimulatorCriterion
from ..evaluator import EvaluationResult
from ..evaluator import Evaluator
from ..evaluator import PerInvocationResult
from ..llm_as_judge import AutoRaterScore
from ..llm_as_judge_utils import get_eval_status
from ..llm_as_judge_utils import get_text_from_content
from ..llm_as_judge_utils import Label

_LATEST_TURN_USER_SIMULATOR_EVALUATOR_PROMPT = """
You are a data scientist tasked with evaluating the quality of a User Simulator that is interacting with an Agent.
Your task is to determine if the Generated User Response is consistent with:
  - The Conversation Plan: A list of high-level goals that the User Simulator is expected to achieve in the conversation.
  - The Conversation History: The exchange between the User Simulator and the Agent so far.
To determine this, we provide specific Evaluation Criteria that must be satisfied by the Generated User Response.

# Definition of Conversation Plan
The Conversation Plan specifies the goals that the User Simulator must execute.
The Conversation Plan also specifies the information and details that are needed to complete the goals.
The Conversation Plan is sequential in nature and the User Simulator must ensure the sequence is followed.

# Definition of Conversation History
The Conversation History is the actual dialogue between the User Simulator and the Agent.
The Conversation History may not be complete, but the exsisting dialogue should adhere to the Conversation Plan.
The Conversation History may contain instances where the User Simulator troubleshoots an incorrect/inappropriate response from the Agent in order to enforce the Conversation Plan.
The Conversation History is finished only when the User Simulator outputs `{stop_signal}` in its response. If this token is missing, the conversation between the User Simulator and the Agent has not finished, and more turns can be generated.

# Definition of Generated User Response
The Generated User Response is a the next user response in the conversation between a User Simulator and an Agent.
The Generated User Response was generated by the User Simulator based on a Conversation Plan and Conversation History.

# Evaluation Criteria
Your task is to evaluate the Generated User Response on a PASS/FAIL basis looking for specific errors.
The Generated User Response is marked as PASS unless it contains any of the Violations listed below, in which case it is marked as FAIL.

** CONVERSATION_PLAN_FOLLOWED **
Does the Generated User Response stick to the Conversation Plan?

Mark as FAIL if any of the following Violations occur:
- The Generated User Response repeats a high-level goal that was already completed in previous turns.
- The Generated User Response provides details for a high-level goal that was already completed.
- The Generated User Response response agrees to change the topic or perform a task not listed in the Conversation Plan.
- The Generated User Response invents a new goal not present in the Conversation Plan.
- The Generated User Response invents details (e.g., a made-up phone number or address) not provided in the Conversation Plan.

** STOP_CONDITION_FOLLOWED **
Did the conversation end exactly when it was supposed to?

Mark as FAIL if any of the following Violations occur:
- The conversation should have ended, but the Generated User Response did not use `{stop_signal}`.
- The Generated User Response used `{stop_signal}`, but tasks in the Conversation Plan are still incomplete AND the Agent has not failed.
- The Agent successfully transferred the User Simulator to a human/live agent, but the Generated User Response continued instead of using `{stop_signal}`.

** USER_GOAL_ORIENTED **
Is the User Simulator acting naturally, or is it "data dumping"?

Mark as FAIL if any of the following Violations occur:
- The Generated User Response provides specific details for a high-level goal (email content, recipient address, phone numbers) BEFORE the Agent has explicitly asked for them.
- The Generated User Response tries to accomplish more than one high-level task in a single turn.

** LIMITED_TROUBLESHOOTING **
Does the User Simulator have the correct amount of patience? (Note: Please check the conversation history and count the number of Agent errors).

Mark as FAIL if any of the following Violations occur:
- The Generated User Response ends the conversation immediately after the first Agent error.
- On the second Agent error, the Generated User Response response continues the conversation without using `{stop_signal}`.
- After the second Agent error, the Generated User Response tries to continue the conversation or continues addressing errors without using `{stop_signal}`.

** RESPONSIVENESS **
Does the User Simulator answer what is asked?

Mark as FAIL if any of the following Violations occur:
- The Agent asked a question (or multiple questions), and the Generated User Response failed to address one or all of them.
- The Agent asked for information NOT in the Conversation Plan, and the Generated User Response made up an answer instead of stating, e.g., "I don't know" or "I don't have that info."

** CORRECTS_AGENT  **
Does the User Simulator catch the Agent's mistakes?

Mark as FAIL if any of the following Violations occur:
- The Agent provided incorrect information, but the Generated User Response continued as if it was correct.
- The Agent made a dangerous assumption (e.g., sending an email without asking for the content first), and the Generated User Response continues without correcting the Agent.

** CONVERSATIONAL_TONE **
Does the User Simulator sound like a human?

Mark as FAIL if any of the following Violations occur:
- The Generated User Response uses overly complex sentence structures, or uses technical jargon inappropriately.
- The Generated User Response is sterile and purely functional (direct commands) with no natural conversational framing.
- The Generated User Response is too formal in nature, employing overly polite phrases and expressions.
- The Generated User Response is a "wall of text" where a simple sentence would suffice.

# Output Format
Format your response in the following JSON format:
{{
    "criteria": [
        {{
          "name": "CRITERIA_NAME_1",
          "reasoning": "reasoning",
          "passes": True or False,
        }},
        {{
          "name": "CRITERIA_NAME_2",
          "reasoning": "reasoning",
          "passes": True or False,
        }},
        ...
    ],
    "is_valid": True or False,
}}

# Conversation Plan
{conversation_plan}

# Conversation History
{conversation_history}

# Generated User Response
{generated_user_response}
""".strip()


def _parse_llm_response(response: str) -> Label:
  """Parses the LLM response and extracts the final label.

  Args:
    response: LLM response.

  Returns:
    The extracted label, either VALID, INVALID, or NOT_FOUND.
  """
  # Regex matching the label field in the response.
  is_valid_match = re.search(
      r'"is_valid":\s*\[*[\n\s]*"*([^"^\]^\s]*)"*[\n\s]*\]*\s*[,\n\}]',
      response,
  )

  # If there was not match for "is_valid", return NOT_FOUND
  if is_valid_match is None:
    return Label.NOT_FOUND

  # Remove any trailing whitespace, commas, or end-brackets from the label.
  label = is_valid_match.group(1).strip(r"\s,\}").lower()
  if label in [
      Label.INVALID.value,
      Label.ALMOST.value,
      Label.FALSE.value,
      *Label.PARTIALLY_VALID.value,
  ]:
    return Label.INVALID
  elif label in [Label.VALID.value, Label.TRUE.value]:
    return Label.VALID
  else:
    return Label.NOT_FOUND


def _format_conversation_history(invocations: list[Invocation]) -> str:
  conversation_history = []
  for invocation in invocations:
    if invocation.user_content is not None:
      conversation_history.append(
          f"user: {get_text_from_content(invocation.user_content)}"
      )

    final_response = invocation.final_response
    if final_response is not None:
      conversation_history.append(
          f"{final_response.role}: {get_text_from_content(final_response)}"
      )
  return "\n\n".join(conversation_history)


def _get_stop_signal_invocation(stop_signal: str) -> Invocation:
  return Invocation(
      invocation_id="stop_signal_proxy_invocation",
      user_content=genai_types.Content(
          parts=[genai_types.Part(text=stop_signal)]
      ),
  )


@experimental
class PerTurnUserSimulatorQualityV1(Evaluator):
  """Per turn user simulator evaluator.

  This evaluator verifies that the conversation from a user simulator sticks
  to the given conversation scenario:
  - In the first turn, it verifies that the user simulator output the
    specified starting prompt.
  - For all the other turns, it verifies that the user simulator stuck to the
    conversation plan.
  - It also verifies that the user simulator finished the conversation
    appropriately.
  This evaluator uses an LLM to verify all turns except the first one. It
  aggregates repeated invocation samples by taking majority vote. The overall
  score is the fraction of turns of the conversation before the verifier
  detects an issue with the user simulator.
  """

  criterion_type: ClassVar[type[LlmBackedUserSimulatorCriterion]] = (
      LlmBackedUserSimulatorCriterion
  )

  def __init__(
      self,
      eval_metric: EvalMetric,
  ):
    self._eval_metric = eval_metric
    self._criterion = self._deserialize_criterion(eval_metric)

    self._prompt_template = _LATEST_TURN_USER_SIMULATOR_EVALUATOR_PROMPT

    self._llm_options = self._criterion.judge_model_options
    self._stop_signal = self._criterion.stop_signal
    self._llm = self._setup_llm()

  def _deserialize_criterion(self, eval_metric: EvalMetric) -> BaseCriterion:
    expected_criterion_type_error = ValueError(
        f"`{eval_metric.metric_name}` metric expects a criterion of type"
        f" `{self.criterion_type}`."
    )
    try:
      if self._eval_metric.criterion is None:
        raise expected_criterion_type_error

      return self.criterion_type.model_validate(
          self._eval_metric.criterion.model_dump()
      )
    except ValidationError as e:
      raise expected_criterion_type_error from e

  @override
  async def evaluate_invocations(
      self,
      actual_invocations: list[Invocation],
      expected_invocations: Optional[list[Invocation]] = None,
      conversation_scenario: Optional[ConversationScenario] = None,
  ) -> EvaluationResult:
    del expected_invocations  # not used by this metric.
    if conversation_scenario is None:
      raise ValueError("conversation_scenario is needed by this metric.")

    # Evaluate the first invocation contains the given starting prompt.
    results = [
        self._evaluate_first_turn(actual_invocations[0], conversation_scenario)
    ]

    # Evaluate the rest of the invocations.
    for i, invocation in enumerate(actual_invocations):
      # skip the first invocation.
      if i == 0:
        continue

      result = await self._evaluate_intermediate_turn(
          invocation_at_step=invocation,
          invocation_history=actual_invocations[:i],
          conversation_scenario=conversation_scenario,
      )
      results.append(result)

    if not results:
      return EvaluationResult()

    # Evaluate whether the conversation ended correctly.
    stop_signal_evaluation = await self._evaluate_stop_signal_turn(
        invocation_history=actual_invocations,
        conversation_scenario=conversation_scenario,
    )

    # If the conversation did not end correctly, indicate so by marking the
    # last user turn as failed.
    if stop_signal_evaluation.eval_status == EvalStatus.FAILED:
      results[-1] = stop_signal_evaluation

    return self._aggregate_conversation_results(results)

  def _setup_llm(self) -> BaseLlm:
    model_id = self._llm_options.judge_model
    llm_registry = LLMRegistry()
    llm_class = llm_registry.resolve(model_id)
    return llm_class(model=model_id)

  def _format_llm_prompt(
      self,
      invocation: Invocation,
      conversation_scenario: ConversationScenario,
      previous_invocations: Optional[list[Invocation]],
  ) -> str:
    if previous_invocations is None:
      raise ValueError(
          "Previous invocations should have a set value when "
          "formatting the LLM prompt. "
          f"Encountered: {previous_invocations}"
      )

    if conversation_scenario is None:
      raise ValueError(
          "Conversation scenario should have a set value when "
          "formatting the LLM prompt. "
          f"Encountered: {conversation_scenario}"
      )

    return self._prompt_template.format(
        conversation_plan=conversation_scenario.conversation_plan,
        conversation_history=_format_conversation_history(previous_invocations),
        generated_user_response=get_text_from_content(invocation.user_content),
        stop_signal=self._stop_signal,
    )

  def _convert_llm_response_to_score(
      self, auto_rater_response: LlmResponse
  ) -> AutoRaterScore:
    response_text = get_text_from_content(auto_rater_response.content)
    if response_text is None or not response_text:
      return AutoRaterScore()
    label = _parse_llm_response(response_text)

    if label == Label.VALID:
      return AutoRaterScore(score=1.0)
    elif label == Label.INVALID:
      return AutoRaterScore(score=0.0)
    else:
      return AutoRaterScore()

  def _aggregate_samples(
      self,
      per_invocation_samples: list[PerInvocationResult],
  ) -> PerInvocationResult:
    """Aggregates samples by taking majority vote."""
    if not per_invocation_samples:
      raise ValueError("No samples to aggregate into a result.")

    positive_results = [s for s in per_invocation_samples if s.score == 1.0]
    negative_results = [s for s in per_invocation_samples if s.score == 0.0]

    if not positive_results and not negative_results:
      return per_invocation_samples[0]
    elif len(positive_results) > len(negative_results):
      return positive_results[0]
    else:  # len(negative_results) >= len(positive_results)
      return negative_results[0]

  def _aggregate_conversation_results(
      self, per_invocation_results: list[PerInvocationResult]
  ) -> EvaluationResult:
    """Computes the fraction of results that resulted in a pass status."""
    num_valid = 0
    num_evaluated = 0
    for result in per_invocation_results:
      if result.eval_status == EvalStatus.PASSED:
        num_valid += result.score

      num_evaluated += 1

    # If no invocation was evaluated, we mark the score as None.
    if num_evaluated == 0:
      return EvaluationResult(
          per_invocation_results=per_invocation_results,
      )

    overall_score = num_valid / num_evaluated
    return EvaluationResult(
        overall_score=overall_score,
        overall_eval_status=get_eval_status(
            overall_score, self._criterion.threshold
        ),
        per_invocation_results=per_invocation_results,
    )

  def _evaluate_first_turn(
      self,
      first_invocation: Invocation,
      conversation_scenario: ConversationScenario,
  ) -> PerInvocationResult:
    if first_invocation.user_content is None:
      return PerInvocationResult(
          actual_invocation=first_invocation,
          eval_status=EvalStatus.NOT_EVALUATED,
      )

    score = int(
        get_text_from_content(first_invocation.user_content).strip()
        == conversation_scenario.starting_prompt.strip()
    )
    return PerInvocationResult(
        actual_invocation=first_invocation,
        score=score,
        eval_status=get_eval_status(score, self._eval_metric.threshold),
    )

  async def _evaluate_intermediate_turn(
      self,
      invocation_at_step: Invocation,
      invocation_history: list[Invocation],
      conversation_scenario: Optional[ConversationScenario],
  ) -> PerInvocationResult:

    auto_rater_prompt = self._format_llm_prompt(
        invocation=invocation_at_step,
        conversation_scenario=conversation_scenario,
        previous_invocations=invocation_history,
    )

    llm_request = LlmRequest(
        model=self._llm_options.judge_model,
        contents=[
            genai_types.Content(
                parts=[genai_types.Part(text=auto_rater_prompt)],
                role="user",
            )
        ],
        config=self._llm_options.judge_model_config,
    )
    add_default_retry_options_if_not_present(llm_request)
    num_samples = self._llm_options.num_samples
    samples = []
    for _ in range(num_samples):
      llm_score = await self._sample_llm(llm_request)
      samples.append(
          PerInvocationResult(
              eval_status=get_eval_status(
                  llm_score.score, self._eval_metric.threshold
              ),
              score=llm_score.score,
              actual_invocation=invocation_at_step,
          )
      )
    if not samples:
      return PerInvocationResult(
          eval_status=EvalStatus.NOT_EVALUATED,
          actual_invocation=invocation_at_step,
      )

    return self._aggregate_samples(samples)

  async def _evaluate_stop_signal_turn(
      self,
      invocation_history: list[Invocation],
      conversation_scenario: ConversationScenario,
  ) -> PerInvocationResult:
    return await self._evaluate_intermediate_turn(
        invocation_at_step=_get_stop_signal_invocation(self._stop_signal),
        invocation_history=invocation_history,
        conversation_scenario=conversation_scenario,
    )

  async def _sample_llm(self, llm_request: LlmRequest) -> AutoRaterScore:
    async with Aclosing(self._llm.generate_content_async(llm_request)) as agen:
      async for llm_response in agen:
        # Non-streaming call, so there is only one response content.
        return self._convert_llm_response_to_score(llm_response)
