"""
 /*************************************************************************
 *
 * CONFIDENTIAL
 * __________________
 *
 *  Copyright (2024-2025) AI Labs, IronOne Technologies, LLC
 *  All Rights Reserved
 *
 *  Author  : Theekshana Samaradiwakara
 *  Description : This file contains the server application routing logic. 
 *  CreatedDate : 17/10/2024
 *  LastModifiedDate : 
 *************************************************************************/
"""

from typing import List

from langchain.schema import HumanMessage
from langchain.prompts import BaseChatPromptTemplate
from langchain.agents import Tool

from reggpt.tools.qa_agent import qa_agent_tools


QA_AGENT_TEMPLATE = """Answer the following questions as best you can.
If the question involves about Rules and regulations made by central bank..
You have access to the following tools:

{tools}


Use the following format:

Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action 
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question

Begin!
Question: {input}
{agent_scratchpad}"""


class CustomPromptTemplate(BaseChatPromptTemplate):
    """
    This is the prompt template for the multi-tool agent. It formats the messages in a specific way.
    """

    template: str
    tools: List[Tool]

    def format_messages(self, **kwargs) -> str:
        """
        This function formats the messages by replacing the variables in the template with the values in kwargs.
        """

        intermediate_steps = kwargs.pop("intermediate_steps")
        thoughts = ""
        for action, observation in intermediate_steps:
            thoughts += action.log
            thoughts += f"\nObservation: {observation}\nThought: "

        kwargs["agent_scratchpad"] = thoughts

        kwargs["tools"] = "\n".join(
            [f"{tool.name}: {tool.description}" for tool in self.tools]
        )

        kwargs["tool_names"] = ", ".join([tool.name for tool in self.tools])
        formatted = self.template.format(**kwargs)
        return [HumanMessage(content=formatted)]


qa_agent_prompt = CustomPromptTemplate(
    template=QA_AGENT_TEMPLATE,
    tools=qa_agent_tools,
    input_variables=["input", "intermediate_steps"],
)