import typing
from typing import List, Union

from langchain_core.messages import AIMessage, HumanMessage, ChatMessage, FunctionMessage, ToolMessage
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, PromptTemplate, MessagesPlaceholder

from agentS.prompts import ranker_json, json_part


class AgentsPromptCreator:
    @staticmethod
    def create_agent_prompt(agent_name, prompt=None):
        # Activate create_{agent_name}_prompt method
        method_name = f"create_{agent_name.lower()}_prompt"
        method = getattr(AgentsPromptCreator, method_name, None)
        if method is None:
            raise ValueError(f"Method {method_name} not found")
        return method(prompt)

    @staticmethod
    def create_ranker_prompt(prompt, env_policies):
        # Modify the prompt to use an f-string for the JSON part
        modified_prompt = f"""
        {prompt.split("**Output your analysis in the following JSON format**:")[0]}
        **Output your analysis in the following JSON format**:
        {{json_part}}
        \n\nYour goal is to streamline the decision-making process by providing a focused, relevant subset of page elements. Accuracy in element selection is crucial for successful task completion. Always prioritize the most specific, interactive elements that directly match the user's intended action, while learning from past interactions to avoid repeating unsuccessful approaches.
        """

        prompt = ChatPromptTemplate(
            input_variables=['elements', 'input', 'url'],
            input_types={
                'scratchpad': List[Union[AIMessage, HumanMessage, ChatMessage, FunctionMessage, ToolMessage]],
                'policy': str,
                'read_page': str,
                'env_policies': str,
            },
            partial_variables={
                'scratchpad': [],
                'policy': '',
                'read_page': '',
                'json_part': ranker_json,
                'env_policies': env_policies,
            },
            messages=[
                SystemMessagePromptTemplate(
                    prompt=PromptTemplate(
                        input_variables=['elements', 'input', 'url'],
                        partial_variables={
                            'scratchpad': '{scratchpad}',
                            'policy': '{policy}',
                            'read_page': '{read_page}',
                            'json_part': '{json_part}',
                            'env_policies': '{env_policies}',
                        },
                        template=modified_prompt
                    )
                ),
                MessagesPlaceholder(variable_name='scratchpad', optional=True)
            ]
        )
        return prompt

    @staticmethod
    def create_action_agent_prompt(prompt, env_policies):
        prompt = ChatPromptTemplate(
            input_variables=['elements', 'input', 'url'],
            input_types={
                'scratchpad': List[Union[AIMessage, HumanMessage, ChatMessage, FunctionMessage, ToolMessage]],
                'policy': str,
                'read_page': str,
                'env_policies': str,
            },
            partial_variables={
                'scratchpad': [],
                'policy': '',
                'read_page': '',
                'env_policies': env_policies,
            },
            messages=[
                SystemMessagePromptTemplate(
                    prompt=PromptTemplate(
                        input_variables=['elements', 'input', 'url'],
                        partial_variables={
                            'scratchpad': '{scratchpad}',
                            'policy': '{policy}',
                            'read_page': '{read_page}',
                            'env_policies': '{env_policies}',
                        },
                        template=prompt
                    )
                ),
                MessagesPlaceholder(variable_name='scratchpad', optional=True)
            ]
        )
        return prompt

    @staticmethod
    def create_planner_agent_prompt(prompt, env_policies):
        # Modify the PLANNER_AGENT_PROMPT to use an f-string for the JSON part
        modified_prompt = f"""
        {prompt.split("**Output - Dynamic Action Plan**:")[0]}
        **Output - Dynamic Action Plan**:
        Generate a comprehensive action plan in JSON format. The plan should include a list of steps, each with an action and details. Also include metadata such as estimated steps, confidence level, potential obstacles, and a fallback strategy. Ensure that your plan directly addresses the reason for the update request and builds upon the successes of the previous policy while rectifying any shortcomings.\n\n
        The JSON structure should be as follows:\n
        {{json_part}}
        \n\n Ensure that your new policy directly addresses the reason for the update request and builds upon the successes of the previous policy while rectifying any shortcomings.
        """
        prompt = ChatPromptTemplate(
            input_variables=['url', 'input'],
            input_types={
                'scratchpad': List[Union[AIMessage, HumanMessage, ChatMessage, FunctionMessage, ToolMessage]],
                'elements': str,
                'policy': str,
                'update_reason': str,
                'read_page': str,
                'env_policies': str,
            },
            partial_variables={
                'scratchpad': [],
                'elements': '',
                'policy': '',
                'update_reason': '',
                'read_page': '',
                'json_part': json_part,
                'env_policies': env_policies,
            },
            messages=[
                SystemMessagePromptTemplate(
                    prompt=PromptTemplate(
                        input_variables=['url', 'input'],
                        partial_variables={
                            'elements': '{elements}',
                            'policy': '{policy}',
                            'update_reason': '{update_reason}',
                            'read_page': '{read_page}',
                            'scratchpad': '{scratchpad}',
                            'json_part': '{json_part}',
                            'env_policies': '{env_policies}',
                        },
                        template=modified_prompt
                    )
                ),
                MessagesPlaceholder(variable_name='scratchpad', optional=True)
            ]
        )
        return prompt

    @staticmethod
    def create_general_prompt(prompt):
        prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    "You are a helpful AI assistant, collaborating with other assistants."
                    " Use the provided tools to progress towards answering the question."
                    " If you are unable to fully answer, that's OK, another assistant with different tools "
                    " will help where you left off. Execute what you can to make progress."
                    " If you or any of the other assistants have the final answer or deliverable,"
                    " prefix your response with FINAL ANSWER so the team knows to stop."
                    " You have access to the following tools: {tool_names}.\n{system_message}",
                ),
                MessagesPlaceholder(variable_name="messages"),
            ]
        )
        return prompt.partial(
            system_message="",
            messages="",
        )
