from utils import get_relevant_history, get_embedding
import torch
from LLM.base_LLM import *
from Memory import Memory
from Prompt import * 
import json
class Environment:
    """
    The place where the agent activities, responsible for storing some shared memories
    """
    def __init__(self, config) -> None:
        self.shared_memory = {"long_term_memory": [], "short_term_memory": None}
        self.agents = None

        self.summary_system_prompt = {}
        self.summary_last_prompt = {}
        self.environment_prompt = {}
        self.environment_type = config["environment_type"] if "environment_type" in config else "cooperative"
        self.current_chat_history_idx = 0
        self.LLMs = {}
        
        # 初始化每个state 的summary 方法
        # Initialize the summary method for each state
        for state_name, state_dict in config["states"].items():
            if state_name != "end_state":
                self.summary_system_prompt[state_name] = (
                    state_dict["summary_system_prompt"]
                    if "summary_system_prompt" in state_dict
                    else eval(Default_environment_summary_system_prompt)
                )

                self.summary_last_prompt[state_name] = (
                    state_dict["summary_last_prompt"]
                    if "summary_last_prompt" in state_dict
                    else eval(Default_environment_summary_last_prompt)
                )

                self.environment_prompt[state_name] = (
                    state_dict["environment_prompt"]
                    if "environment_prompt" in state_dict
                    else " "
                )
                self.LLMs[state_name] = init_LLM(f"logs/{state_name}",**state_dict)
        self.roles_to_names = None
        self.names_to_roles = None

    @classmethod
    def from_config(cls, config_path):
        with open(config_path) as f:
            config = json.load(f)
        return cls(config)

    def summary(self, current_state):
        """
        Summarize the situation in the current environment every once in a while
        """
        MAX_CHAT_HISTORY = eval(os.environ["MAX_CHAT_HISTORY"])
        current_state_name = current_state.name

        query = self.shared_memory["long_term_memory"][-1].content
        relevant_history = get_relevant_history(
            query,
            self.shared_memory["long_term_memory"][:-1],
            self.shared_memory["chat_embeddings"][:-1],
        )

        relevant_history = Memory.get_chat_history(relevant_history)
        chat_history = Memory.get_chat_history(
            self.shared_memory["long_term_memory"][-MAX_CHAT_HISTORY + 1 :]
        )
        summary = self.shared_memory["short_term_memory"]
        
        
        # system prompt = environment prompt + current memory + system prompt
        # current_memory = summary + chat history + relevant history
        current_memory = eval(Environment_summary_memory)
        environment_prompt = self.environment_prompt[current_state_name]
        summary_system_prompt = self.summary_system_prompt[current_state_name]
        
        environment_summary_system_prompt = eval(Environment_summary_system_prompt)
        response = self.LLMs[current_state_name].get_response(None, environment_summary_system_prompt, stream=False)
        return response

    def update_memory(self, memory, current_state):
        """
        update chat embbedings and long term memory,short term memory,agents long term memory
        """
        MAX_CHAT_HISTORY = eval(os.environ["MAX_CHAT_HISTORY"])
        self.shared_memory["long_term_memory"].append(memory)
        current_embedding = get_embedding(memory.content)
        if "chat_embeddings" not in self.shared_memory:
            self.shared_memory["chat_embeddings"] = current_embedding
        else:
            self.shared_memory["chat_embeddings"] = torch.cat(
                [self.shared_memory["chat_embeddings"], current_embedding], dim=0
            )
        if len(self.shared_memory["long_term_memory"]) % MAX_CHAT_HISTORY == 0:
            summary = self.summary(current_state)
            self.shared_memory["short_term_memory"] = summary

        self.agents[memory.send_name].update_memory(memory)
    
    
    def _get_agent_last_conversation_idx(self,agent,current_long_term_memory):
        last_conversation_idx = -1
        for i, history in enumerate(current_long_term_memory):
            if history.send_name == agent.name:
                last_conversation_idx = i
        return last_conversation_idx
    
    
    def _get_agent_new_memory(self,agent,current_long_term_memory):
        # get new conversation
        last_conversation_idx = self._get_agent_last_conversation_idx(agent,current_long_term_memory)

        if last_conversation_idx == -1:
            new_conversation =current_long_term_memory
        elif (
            last_conversation_idx
            == len(current_long_term_memory) - 1
        ):
            new_conversation = []
        else:
            new_conversation = current_long_term_memory[
                last_conversation_idx + 1 :
            ]

        # get chat history from new conversation
        return Memory.get_chat_history(new_conversation)
    
    
    def _observe(self,agent):
        MAX_CHAT_HISTORY = eval(os.environ["MAX_CHAT_HISTORY"])
        current_state = agent.current_state
        current_role = agent.state_roles[current_state.name]
        current_component_dict = current_state.components[current_role]
        
        # cooperative:Sharing information between different states ;  competive: No information is shared between different states
        current_chat_history_idx = self.current_chat_history_idx if self.environment_type == "competive" else 0
        current_long_term_memory = self.shared_memory["long_term_memory"][current_chat_history_idx:]
        current_chat_embbedings = self.shared_memory["chat_embeddings"][current_chat_history_idx:]
            
        
        # relevant_memory
        query = current_long_term_memory[-1].content

        relevant_memory = get_relevant_history(
            query,
            current_long_term_memory[:-1],
            current_chat_embbedings[:-1],
        )
        relevant_memory = Memory.get_chat_history(relevant_memory,agent.name)
        
        relevant_memory = eval(Agent_observe_relevant_memory)
        agent.relevant_memory = relevant_memory
        
        
        # get chat history from new conversation
        conversations = self._get_agent_new_memory(agent,current_long_term_memory)

        # memory = relevant_memory + summary + history + query
        query = current_long_term_memory[-1]
        current_memory = eval(Agent_observe_memory)

        return {"role": "user", "content": current_memory}