# 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 importlib.util
import logging
import os
import sys
from types import ModuleType
from typing import Any
from typing import cast
from typing import Optional

import click
from google.genai import types as genai_types

from ..agents.llm_agent import Agent
from ..evaluation.base_eval_service import BaseEvalService
from ..evaluation.base_eval_service import EvaluateConfig
from ..evaluation.base_eval_service import EvaluateRequest
from ..evaluation.base_eval_service import InferenceRequest
from ..evaluation.base_eval_service import InferenceResult
from ..evaluation.constants import MISSING_EVAL_DEPENDENCIES_MESSAGE
from ..evaluation.eval_case import get_all_tool_calls
from ..evaluation.eval_case import IntermediateDataType
from ..evaluation.eval_metrics import EvalMetric
from ..evaluation.eval_metrics import Interval
from ..evaluation.eval_metrics import MetricInfo
from ..evaluation.eval_metrics import MetricValueInfo
from ..evaluation.eval_result import EvalCaseResult
from ..evaluation.eval_sets_manager import EvalSetsManager
from ..utils.context_utils import Aclosing

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


TOOL_TRAJECTORY_SCORE_KEY = "tool_trajectory_avg_score"
RESPONSE_MATCH_SCORE_KEY = "response_match_score"
SAFETY_V1_KEY = "safety_v1"
FINAL_RESPONSE_MATCH_V2 = "final_response_match_v2"
# This evaluation is not very stable.
# This is always optional unless explicitly specified.
RESPONSE_EVALUATION_SCORE_KEY = "response_evaluation_score"

EVAL_SESSION_ID_PREFIX = "___eval___session___"
DEFAULT_CRITERIA = {
    TOOL_TRAJECTORY_SCORE_KEY: 1.0,  # 1-point scale; 1.0 is perfect.
    RESPONSE_MATCH_SCORE_KEY: 0.8,
}


def _import_from_path(module_name: str, file_path: str) -> ModuleType:
  spec = importlib.util.spec_from_file_location(module_name, file_path)
  if spec is None or spec.loader is None:
    raise ImportError(f"Cannot import module {module_name} from {file_path}")
  module = importlib.util.module_from_spec(spec)
  sys.modules[module_name] = module
  spec.loader.exec_module(module)
  return module


def _get_agent_module(agent_module_file_path: str) -> ModuleType:
  file_path = os.path.join(agent_module_file_path, "__init__.py")
  module_name = "agent"
  return _import_from_path(module_name, file_path)


def get_default_metric_info(
    metric_name: str, description: str = ""
) -> MetricInfo:
  """Returns a default MetricInfo for a metric."""
  return MetricInfo(
      metric_name=metric_name,
      description=description,
      metric_value_info=MetricValueInfo(
          interval=Interval(min_value=0.0, max_value=1.0)
      ),
  )


def get_root_agent(agent_module_file_path: str) -> Agent:
  """Returns root agent given the agent module."""
  agent_module = _get_agent_module(agent_module_file_path)
  root_agent = agent_module.agent.root_agent
  return cast(Agent, root_agent)


def try_get_reset_func(agent_module_file_path: str) -> Any:
  """Returns reset function for the agent, if present, given the agent module."""
  agent_module = _get_agent_module(agent_module_file_path)
  reset_func = getattr(agent_module.agent, "reset_data", None)
  return reset_func


def parse_and_get_evals_to_run(
    evals_to_run_info: list[str],
) -> dict[str, list[str]]:
  """Returns a dictionary of eval set info to evals that should be run.

  Args:
    evals_to_run_info: While the structure is quite simple, a list of string,
      each string actually is formatted with the following convention:
      <eval_set_file_path | eval_set_id>:[comma separated eval case ids]
  """
  eval_set_to_evals: dict[str, list[str]] = {}
  for input_eval_set in evals_to_run_info:
    evals = []
    if ":" not in input_eval_set:
      # We don't have any eval cases specified. This would be the case where the
      # the user wants to run all eval cases in the eval set.
      eval_set = input_eval_set
    else:
      # There are eval cases that we need to parse. The user wants to run
      # specific eval cases from the eval set.
      eval_set = input_eval_set.split(":")[0]
      evals = input_eval_set.split(":")[1].split(",")
      evals = [s for s in evals if s.strip()]

    if eval_set not in eval_set_to_evals:
      eval_set_to_evals[eval_set] = []

    eval_set_to_evals[eval_set].extend(evals)

  return eval_set_to_evals


async def _collect_inferences(
    inference_requests: list[InferenceRequest],
    eval_service: BaseEvalService,
) -> list[InferenceResult]:
  """Simple utility methods to collect inferences from an eval service.

  The method is intentionally kept private to prevent general usage.
  """
  inference_results = []
  for inference_request in inference_requests:
    async with Aclosing(
        eval_service.perform_inference(inference_request=inference_request)
    ) as agen:
      async for inference_result in agen:
        inference_results.append(inference_result)
  return inference_results


async def _collect_eval_results(
    inference_results: list[InferenceResult],
    eval_service: BaseEvalService,
    eval_metrics: list[EvalMetric],
) -> list[EvalCaseResult]:
  """Simple utility methods to collect eval results from an eval service.

  The method is intentionally kept private to prevent general usage.
  """
  eval_results = []
  evaluate_request = EvaluateRequest(
      inference_results=inference_results,
      evaluate_config=EvaluateConfig(eval_metrics=eval_metrics),
  )
  async with Aclosing(
      eval_service.evaluate(evaluate_request=evaluate_request)
  ) as agen:
    async for eval_result in agen:
      eval_results.append(eval_result)

  return eval_results


def _convert_content_to_text(
    content: Optional[genai_types.Content],
) -> str:
  if content and content.parts:
    return "\n".join([p.text for p in content.parts if p.text])
  return ""


def _convert_tool_calls_to_text(
    intermediate_data: Optional[IntermediateDataType],
) -> str:
  tool_calls = get_all_tool_calls(intermediate_data)
  return "\n".join([str(t) for t in tool_calls])


def pretty_print_eval_result(eval_result: EvalCaseResult) -> None:
  """Pretty prints eval result."""
  try:
    import pandas as pd
    from tabulate import tabulate
  except ModuleNotFoundError as e:
    raise ModuleNotFoundError(MISSING_EVAL_DEPENDENCIES_MESSAGE) from e

  click.echo(f"Eval Set Id: {eval_result.eval_set_id}")
  click.echo(f"Eval Id: {eval_result.eval_id}")
  click.echo(f"Overall Eval Status: {eval_result.final_eval_status.name}")

  for metric_result in eval_result.overall_eval_metric_results:
    click.echo(
        "---------------------------------------------------------------------"
    )
    click.echo(
        f"Metric: {metric_result.metric_name}, "
        f"Status: {metric_result.eval_status.name}, "
        f"Score: {metric_result.score}, "
        f"Threshold: {metric_result.threshold}"
    )
    if metric_result.details and metric_result.details.rubric_scores:
      click.echo("Rubric Scores:")
      rubrics_by_id = {
          r["rubric_id"]: r["rubric_content"]["text_property"]
          for r in metric_result.criterion.rubrics
      }
      for rubric_score in metric_result.details.rubric_scores:
        rubric_text = rubrics_by_id.get(rubric_score.rubric_id)
        if not rubric_text:
          rubric_text = rubric_score.rubric_id
        click.echo(
            f"Rubric: {rubric_text}, "
            f"Score: {rubric_score.score}, "
            f"Reasoning: {rubric_score.rationale}"
        )

  data = []
  for per_invocation_result in eval_result.eval_metric_result_per_invocation:
    actual_invocation = per_invocation_result.actual_invocation
    expected_invocation = per_invocation_result.expected_invocation
    row_data = {
        "prompt": _convert_content_to_text(actual_invocation.user_content),
        "expected_response": (
            _convert_content_to_text(expected_invocation.final_response)
            if expected_invocation
            else None
        ),
        "actual_response": _convert_content_to_text(
            actual_invocation.final_response
        ),
        "expected_tool_calls": (
            _convert_tool_calls_to_text(expected_invocation.intermediate_data)
            if expected_invocation
            else None
        ),
        "actual_tool_calls": _convert_tool_calls_to_text(
            actual_invocation.intermediate_data
        ),
    }
    for metric_result in per_invocation_result.eval_metric_results:
      row_data[metric_result.metric_name] = (
          f"Status: {metric_result.eval_status.name}, "
          f"Score: {metric_result.score}"
      )
      if metric_result.details and metric_result.details.rubric_scores:
        rubrics_by_id = {
            r["rubric_id"]: r["rubric_content"]["text_property"]
            for r in metric_result.criterion.rubrics
        }
        for rubric_score in metric_result.details.rubric_scores:
          rubric = rubrics_by_id.get(rubric_score.rubric_id)
          if not rubric:
            rubric = rubric_score.rubric_id
          row_data[f"Rubric: {rubric}"] = (
              f"Reasoning: {rubric_score.rationale}, "
              f"Score: {rubric_score.score}"
          )
    data.append(row_data)
  if data:
    click.echo(
        "---------------------------------------------------------------------"
    )
    click.echo("Invocation Details:")
    df = pd.DataFrame(data)

    # Identify columns where ALL values are exactly None
    columns_to_keep = []
    for col in df.columns:
      # Check if all elements in the column are NOT None
      if not df[col].apply(lambda x: x is None).all():
        columns_to_keep.append(col)

    # Select only the columns to keep
    df_result = df[columns_to_keep]

    for col in df_result.columns:
      if df_result[col].dtype == "object":
        df_result[col] = df_result[col].str.wrap(40)

    click.echo(
        tabulate(df_result, headers="keys", tablefmt="grid", maxcolwidths=25)
    )
    click.echo("\n\n")  # Few empty lines for visual clarity


def get_eval_sets_manager(
    eval_storage_uri: Optional[str], agents_dir: str
) -> EvalSetsManager:
  """Returns an instance of EvalSetsManager."""
  try:
    from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager
    from .utils import evals
  except ModuleNotFoundError as mnf:
    raise click.ClickException(MISSING_EVAL_DEPENDENCIES_MESSAGE) from mnf

  if eval_storage_uri:
    gcs_eval_managers = evals.create_gcs_eval_managers_from_uri(
        eval_storage_uri
    )
    return gcs_eval_managers.eval_sets_manager
  else:
    return LocalEvalSetsManager(agents_dir=agents_dir)
