"""Trajectory evaluator for analyzing execution paths and decision sequences."""

from typing import Any, Optional

from opentelemetry.sdk.trace import ReadableSpan
from pydantic import field_validator

from uipath.eval.models import EvaluationResult

from ..._utils.constants import COMMUNITY_agents_SUFFIX
from ...platform.chat import UiPathLlmChatService
from ...platform.chat.llm_gateway import RequiredToolChoice
from ..models.models import (
    AgentExecution,
    LLMResponse,
    NumericEvaluationResult,
    TrajectoryEvaluationTrace,
)
from .legacy_base_evaluator import (
    LegacyBaseEvaluator,
    LegacyEvaluationCriteria,
    LegacyEvaluatorConfig,
)
from .legacy_llm_helpers import create_evaluation_tool, extract_tool_call_response


class LegacyTrajectoryEvaluatorConfig(LegacyEvaluatorConfig):
    """Configuration for legacy trajectory evaluators."""

    name: str = "LegacyTrajectoryEvaluator"


class LegacyTrajectoryEvaluator(LegacyBaseEvaluator[LegacyTrajectoryEvaluatorConfig]):
    """Legacy evaluator that analyzes the trajectory/path taken to reach outputs."""

    prompt: str
    model: str
    expected_agent_behavior_placeholder: str = "{{ExpectedAgentBehavior}}"
    agent_run_history_placeholder: str = "{{AgentRunHistory}}"
    llm: Optional[UiPathLlmChatService] = None

    @field_validator("prompt")
    @classmethod
    def validate_prompt_placeholder(cls, v: str) -> str:
        """Validate that prompt contains required placeholders."""
        if "{{ExpectedAgentBehavior}}" not in v or "{{AgentRunHistory}}" not in v:
            raise ValueError(
                "Prompt must contain {ExpectedAgentBehavior} and {{AgentRunHistory}} placeholders"
            )
        return v

    def model_post_init(self, __context: Any):
        """Initialize the evaluator after model creation."""
        super().model_post_init(__context)

    def _initialize_llm(self):
        """Initialize the LLM used for evaluation."""
        from uipath.platform import UiPath

        uipath = UiPath()
        self.llm = uipath.llm

    async def evaluate(
        self,
        agent_execution: AgentExecution,
        evaluation_criteria: LegacyEvaluationCriteria,
    ) -> EvaluationResult:
        """Evaluate using trajectory analysis.

        Analyzes the execution path and decision sequence taken by the agent.

        Args:
            agent_execution: The execution details containing:
                - agent_input: The input received by the agent
                - actual_output: The actual output from the agent
                - agent_trace: The execution spans to use for the evaluation
                - expected_agent_behavior: The expected agent behavior
            evaluation_criteria: The criteria to evaluate
        Returns:
            EvaluationResult: Score based on trajectory analysis

        Raises:
            NotImplementedError: This evaluator is not yet implemented
        """
        # Lazily initialize the LLM on first evaluation call
        if self.llm is None:
            self._initialize_llm()

        evaluation_prompt = self._create_evaluation_prompt(
            expected_agent_behavior=agent_execution.expected_agent_behavior,
            agent_run_history=agent_execution.agent_trace,
        )
        llm_response = await self._get_llm_response(evaluation_prompt)

        return NumericEvaluationResult(
            score=llm_response.score,
            details=llm_response.justification,
        )

    def _create_evaluation_prompt(
        self,
        expected_agent_behavior: Any,
        agent_run_history: Any,
    ) -> str:
        """Create the evaluation prompt for the LLM."""
        formatted_prompt = self.prompt.replace(
            self.expected_agent_behavior_placeholder,
            str(expected_agent_behavior),
        )

        # Trim extra properties from the spans (such as timestamps which are not relevant to the eval)
        if (
            isinstance(agent_run_history, list)
            and agent_run_history
            and isinstance(agent_run_history[0], ReadableSpan)
        ):
            trajectory_trace = TrajectoryEvaluationTrace.from_readable_spans(
                agent_run_history
            )
            agent_run_history = str(trajectory_trace.spans)
        else:
            agent_run_history = str(agent_run_history)

        formatted_prompt = formatted_prompt.replace(
            self.agent_run_history_placeholder,
            agent_run_history,
        )

        return formatted_prompt

    async def _get_llm_response(self, evaluation_prompt: str) -> LLMResponse:
        """Get response from the LLM using universal function calling.

        Args:
            evaluation_prompt: The formatted prompt to send to the LLM

        Returns:
            LLMResponse with score and justification
        """
        assert self.llm, "LLM should be initialized before calling this method."

        model = self.model
        if model.endswith(COMMUNITY_agents_SUFFIX):
            model = model.replace(COMMUNITY_agents_SUFFIX, "")

        # Create evaluation tool for function calling (works across all models)
        evaluation_tool = create_evaluation_tool()
        tool_choice = RequiredToolChoice()

        # Prepare the request with function calling
        request_data = {
            "model": model,
            "messages": [{"role": "user", "content": evaluation_prompt}],
            "tools": [evaluation_tool],
            "tool_choice": tool_choice,
        }

        response = await self.llm.chat_completions(**request_data)
        return extract_tool_call_response(response, model)
