from enum import Enum
from typing import Optional
from pydantic import BaseModel
from llmops.utils.llm_client import LlmClient, LlmConfig

from llmops.business.business_expert import BusinessExpert
from llmops.config import prompt_path
from llmops.storage.transition_store import TransitionStore


class OperationSelection4GPT(BaseModel):
    reason: str
    operation_id: Enum

    def model_dump(self, **kwargs):
        data = super().model_dump(**kwargs)
        data["operation_id"] = self.operation_id.value
        return data


class OperationSelector:
    def __init__(self, business_expert: BusinessExpert, model: str, transition_store: Optional[TransitionStore] = None):
        self.system_prompt_template = prompt_path / "select_operation.md"
        self.business_expert = business_expert
        self.client = LlmClient(self.system_prompt_template, llm_config=LlmConfig(model=model))
        self.all_operations = business_expert.operation_ids
        self.trajectory = []
        self.transition_store = transition_store
        if business_expert.biz_system_prompt_operation_selector != "":
            self.client.update_system_prompt(
                self.client.system_prompt
                + "\nBusiness Logic:\n"
                + business_expert.biz_system_prompt_operation_selector
            )

    def _filter_operations(self, current_state, target_operation, feedback):
        # TODO: filter and return part of all_operations
        return self.all_operations

    def select_operation(self, current_state, target_operation, feedback):
        filtered_operations = self._filter_operations(
            current_state, target_operation, feedback
        )
        Selected_OperationID = Enum(
            "Selected_OperationID", {op: op for op in filtered_operations}
        )

        class FilteredOperationSelection4GPT(OperationSelection4GPT):
            reason: str
            operation_id: Selected_OperationID

        historical_data = None
        if self.transition_store:
            try:
                best_ops = self.transition_store.get_best_operations_for_state(
                    current_state, limit=3
                )
                if best_ops:
                    historical_data = {
                        "historical_successful_operations": [
                            {
                                "operation_id": op["operation_id"],
                                "success_rate": f"{op['success_rate']:.2%}",
                                "attempts": op["attempts"]
                            }
                            for op in best_ops
                        ],
                        "note": "These operations have historically succeeded from similar states. Consider them if they match your target."
                    }
            except Exception:
                pass

        user_prompt = {
            "current_state": current_state,
            "target_operation": target_operation.model_dump(),
            "feedback": feedback,
            "available_operations": self.business_expert.extract_ids_and_descriptions,
            "history": self.trajectory[-3:],
        }
        
        if historical_data:
            user_prompt["historical_data"] = historical_data

        result = self.client.generate(user_prompt, FilteredOperationSelection4GPT)

        self.trajectory.append(
            {
                "input": {
                    "current_state": current_state,
                    "target_operation": target_operation.model_dump(),
                    "feedback": feedback,
                },
                "output": result.model_dump(),
            }
        )

        return result

    def get_trajectory(self):
        return self.trajectory