from typing import Union, Dict, Tuple
from datetime import datetime, timezone

from app.utils.prompts import create_agent_prompt, get_agent_prompt
from app.rag.tools import gen_report, fact_report, multi_fact_report
from app.schemas.chat_schema import factSchema, reportSchema, multiFactSchema

from langchain.tools import Tool
from langchain.agents import AgentExecutor
from langchain.prompts.chat import MessagesPlaceholder
from langchain_openai.chat_models.base import ChatOpenAI
from langchain.agents.openai_functions_agent.base import create_openai_functions_agent

class CustomTool(Tool):
    """Custom Tool Just For Override _to_args_and_kwargs function."""
    def _to_args_and_kwargs(self, tool_input: Union[str, Dict]) -> Tuple[Tuple, Dict]:
        if isinstance(tool_input, str):
            return (tool_input,), {}
        else:
            return (), tool_input

class RAG:
    
    def __init__(
        self,
        llm: ChatOpenAI,
        query:str,
        db_user: str,
        db_password: str,
        db_name: str,
        db_host :str
        ) -> None:
        self.llm = llm
        self.query = query
        self.db_user = db_user
        self.db_password = db_password
        self.db_name = db_name
        self.db_host = db_host
        pass
    
    async def create_tool(self, tool_name:str) -> CustomTool:

        tool_details = {
           'genReportTool' : {
                'description':'Always use this tool to fetch data from sql database or user gives report name that we have.',
                'func': gen_report,
                'schema': reportSchema
            },
           'factReportTool': {
               'description':'Use this tool to fetch the past dates data from the fact file for particular date.',
               'func':fact_report,
               'schema':factSchema
           },
           'multiFactReportTool': {
               'description':'Use this tool to fetch the past data from the fact file in date range.',
               'func':multi_fact_report,
               'schema':multiFactSchema
           }
        }
        
        return CustomTool(
            name=tool_name,
            description= tool_details.get(tool_name).get('description'),
            func= tool_details.get(tool_name).get('func'),
            coroutine= tool_details.get(tool_name).get('func'),
            return_direct=True,
            args_schema=tool_details.get(tool_name).get('schema')
        )
    
    async def ainvoke(self, chat_history:list = [], callbacks = None) -> dict:
        
        self.callbacks = callbacks
        prompt = create_agent_prompt(system_message=get_agent_prompt(),extra_prompt_messages=[MessagesPlaceholder(variable_name='chat_history')])
        tools = [await self.create_tool('genReportTool'),
                 await self.create_tool('factReportTool'),
                 await self.create_tool('multiFactReportTool')]
        
        agent = AgentExecutor(
            agent=create_openai_functions_agent(llm=self.llm, tools=tools, prompt=prompt),
            tools=tools,
            return_intermediate_steps=True,
        )
        schema = ''
        fact = ''
        with open("/var/www/html/chat-with-data/test/fullSchema.txt", 'r') as file:
            schema = file.read()
        with open("/var/www/html/chat-with-data/test/hotel_2023-06-08.txt", 'r') as factFile:
            fact = factFile.read()
        
        response = await agent.ainvoke({"input": self.query, "chat_history": chat_history, "context": schema, "fact": fact, "date": datetime.today().strftime('%Y-%m-%d') }, config={"callbacks":callbacks, "run_name":"SuperAgent"})
        return { 'role': 'bot', 'answer': response.get('output'), 'timestamp': datetime.now(tz=timezone.utc).timestamp()}