from typing import Dict, Any, List, Optional
from llmops.base import Step
from llmops.base import ExecutionAnalysisType, ExecutionJudgement
from llmops.agents.operation_selector import OperationSelector
from llmops.agents.request_generator import RequestGenerator
from llmops.business.business_expert import BusinessExpert
from llmops.storage.transition_store import TransitionStore


class LLMAgent:
    def __init__(self, business_expert: BusinessExpert, model: str, transition_store: Optional[TransitionStore] = None):
        self.operation_selector = OperationSelector(business_expert, model, transition_store)
        self.request_generator = RequestGenerator(business_expert, model)
        self.trajectory: List[Dict[str, Any]] = []

    def get_action(
        self, current_state: Dict[str, Any], target_operation, feedback: dict
    ) -> Step:
        if (
            isinstance(feedback.get("execution_analyzer"), ExecutionJudgement)
            and feedback.get("execution_analyzer").result_type
            == ExecutionAnalysisType.REQUEST_ERROR
        ):
            selected_operation_id = self.trajectory[-1]["output"]["selected_operation"][
                "operation_id"
            ]
        else:
            selected_operation = self.operation_selector.select_operation(
                current_state, target_operation, feedback
            )
            selected_operation_id = selected_operation.operation_id.value
        complete_step = self.request_generator.generate_step(
            selected_operation_id,
            current_state,
            target_operation.model_dump(),
            feedback,
            selected_operation.reason,
        )
        self.trajectory.append(
            {
                "input": {
                    "current_state": current_state,
                    "target_operation": target_operation.model_dump(),
                    "feedback": feedback,
                },
                "output": {
                    "selected_operation": selected_operation.model_dump(),
                    "step": complete_step.model_dump(),
                },
            }
        )
        return complete_step