# 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

import json
import logging
from typing import ClassVar
from typing import Optional

from typing_extensions import override

from .eval_case import ConversationScenario
from .eval_case import Invocation
from .eval_case import InvocationEvents
from .eval_metrics import EvalMetric
from .eval_metrics import EvalStatus
from .eval_metrics import RubricsBasedCriterion
from .evaluator import EvaluationResult
from .evaluator import PerInvocationResult
from .rubric_based_evaluator import RubricBasedEvaluator

logger = logging.getLogger("google_adk." + __name__)

_RUBRIC_BASED_MULTI_TURN_TRAJECTORY_QUALITY_V1_PROMPT = """# Mission
Your mission is to evaluate the quality of responses generated by an AI agent in a multi-turn conversation. Your task is to analyze the entire conversation history, focusing on all assistant responses, and rate them against the provided rubric criteria. The conversation may include text-based interactions, transcribed audio interactions, and records of tool calls made by the assistant.
One turn is defined as a user and assistant response pair.

# Instructions:
- Analyze the entire conversation: Carefully review every turn in the <conversation_history>, paying attention to each user query and assistant response turn within <conversation_history>.
- Evaluate ALL assistant responses cumulatively: Your assessment for each criterion should reflect the assistant's overall performance across all its turns. For instance, if an assistant made a factual error in Turn 2 but corrected it in Turn 4, the "Factual Accuracy" score should reflect this nuanced performance.
- Consider Modality and Tool Use: Take into account the different modalities presented (text, audio transcripts) and the effective or ineffective use of tools by the assistant as described in the conversation history.
- Refer to the Rubric: For each criterion listed under <properties>, determine how well the assistant's collective performance across all its turns satisfies that criterion.
- If a specific property is not fulfilled, you must provide the agent response turn number that violated the property.
- Your mission is to evaluate the quality of responses generated by an AI agent. You will be presented with the conversation history (<conversation_history>) which includes a set of user and assistant turns, and a set of properties (<property>) that you must use to objectively assess the validity of the agent's response.
- Only use the properties provided. Do not make up new properties.
- Keep the critiques across each property mutually exclusive.
- IMPORTANT: Assess all of the provided properties. Do not drop any of the properties from your response.

# Rubric:
"yes": The agent's response fulfilled the property or the property is not applicable to the user prompt.
"no": The agent's response did not fulfill the property.

# For each property starting with a new line, follow these steps:
STEP 1: Repeat the property, word for word, without making any changes. Keep everything including punctuation and capitalization as-is.
STEP 2: Determine the steps needed to check if all the assistant responses in the conversation history fulfill the *intent* of the property. Refer back to the `<user_prompt>` to understand the user's original goal.
STEP 3: Follow the steps outlined in STEP 2, thinking out loud. As you think, refer to specific assistant responses and count their turn numbers within the <conversation_history>.
STEP 4: Review the thoughts and the original property.
STEP 5: Output the final verdict.
Property: [[Repeat the property in STEP 1 again.]]
Rationale: [[Explain your reasoning for the verdict.]]
Verdict: [[yes|no]]

# Output format (repeat this format for every property started with a new line):
STEP 1: ...
STEP 2: ...
STEP 3: ...
STEP 4: ...
STEP 5: ...
Property: ...
Rationale: ...
Verdict: ...

# Example output 1

STEP 1: Does the agent run function call 'default_api.grammar_check'?
STEP 2: I need to check if any assistant response in the <conversation_history> includes the exact function name 'default_api.grammar_check'.
STEP 3: Reviewing the <conversation_history>, I find that in Assistant Turn 2, the agent includes the function call default_api.grammar_check(sentence="the dog walks on the a park"). This matches the required function name.
STEP 4: The assistant successfully called the specified function in Turn 2.
STEP 5: yes
Property: Does the agent run function call 'default_api.grammar_check'?
Rationale: The agent's response in Assistant Turn 2 contains the function call 'default_api.grammar_check' within a proper code block and with the correct function name.
Verdict: yes

STEP 1: The sentence parameter in the grammar_check function call must be 'the dog walks on the a park'.
STEP 2: I need to check if the function call 'default_api.grammar_check' (if present in any assistant turn) includes the parameter 'sentence' and whether the value assigned to 'sentence' is valid according to the provided guideline in the property. The property specifies the valid value as 'the dog walks on the a park'.
STEP 3: In Assistant Turn 2, the agent's response includes the function call default_api.grammar_check(sentence="the dog walks on the a park"). The parameter 'sentence' is present, and the value assigned to it is "the dog walks on the a park", which is identical to the reference value.
STEP 4: The parameter 'sentence' is present and its value is exactly the same as the reference value in Assistant Turn 2.
STEP 5: yes
Property: The sentence parameter in the grammar_check function call must be 'the dog walks on the a park'.
Rationale: The agent's response in Assistant Turn 2 includes the 'sentence' parameter in the function call 'default_api.grammar_check', and the value assigned to it is exactly the same as the reference value.
Verdict: yes

# Example output 2

STEP 1: Does the agent run function call 'default_api.search_via_perplexity'?
STEP 2: I need to check if any assistant response in the <conversation_history> includes the exact function name 'default_api.search_via_perplexity'.
STEP 3: Reviewing the <conversation_history>, I find that in Assistant Turn 2, the agent includes a function call default_api.get_web_search_results. This does not match the required 'default_api.search_via_perplexity'. I checked all other assistant turns and no other function call matches.
STEP 4: The agent did not use the function 'default_api.search_via_perplexity' in any turn.
STEP 5: no
Property: Does the agent run function call 'default_api.search_via_perplexity'?
Rationale: The agent called 'default_api.get_web_search_results' in Assistant Turn 2, not 'default_api.search_via_perplexity'. The correct function was not called in any turn.
Verdict: no

STEP 1: Does the agent provide function call 'default_api.search_via_perplexity' with input parameter 'keyword' that is valid compared to the reference 'keyword'= 'GPT-4o vs GPT-3.5 cost comparison' and based on the following guideline? Guideline for 'keyword': 'The wording can differ. The agent response is valid if it conveys similar core content as the reference response. Less efficient and minor inaccurate phrasing is acceptable.'
STEP 2: I need to check if the agent runs the function call with exact function name as 'default_api.search_via_perplexity' in any turn, and if so, check if the value assigned to 'keyword' is valid according to the provided guideline in the property.
STEP 3: Based on the previous evaluation, the agent did not include a function call 'default_api.search_via_perplexity' in any of its turns. Therefore, this property cannot be fulfilled.
STEP 4: The property cannot be fulfilled because the agent did not use the specified function call.
STEP 5: no
Property: Does the agent provide function call 'default_api.search_via_perplexity' with input parameter 'keyword' that is valid compared to the reference 'keyword'= 'GPT-4o vs GPT-3.5 cost comparison' and based on the following guideline? Guideline for 'keyword': 'The wording can differ. The agent response is valid if it conveys similar core content as the reference response. Less efficient and minor inaccurate phrasing is acceptable.'
Rationale: The agent did not use the function call 'default_api.search_via_perplexity' in any of its responses throughout the conversation.
Verdict: no

<agent_system_instructions>
{agent_instructions}
</agent_system_instructions>

<agent_tool_definitions>
{agent_tool_definitions}
</agent_tool_definitions>

<conversation_history>
{user_agent_dialogue}
</conversation_history>

<properties>
{properties}
</properties>
"""


class RubricBasedMultiTurnTrajectoryEvaluator(RubricBasedEvaluator):
  """An Evaluator for rubric based assessment of an agent's multi-turn trajectory using a LLM.

  This evaluator assesses the quality of an agent's behavior across an entire
  multi-turn conversation, including intermediate tool calls, function
  responses, and the final textual replies. It uses a set of rubrics (defined
  as properties) to judge whether the agent's cumulative performance satisfies
  each criterion.

  Unlike single-turn evaluators that assess each invocation independently, this
  evaluator accumulates the full dialogue history (user turns, agent turns, and
  tool interactions) and performs a single LLM-based evaluation on the last
  turn using the complete conversation context. The first N-1 turns are marked
  as NOT_EVALUATED, and the final turn carries the aggregate score.

  Example: For a travel-booking agent that handles multi-step itinerary
  planning, one could specify the following rubrics:

  Rubric 1: The agent correctly called the flight search tool with the
  user-specified origin, destination, and dates.
  Rubric 2: The agent confirmed the booking details with the user before
  finalizing the reservation.
  Rubric 3: The agent's final response included a complete itinerary summary
  with flight numbers, times, and confirmation codes.

  For each rubric, the LLM judge will produce a binary verdict ("yes" or "no")
  indicating whether the agent's trajectory satisfied that property. A "yes"
  verdict maps to a score of 1.0 and "no" maps to 0.0.

  A combined score using individual rubric verdicts will also be generated.
  This overall score is the average of all individual rubric scores, giving a
  value between 0 and 1 where values closer to 1 indicate better adherence to
  the specified rubrics.
  """

  criterion_type: ClassVar[type[RubricsBasedCriterion]] = RubricsBasedCriterion
  RUBRIC_TYPE: ClassVar[str] = "TRAJECTORY_QUALITY"

  def __init__(self, eval_metric: EvalMetric):
    super().__init__(
        eval_metric,
        criterion_type=RubricBasedMultiTurnTrajectoryEvaluator.criterion_type,
        rubric_type=RubricBasedMultiTurnTrajectoryEvaluator.RUBRIC_TYPE,
    )
    self._auto_rater_prompt_template = (
        _RUBRIC_BASED_MULTI_TURN_TRAJECTORY_QUALITY_V1_PROMPT
    )

  def _assemble_dialogue_history(
      self, actual_invocations: list[Invocation]
  ) -> None:
    """Assembles the dialogue history, instructions, and tools from invocations."""
    dialogue_lines = []
    for turn_index, invocation in enumerate(actual_invocations):
      # USER TURN
      if invocation.user_content and invocation.user_content.parts:
        text_parts = [p.text for p in invocation.user_content.parts if p.text]
        if text_parts:
          dialogue_lines.append(
              f"USER TURN {turn_index + 1}: {' '.join(text_parts)}"
          )

      # INTERMEDIATE TOOL EVENTS
      if isinstance(invocation.intermediate_data, InvocationEvents):
        for event in invocation.intermediate_data.invocation_events:
          role = (
              "USER"
              if event.author and event.author.lower() == "user"
              else f"AGENT ({event.author})"
          )
          if event.content and event.content.parts:
            text_parts = [p.text for p in event.content.parts if p.text]
            if text_parts:
              dialogue_lines.append(
                  f"{role} TURN {turn_index + 1}: {' '.join(text_parts)}"
              )
            for p in event.content.parts:
              if p.function_call:
                args = (
                    json.dumps(dict(p.function_call.args))
                    if p.function_call.args
                    else "{}"
                )
                dialogue_lines.append(
                    f"{role} TURN {turn_index + 1} (tool call):"
                    f" {p.function_call.name}({args})"
                )
              if p.function_response:
                resp = (
                    json.dumps(dict(p.function_response.response))
                    if p.function_response.response
                    else "{}"
                )
                dialogue_lines.append(
                    f"{role} TURN {turn_index + 1} (tool output):"
                    f" {p.function_response.name} -> {resp}"
                )

      # FINAL AGENT TURN
      if invocation.final_response and invocation.final_response.parts:
        try:
          agent_name = invocation.intermediate_data.invocation_events[0].author
        except (AttributeError, IndexError):
          agent_name = "agent"
        role = f"AGENT ({agent_name})"
        text_parts = [p.text for p in invocation.final_response.parts if p.text]
        if text_parts:
          dialogue_lines.append(
              f"{role} TURN {turn_index + 1}: {' '.join(text_parts)}"
          )

    self._formatted_dialogue = "\n".join(dialogue_lines)

    # Grab instructions and tools locally from Invocation
    instructions_parts = []
    tools_parts = []
    for invocation in actual_invocations:
      if invocation.app_details and invocation.app_details.agent_details:
        for agent_id, details in invocation.app_details.agent_details.items():
          instructions_parts.append(
              f"Agent {agent_id} Instructions:\n{details.instructions}"
          )
          tools_parts.append(f"Agent: {agent_id}")
          if details.tool_declarations:
            for tool in details.tool_declarations:
              for func in tool.function_declarations:
                tools_parts.append(f"- {func.name}: {func.description}")

    self._formatted_instructions = "\n\n".join(
        dict.fromkeys(instructions_parts)
    )
    self._formatted_tools = "\n".join(dict.fromkeys(tools_parts))

  @override
  async def evaluate_invocations(
      self,
      actual_invocations: list[Invocation],
      expected_invocations: Optional[list[Invocation]] = None,
      conversation_scenario: Optional[ConversationScenario] = None,
  ) -> EvaluationResult:
    """Evaluates multiple turns locally, collecting full dialogue history."""
    # Parse out Pydantic types locally without server-side imports
    # Re-implement dialogue parsing logic fully locally
    # to have zero deps on cloud.ai.
    logger.debug(
        "Local evaluator start (invocations: %d)",
        len(actual_invocations),
    )
    self._assemble_dialogue_history(actual_invocations)

    # If expected_invocations are not supplied, provide a list of None.
    expected_invocations = (
        [None] * len(actual_invocations)
        if expected_invocations is None
        else expected_invocations
    )

    # Mark the first N-1 turns as NOT_EVALUATED.
    per_invocation_results = []
    for actual, expected in zip(
        actual_invocations[:-1], expected_invocations[:-1]
    ):
      per_invocation_results.append(
          PerInvocationResult(
              actual_invocation=actual,
              expected_invocation=expected,
              score=None,
              eval_status=EvalStatus.NOT_EVALUATED,
          )
      )

    # Conversation-level evaluation: run the LLM judge
    # once on the last turn with full dialogue context.
    last_expected = (
        [expected_invocations[-1]] if expected_invocations[-1] else None
    )
    last_turn_result = await super().evaluate_invocations(
        [actual_invocations[-1]],
        last_expected,
        conversation_scenario,
    )

    # Append the evaluated last-turn result.
    if last_turn_result.per_invocation_results:
      per_invocation_results.extend(last_turn_result.per_invocation_results)
    else:
      per_invocation_results.append(
          PerInvocationResult(
              actual_invocation=actual_invocations[-1],
              expected_invocation=expected_invocations[-1],
              score=last_turn_result.overall_score,
              eval_status=last_turn_result.overall_eval_status,
              rubric_scores=last_turn_result.overall_rubric_scores,
          )
      )

    return EvaluationResult(
        overall_score=last_turn_result.overall_score,
        overall_eval_status=last_turn_result.overall_eval_status,
        per_invocation_results=per_invocation_results,
        overall_rubric_scores=last_turn_result.overall_rubric_scores,
    )

  @override
  def format_auto_rater_prompt(
      self,
      actual_invocation: Invocation,
      _: Optional[Invocation],
  ) -> str:
    """Returns the fully rendered validation prompt using locally configured rubrics."""
    self.create_effective_rubrics_list(actual_invocation.rubrics)
    logger.debug(
        "format_auto_rater_prompt called (effective rubrics: %d)",
        len(self._effective_rubrics_list),
    )

    rubrics_list = []
    for r in self._effective_rubrics_list:
      rubrics_dict = {
          "property": r.rubric_content.text_property,
      }
      if r.type:
        rubrics_dict["type"] = r.type
      rubrics_list.append(rubrics_dict)
    formatted_rubrics = json.dumps(rubrics_list, indent=2)

    prompt = self._auto_rater_prompt_template.format(
        user_agent_dialogue=self._formatted_dialogue,
        properties=formatted_rubrics,
        agent_instructions=self._formatted_instructions,
        agent_tool_definitions=self._formatted_tools,
    )
    logger.debug("Generated rubric validator prompt:\n%s", prompt)
    return prompt
