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

from typing import Optional

from .user_simulator_personas import UserPersona

_LATEST_TURN_USER_SIMULATOR_EVALUATOR_PROMPT_TEMPLATE = """
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 existing 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()


_LATEST_TURN_USER_SIMULATOR_WITH_PERSONA_EVALUATOR_PROMPT_TEMPLATE = """
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.
  - A Persona: A set of behaviours that the User Simulator is expected to exhibit in the conversation.
To determine this, we provide specific Evaluation Criteria that you must use to evaluate 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.
The Conversation Plan is not a script.

# 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 existing 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 Persona
The Persona is a description of how the User Simulator should behave in a conversation with the Agent.
A Persona specifies behaviors, not goals.
If the Persona contradicts the Conversation Plan, the Conversation Plan has precedence.

# Definition of Generated User Response
The Generated User Response is 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 the 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.
{% for b in persona.behaviors %}
## Criteria: {{ b.name | render_string_filter}}
{{ b.description | render_string_filter}}

Mark as FAIL if any of the following Violations occur:
{{ b.get_violation_rubrics_str() | render_string_filter}}
{% endfor %}
# 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 if it passes all criteria, False otherwise
}

# Conversation Plan
{{ conversation_plan }}

# Conversation History
{{ conversation_history }}

# Persona Description
{{ persona.description }}
The Evaluation Criteria above already specify how to evaluate whether the Generated User Response satisfies this persona.

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


def _get_latest_turn_user_simulator_quality_prompt_template(
    user_persona: Optional[UserPersona] = None,
) -> str:
  """Returns the appropriate prompt for user simulator quality"""
  if user_persona is None:
    return _LATEST_TURN_USER_SIMULATOR_EVALUATOR_PROMPT_TEMPLATE
  return _LATEST_TURN_USER_SIMULATOR_WITH_PERSONA_EVALUATOR_PROMPT_TEMPLATE


def get_per_turn_user_simulator_quality_prompt(
    conversation_plan: str,
    conversation_history: str,
    generated_user_response: str,
    stop_signal: str,
    user_persona: Optional[UserPersona] = None,
):
  """Formats the prompt for the per turn user simulator evaluator"""
  from jinja2 import DictLoader
  from jinja2 import pass_context
  from jinja2.sandbox import SandboxedEnvironment

  templates = {
      "verifier_instructions": (
          _get_latest_turn_user_simulator_quality_prompt_template(
              user_persona=user_persona
          )
      ),
  }
  template_env = SandboxedEnvironment(loader=DictLoader(templates))

  @pass_context
  def _render_string_filter(context, template_string):
    if not template_string:
      return ""
    return template_env.from_string(template_string).render(context.get_all())

  template_env.filters["render_string_filter"] = _render_string_filter

  template_parameters = {
      "conversation_plan": conversation_plan,
      "conversation_history": conversation_history,
      "generated_user_response": generated_user_response,
      "stop_signal": stop_signal,
  }
  if user_persona is not None:
    template_parameters["persona"] = user_persona

  return template_env.get_template("verifier_instructions").render(
      template_parameters
  )
