# 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 typing_extensions import override

from ..models.llm_response import LlmResponse
from ..utils.feature_decorator import experimental
from .common import EvalBaseModel
from .eval_case import Invocation
from .eval_case import InvocationEvents
from .eval_metrics import EvalMetric
from .eval_metrics import Interval
from .eval_metrics import MetricInfo
from .eval_metrics import MetricValueInfo
from .eval_metrics import PrebuiltMetrics
from .eval_metrics import RubricsBasedCriterion
from .eval_rubrics import Rubric
from .eval_rubrics import RubricScore
from .evaluator import EvaluationResult
from .evaluator import PerInvocationResult
from .llm_as_judge import AutoRaterScore
from .llm_as_judge import LlmAsJudge
from .llm_as_judge_utils import get_average_rubric_score
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 get_tool_calls_and_responses_as_json_str
from .llm_as_judge_utils import get_tool_declarations_as_json_str

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

_RUBRIC_BASED_FINAL_RESPONSE_QUALITY_V1_PROMPT = """
SPECIAL INSTRUCTION: think silently. Silent thinking token budget: 10240 tokens.

# Mission
Your mission is to evaluate the final answer quality of responses generated by an AI agent. You will be presented with a user prompt (<user_prompt>), the agent's response (<response>) to that user prompt, and a set of properties (<property>) that you must use to objectively assess the validity of the agent's response.
Only respond to the properties provided. Do not make up new properties.

# Rubric
"yes": The model's response fulfilled the property, OR the property's condition was not applicable to the response.
"no": The model's response met the conditions for the property to be applicable, but failed to fulfill it, or the property applies to a claim in the model's response that cannot be unambiguously verified using trusted evidence.

# Key Evaluation Principles
Your evaluation must follow a two-part process: first, collect trusted evidence from the agent's work, and second, judge the final answer against it.
1. **Establish Trusted Evidence from Tool Calls**: You must first examine the agent's tool calls to determine if they are procedurally sound, meaning that the agent used the appropriate tools with logical parameters to address the user's prompt.
  * Your ONLY sources of truth are the <user_prompt> and the direct output ('tool_response') from PROCEDURALLY SOUND tool calls found in the <response_steps>. Examples of procedural flaws include:
    * The agent failed to call a tool that will enable it to answer the user's prompt despite having all the necessary parameters to do so.
    * The agent called the tool with incorrect or missing parameters.
    * The agent called a tool that does not exist, or called a tool with a parameter that does not exist.
    * The agent's sequence of tool calls contains a logical error.
  * The following kinds of information ABSOLUTELY CANNOT BE USED to derive trusted evidence:
    * The agent's final answer.
    * The agent's reasoning, summaries, or any interpretations of the tool responses by the agent.
    * Any tool call that is flawed (e.g., queries the wrong file, contains incorrect logic).
  * You may not have access to the same tools as the agent, so do not attempt to call any tools yourself.
2. **Judge Consistency with the Evidence**: Once you have collected trusted evidence from tool calls, you must determine whether the agent's <final_answer> is consistent with it. A claim in the final answer is only considered correct if it can be unambiguously verified using this evidence.
  * If the necessary evidence is missing because the agent failed to make a correct and sound tool call, the final answer must be judged as failing the property.

While judging the final answer against the evidence, be flexible about how it is conveyed. Accept answers that are semantically equivalent (e.g., different phrasing) as long as they still fulfill the property. For numbers, accept answers that are numerically equivalent, allowing for minor differences in rounding or precision, as long as they do not alter a final conclusion (e.g., the outcome of a statistical test).

For each property follow these internal steps:
1. Understand the property and the key evaluation principles.
2. Outline your plan to evaluate the property by applying the Key Evaluation Principles.
3. Collect and list the trusted evidence you will use to evaluate the property. Note any procedural flaws in the tool calls.
4. Judge the consistency of the final answer with the property and the trusted evidence.
5. Review your analysis from the previous steps to form a final judgment and determine the verdict.
6. Output the final verdict in the required output format.

# Output Format (repeat this format for every property, starting with a new line):
Property: [Repeat the property, word for word, without making any changes. Keep everything including punctuation and capitalization as-is.]
Evidence: [List all trusted evidence from tool calls or the user prompt that is relevant to the property (referencing the Step Index). Alternatively, if either no trusted evidence is required, or no trusted evidence exists (e.g., flawed process, missing tool call, tool error), explain why.]
Rationale: [Explain your reasoning, detailing how the evidence (or lack thereof) supports or contradicts the final answer, or why the property is not applicable.]
Verdict: [yes|no]

REMEMBER: Your answer will help improve the AI agent. It is important to determine the fulfillment of the properties correctly. Even answering "no" will improve the agent! Respond in pure text, not json.

# Example
## Input
<user_prompt>
  <developer_instructions>
  You are an AI agent who is an expert in HR data analysis.
  If a company has fewer than 100 employees, then the final answer should alert the user that there are fewer than 100 employees.
  If you have sufficient information and tools to respond to the user's question, then do not ask for further clarification.
  </developer_instructions>
  <available_tools>
  {{
    'name': 'load_hr_data_from_file',
    'description': 'Reads a data file from the company's HR database into a Pandas DataFrame.'
    'parameters': [
        {{
          'type': 'string',
          'name': 'file_name',
          'description': 'The name of the data file.'
        }},
    ],
    'required': ['file_name']
  }},
  {{
    'name': 'get_manager',
    'description': 'Returns the manager of a given employee.',
    'parameters': [
        {{
          'type': 'string',
          'name': 'employee_name',
          'description': 'The name of the employee.'
        }},
    ],
    'required': ['employee_name']
  }}
  </available_tools>
  <main_prompt>
  Using the employees.csv file, determine:
  1. the total number of employees
  2. the name of Alice Smith's manager
  3. the name of the employee with the highest salary, and their gender
  4. the average salary for the "Marketing" department
  Please format your final answer as a numbered list.
  </main_prompt>
</user_prompt>
<response>
  <response_steps>
  [
    {{
      "step_index": 0,
      "tool_call": "df = load_hr_data_from_file('employees.csv')\nprint(len(df))",
      "tool_response": "110",
    }},
    {{
      "step_index": 1,
      "tool_call": "print(df[df['Department'] == 'Engineering']['Salary'].mean())",
      "tool_response": "155000",
    }},
    {{
      "step_index": 2,
      "tool_call="print(df.loc[df['Salary'].idxmax(), 'Name'])",
      "tool_response": "John Smith",
    }},
  ]
  </response_steps>
  <final_answer>
  1. The total number of employees is 110.
  2. Please provide Alice Smith's employee ID so that I can find her manager.
  3. The employee with the highest salary is John Doe, and this employee's gender is male.
  4. The average salary for the Marketing department is 155000.
  </final_answer>
</response>

<properties>
* The final answer correctly identifies the total number of employees.
* The final answer correctly identifies the name of Alice Smith's manager, or correctly states that it cannot be determined and why.
* The final answer correctly states the average salary for the Marketing department.
* The final answer correctly identifies the employee with the highest salary.
* The final answer correctly identifies the gender of the employee with the highest salary, or correctly states that it cannot be determined and why.
* The final answer is formatted as a numbered list.
* If the company has fewer than 100 employees, then the final answer states that it has fewer than 100 employees.
</properties>

## Output
Property: The final answer correctly identifies the total number of employees.
Evidence: The trusted evidence is "110 employees". The tool call in Step 0 is procedurally sound and provides the total number of employees (110) by calling the load_hr_data_from_file tool with the correct file name.
Rationale: The final answer's claim ("110 employees") is fully consistent with the trusted evidence.
Verdict: yes

Property: The final answer correctly identifies the name of Alice Smith's manager, or correctly states that it cannot be determined and why.
Evidence: No trusted evidence exists. The agent did not perform a tool call to determine the manager of Alice Smith, despite having the necessary information (the employee name) and access to the necessary tools (get_manager) to do so.
Rationale: The agent incorrectly stated that the final answer cannot be determined, despite having the necessary information (the employee name) and tools (get_manager) to determine it.
Verdict: no

Property: The final answer correctly states the average salary for the Marketing department.
Evidence: No trusted evidence exists for the Marketing department's average salary. The tool call in Step 1 is procedurally flawed; the agent searched for "Engineering" instead of "Marketing".
Rationale: There is no trusted evidence for the Marketing department's average salary.
Verdict: no

Property: The final answer correctly identifies the employee with the highest salary.
Evidence: The trusted evidence is "John Smith". The tool call in Step 2 produces trusted evidence for the employee with the highest salary by calling the load_hr_data_from_file tool with the correct file name and then using the idxmax() method to find the employee with the highest salary.
Rationale: The final answer's claim ("John Doe") is inconsistent with the trusted evidence ("John Smith").
Verdict: no

Property: The final answer correctly identifies the gender of the employee with the highest salary, or correctly states that it cannot be determined and why.
Evidence: No trusted evidence exists. The agent did not perform a tool call to determine the gender of the employee with the highest salary.
Rationale: There is no trusted evidence to confirm the gender of the employee with the highest salary that the final answer states (male). Even if the gender is coincidentally actually male, the claim in the final answer cannot be unambiguously verified using the evidence.
Verdict: no

Property: If the company has fewer than 100 employees, then the final answer should state that it has fewer than 100 employees.
Evidence: The trusted evidence is "110 employees". The tool call in Step 0 correctly counts the total number of employees as 110 by calling the load_hr_data_from_file tool with the correct file name.
Rationale: The total number of employees is 110, so the condition for this property (fewer than 100 employees) was not met. Therefore, the property is not applicable to this response.
Verdict: yes

Property: The final answer is formatted as a numbered list.
Evidence: N/A. Trusted evidence from tool calls or the user prompt is not required in order to determine the format of the final answer.
Rationale: The final answer is formatted as a numbered list from 1 to 4, e.g. "1. The total number of employees is 110\n2...".
Verdict: yes

# Your Turn
## Input
<user_prompt>
  <developer_instructions>
  {developer_instructions}
  </developer_instructions>

  <available_tools>
  {tool_declarations}
  </available_tools>

  <main_prompt>
  {user_input}
  </main_prompt>
</user_prompt>

<response>
  <response_steps>
  {response_steps}
  </response_steps>
  <final_answer>
  {final_response}
  </final_answer>
</response>

<properties>
{rubrics}
</properties>

## Output
"""


_PROPERTY_PATTERN = r"(?<=Property: )(.*)"
_RATIONALE_PATTERN = r"(?<=Rationale: )(.*)"
_VERDICT_PATTERN = r"(?<=Verdict: )(.*)"


class _RubricResponse(EvalBaseModel):
  """Internal data model to represent a rubric's response from the auto-rater."""

  property_text: Optional[str] = None
  rationale: Optional[str] = None
  score: Optional[float] = None


def _normalize_text(text: str) -> str:
  """Returns a normalized version of the passed in text."""
  if not isinstance(text, str):
    return ""
  return text.lower().strip()


def _parse_auto_rater_response(
    auto_rater_response: str,
) -> list[_RubricResponse]:
  """Returns a list of _RubricResponse parsed from the AutoRater's response."""
  properties = re.findall(_PROPERTY_PATTERN, auto_rater_response)
  rationales = re.findall(_RATIONALE_PATTERN, auto_rater_response)
  scores = []

  for verdict in re.findall(_VERDICT_PATTERN, auto_rater_response):
    if "yes" in verdict.lower():
      score = 1.0
    elif "no" in verdict.lower():
      score = 0.0
    else:
      score = None

    scores.append(score)

  rubric_responses = []
  for p, r, s in zip(properties, rationales, scores):
    rubric_responses.append(
        _RubricResponse(property_text=p.strip(), rationale=r.strip(), score=s)
    )

  return rubric_responses


@experimental
class RubricBasedFinalResponseQualityV1Evaluator(LlmAsJudge):
  """An Evaluator for rubric based assessment of the agent's final response using a LLM.

  The evaluator uses a set of rubrics to assess the quality of the agent's
  final response.

  Example: For a weather agent that responds to weather related queries of the
  user, one could specify following rubrics:

  Rubric 1: Agent's response is direct and to the point.
  Rubric 2: Agent's response accurately inferred user's underlying goal from
  ambiguous queries (e.g. "is it a beach weather?" would mean sun, warmth and
  low wind)

  For each rubric, this evaluator will generate a confidence score between 0
  and 1, where 0 means that agent's response did not satisfy the rubric at all
  and 1 means complete adherence. Value closer to 1 are desirable.

  A combined score using individual rubric confidences will also be generated.
  Like individual rubric confidence scores, the range for this value will be
  between 0 and 1, and it will have the same interpretation.
  """

  criterion_type: ClassVar[type[RubricsBasedCriterion]] = RubricsBasedCriterion

  def __init__(self, eval_metric: EvalMetric):
    super().__init__(
        eval_metric,
        criterion_type=RubricBasedFinalResponseQualityV1Evaluator.criterion_type,
    )
    self._auto_rater_prompt_template = (
        _RUBRIC_BASED_FINAL_RESPONSE_QUALITY_V1_PROMPT
    )

    assert self._criterion.rubrics, "Rubrics are required."

    self._rubrics: list[Rubric] = self._criterion.rubrics

    self._normalized_rubric_to_id_map = {
        _normalize_text(r.rubric_content.text_property): r.rubric_id
        for r in self._rubrics
    }

  @staticmethod
  def get_metric_info() -> MetricInfo:
    return MetricInfo(
        metric_name=PrebuiltMetrics.RUBRIC_BASED_FINAL_RESPONSE_QUALITY_V1.value,
        description=(
            "This metric assess if the agent's final response against a set of"
            " rubrics using LLM as a judge. Value range for this metric is"
            " [0,1], with values closer to 1 more desirable."
        ),
        metric_value_info=MetricValueInfo(
            interval=Interval(min_value=0.0, max_value=1.0)
        ),
    )

  @override
  def format_auto_rater_prompt(
      self, actual_invocation: Invocation, _: Invocation
  ) -> str:
    """Returns the autorater prompt."""

    user_input = get_text_from_content(actual_invocation.user_content)
    final_response = get_text_from_content(actual_invocation.final_response)
    rubrics = "\n*  ".join(
        [r.rubric_content.text_property for r in self._rubrics]
    )

    developer_instructions = ""
    tool_declarations = "Agent has no tools."
    response_steps = get_tool_calls_and_responses_as_json_str(
        actual_invocation.intermediate_data
    )

    app_details = actual_invocation.app_details
    if app_details:
      if (
          isinstance(actual_invocation.intermediate_data, InvocationEvents)
          and actual_invocation.intermediate_data.invocation_events
      ):
        developer_instructions = app_details.get_developer_instructions(
            agent_name=actual_invocation.intermediate_data.invocation_events[
                0
            ].author
        )
      tool_declarations = get_tool_declarations_as_json_str(app_details)

    auto_rater_prompt = self._auto_rater_prompt_template.format(
        developer_instructions=developer_instructions,
        tool_declarations=tool_declarations,
        user_input=user_input,
        response_steps=response_steps,
        final_response=final_response,
        rubrics=rubrics,
    )

    return auto_rater_prompt

  @override
  def convert_auto_rater_response_to_score(
      self, auto_rater_response: LlmResponse
  ) -> AutoRaterScore:
    """Returns an AutoRaterScore generated from AutoRater's response."""
    response_text = get_text_from_content(auto_rater_response.content)
    rubric_responses = _parse_auto_rater_response(response_text)
    rubric_scores = []

    for rubric_response in rubric_responses:
      normalized_rubric = _normalize_text(rubric_response.property_text)
      rubric_id = self._normalized_rubric_to_id_map.get(normalized_rubric, None)
      if rubric_id:
        rubric_scores.append(
            RubricScore(
                rubric_id=rubric_id,
                rationale=rubric_response.rationale,
                score=rubric_response.score,
            )
        )
      else:
        logger.warning(
            f"Rubric {rubric_response.property_text} not found in the rubrics"
            " provided to the metric."
        )

    aggregated_score = get_average_rubric_score(rubric_scores)
    return AutoRaterScore(score=aggregated_score, rubric_scores=rubric_scores)

  @override
  def aggregate_per_invocation_samples(
      self,
      per_invocation_samples: list[PerInvocationResult],
  ) -> PerInvocationResult:
    """Returns a combined result for the invocation.

    This AutoRater is backed by an LLM that are known to have certain degree of
    unreliabilty to their responses. In order to counter that we sample the
    autorater more than once for a single invocation.

    This method takes all those samples for a single invocation and combines
    them to genreate one single result for the invocation.

    This method specifically uses majority vote to aggregate scores for a
    rubric. Take following Invocation and Rubric for example:

      Invocation:
         User: Is it going to be cold in Seattle tomorrow?
         Weather Agent: No, it will be moderately warm as predicted temperature
         for Seattle, WA tomorrow is 88F.

      Rubric: Agent's reponse was concise and to the point.

      We will sample the AutoRater 5 times, and the AutoRater responds
      with (skipping the rationale field for now):
        Sample 1:
          Verdict: Yes
        Sample 2:
          Verdict: No
        Sample 3:
          Verdict: Yes
        Sample 4:
          Verdict: Yes
        Sample 5:
          Verdict: No

      This method will use majority vote and combine the results of 5 samples
      into one, and it will report "Yes" as the final verdict.
    """
    score_category_by_rubric_id = {}

    # We go over each rubric for each sample, and categorize the rubric into
    # one of the following buckets:
    #  - Bucket 0: No score was generated for the rubric
    #  - Bucket 1: Score was generated and it was positive (1.0)
    #  - Bucket 2: Score was generated and it was negative (0.0)
    for sample in per_invocation_samples:
      if not sample.rubric_scores:
        continue

      for rubric_score in sample.rubric_scores:
        rubric_id = rubric_score.rubric_id
        if rubric_id not in score_category_by_rubric_id:
          score_category_by_rubric_id[rubric_id] = ([], [], [])

        if rubric_score.score is None:  # No score
          score_category_by_rubric_id[rubric_id][0].append(rubric_score)
        elif rubric_score.score == 1.0:  # Positive Result
          score_category_by_rubric_id[rubric_id][1].append(rubric_score)
        else:  # Negative result
          score_category_by_rubric_id[rubric_id][2].append(rubric_score)

    aggregated_rubric_scores = []
    for rubric_id in score_category_by_rubric_id:
      no_scores, positives, negatives = score_category_by_rubric_id[rubric_id]

      if not positives and not negatives:
        # There has to be at least a no score rubric!
        aggregated_rubric_scores.append(no_scores[0])

      # This is where we are taking a majority vote.
      elif len(positives) > len(negatives):
        aggregated_rubric_scores.append(positives[0])
      else:
        aggregated_rubric_scores.append(negatives[0])

    aggregated_overall_score = get_average_rubric_score(
        aggregated_rubric_scores
    )

    return PerInvocationResult(
        actual_invocation=per_invocation_samples[0].actual_invocation,
        expected_invocation=per_invocation_samples[0].expected_invocation,
        score=aggregated_overall_score,
        rubric_scores=aggregated_rubric_scores,
        eval_status=get_eval_status(
            aggregated_overall_score, self._eval_metric.threshold
        ),
    )

  @override
  def aggregate_invocation_results(
      self, per_invocation_results: list[PerInvocationResult]
  ) -> EvaluationResult:
    """Aggregates per invocation evaluation results into a single score.

    A single eval case can have multiple invocations and the eval metric is
    assessed for each invocation. But, we do want to make an aggregate
    statement on how the eval case as a whole performed on the metric.

    This method helps us aggreate rubric scores across invocation.

    Do note that the aggregation strategy used here is different from the one
    that is used in `aggregate_per_invocation_samples` method, where we used
    majority vote. In this method, we actually calculate the mean score of a
    rubric across several invocations, as majority score would be misleading.
    """

    unaggregated_rubric_scores = []  # Later used to calculate average.

    # Collect rubric scores by id, so that we can calculate average score
    # for each rubric id.
    rubric_scores_by_id = {}
    for sample in per_invocation_results:
      if not sample.rubric_scores:
        continue

      for rubric_score in sample.rubric_scores:
        rubric_id = rubric_score.rubric_id
        if rubric_id not in rubric_scores_by_id:
          rubric_scores_by_id[rubric_id] = []

        rubric_scores_by_id[rubric_id].append(rubric_score)
        unaggregated_rubric_scores.append(rubric_score)

    aggregated_rubric_scores = []
    for rubric_id, rubric_scores in rubric_scores_by_id.items():
      overall_score = get_average_rubric_score(rubric_scores)
      aggregated_rubric_scores.append(
          RubricScore(
              rubric_id=rubric_id,
              score=overall_score,
              # There is no real way for us generate a rationale here, so we
              # make is clear to the consumer of the result.
              rationale=(
                  "This is an aggregated score derived from individual entries."
                  " Please refer to individual entries in each invocation for"
                  " actual rationale from the model."
              ),
          )
      )

    # Use unaggregate rubric score to calculate overall score.
    aggregated_overall_score = get_average_rubric_score(
        unaggregated_rubric_scores
    )
    return EvaluationResult(
        overall_score=aggregated_overall_score,
        overall_eval_status=get_eval_status(
            aggregated_overall_score, self._eval_metric.threshold
        ),
        per_invocation_results=per_invocation_results,
        overall_rubric_scores=aggregated_rubric_scores,
    )
