from llmops.config import prompt_path
from llmops.base import ResourceList4GPT
from llmops.utils.llm_client import LlmClient, LlmConfig
from llmops.business.business_expert import BusinessExpert
from typing import Dict, Any
import copy


class ResourceManager:
    """Manage resources"""

    def __init__(self, business_expert: BusinessExpert, model: str):
        self.system_prompt_template = prompt_path / "resource_manager.md"
        self.business_expert = business_expert
        self.llm_client = LlmClient(self.system_prompt_template, llm_config=LlmConfig(model=model))
        self._resources = self.business_expert.list_all_resources()
        self.trajectory = []
        if business_expert.biz_system_prompt_resource_manager != "":
            self.llm_client.update_system_prompt(
                self.llm_client.system_prompt
                + "\nBusiness Logic:\n"
                + business_expert.biz_system_prompt_resource_manager
            )

    def update_resources(self, last_request, last_response) -> ResourceList4GPT:
        # Store the current resources before any updates
        previous_resources = copy.deepcopy(self._resources)
        self._resources = self.business_expert.list_all_resources()

        user_prompt = {
            "last_request": last_request,
            "last_response": last_response,
            "available_resource_types": list(self._resources.keys()),
            "previous_resources": previous_resources,
        }

        result = self.llm_client.generate(user_prompt, ResourceList4GPT)

        for resource in result.resources:
            if resource.resource_type not in self._resources:
                self._resources[resource.resource_type] = {}
            
            self._resources[resource.resource_type][resource.id] = {
                "attributes": resource.attributes,
                "state": resource.state,
                "references": resource.reference_ids,
                "parent_id": resource.parent_id,
            }

        # Record the trajectory with the correct previous_resources
        self.trajectory.append(
            {
                "input": {
                    "last_request": last_request,
                    "last_response": last_response,
                    "previous_resources": previous_resources,
                },
                "output": result.model_dump(),
            }
        )
        return self._resources

    def list_all_resources(self) -> Dict[str, Any]:
        return self._resources

    def cleanup(self):
        pass

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