from typing import Dict, Any

from llmops.config import prompt_path
from llmops.base import ExecutionJudgement
from llmops.utils.llm_client import LlmClient, LlmConfig
from llmops.business.business_expert import BusinessExpert


class ExecutionAnalyzer:
    def __init__(self, business_expert: BusinessExpert, model: str):
        self.system_prompt_template = prompt_path / "execution_analyzer.md"
        self.client = LlmClient(self.system_prompt_template, llm_config=LlmConfig(model=model))
        self.business_expert = business_expert
        self.trajectory = []

    def evaluate(
        self,
        request: Dict[str, Any],  # exclude auth
        response: Dict[str, Any],
        prev_state: Dict[str, Any],
        current_state: Dict[str, Any],
        target_operation: Dict[str, Any],
        operation_id: str,
    ) -> ExecutionJudgement:
        """Evaluate the execution result"""
        # TODO: consider status_code / exit_code? Differentiate responsibility when non-200 / 0
        # Prepare prompt
        if request.get("auth"):
            # 如果请求中包含 auth, 移除
            request.pop("auth")
        operation = self.business_expert.extract_operation(operation_id)
        user_prompt = {
            "request": request,
            "response": response,
            "prev_state": prev_state,
            "current_state": current_state,
            "target_operation": target_operation,
            "available_operations": self.business_expert.extract_ids_and_descriptions,
            "operation_spec": operation,
        }
        result = self.client.generate(user_prompt, ExecutionJudgement)
        self.trajectory.append(
            {
                "input": {
                    "request": request,
                    "response": response,
                    "prev_state": prev_state,
                    "current_state": current_state,
                    "target_operation": target_operation,
                    "available_operations": self.business_expert.extract_ids_and_descriptions,
                    "operation_spec": operation,
                },
                "output": result.model_dump(),
            }
        )
        return result

    def get_trajectory(self) -> Dict[str, Any]:
        """Get the trajectory of execution analysis process"""
        return self.trajectory