from datetime import datetime
from typing import Union, List, Literal
from app.agents.tools.book_room import BookRoomSchema
from app.agents.tools.show_suites import ShowRoomSchema
from app.configs.settings import settings
from app.schemas.chat import ChatResponse
from .tools.search_rooms import SearchRoomSchema
from . import create_agent_prompt
from datetime import date
from langchain.agents import AgentExecutor, create_openai_functions_agent
from langchain_openai.chat_models import ChatOpenAI
from langchain.schema import SystemMessage, AIMessage, HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain.callbacks.manager import collect_runs, atrace_as_chain_group


today = date.today()
HOTEL_AGENT_PROMPT = """You are a helpful AI hotel manager at the "Bay View Hotel". Your responsibilities include answering user queries about searching/checking room availability, booking rooms, canceling room bookings, and providing information on the cheapest available rooms. Provide concise answers, not exceeding 10 words unless more detail is required. Feel free to use any necessary tools like "SEARCH_ROOM", "BOOK_ROOM", "CANCEL_ROOM", and "SHOW_SUITES" to assist with queries. Transfer the conversation flow to the previous tool when needed. Answer in the same language as the conversation.
The current date is """+today.strftime("%d%t %B, %Y")+""" and never provide room availability for past dates. Include relevant images or URLs only if they were in the previous conversation.
Always promote our hotel by highlighting facilities, nearby popular locations, and other attractions. Catch users' attention and encourage them to visit our hotel including the cheapest available rooms.. Provide concise answers following the law of least action. Deafult members are 2 adults and 0 child and default check in date is today and checkout date is tommorow and at Lastly, confirm all the details for booking or checkout.
<HOTEL INFORMATION>
Name: Bay View Hotel,
Address: 4321 Gulf Bivd St Pete Beach, FL 33706,
Email: frontoffice@plazabeach.com,
phone: +1 (800) 257-8998
</HOTEL INFORMATION>"""

class HotelAgent: 
    def __init__(
            self,
            question: str,
            chat_history: List[HumanMessage | AIMessage | SystemMessage] = None,
            tools_used: List[Literal["SEARCH_ROOM", "BOOK_ROOM", "CANCEL_ROOM", "SHOW_SUITES", "missing_parameters"]] = ["SHOW_SUITES", "SEARCH_ROOM", "BOOK_ROOM", "CANCEL_ROOM", "missing_parameters"],
            stream: bool = False,
        ) -> None:       
        self.llm  = ChatOpenAI(
            openai_api_key=settings.OPENAI_KEY,
            model_name="gpt-4-turbo-preview",
            callbacks=[], 
        )
        
        self.question = question
        self.tools_used = tools_used
        self.chat_history = chat_history

    async def ainvoke(self) -> ChatResponse:
        today = date.today()
        formatted_date = today.strftime("%d%t %B, %Y")
        print(formatted_date)
        agent_executor = self._create_agent()
        with collect_runs() as runs_cb:
            async with atrace_as_chain_group("HotelBot", inputs={"question": self.question}) as group_manager:
                result = await agent_executor.ainvoke(
                input={
                    'question': self.question,
                    'today_date': formatted_date,
                },
                config={'callbacks':group_manager,"run_name":"AgentExecutor"})
                await group_manager.on_chain_end({"output": result})
                run_id = runs_cb.traced_runs[0].id.__str__()

        return self.__parse_results(result)

    
    def _create_agent(self) -> AgentExecutor:
        tools = self._setup_agent_tools() 
        agent=create_openai_functions_agent(
                llm=self.llm,
                tools=tools,
                prompt=create_agent_prompt(system_message=SystemMessage(content=HOTEL_AGENT_PROMPT), extra_prompt_messages=self.chat_history)
        )
        
        agentExecutor = AgentExecutor(
            agent=agent,
            tools=tools,
            return_intermediate_steps=True,
        ) 
        return agentExecutor
    
    def _setup_agent_tools(self) -> None:
        tools = []
        if "SHOW_SUITES" in self.tools_used:
            from .tools.show_suites import ShowRoomTool
            show_room_tool = ShowRoomTool(question=self.question, llm= self.llm, history=self.chat_history)
            tools.append(show_room_tool())
        
        if "SEARCH_ROOM" in self.tools_used:
            from .tools.search_rooms import SearchRoomTool
            search_room_tool = SearchRoomTool(question=self.question, llm= self.llm, history=self.chat_history)
            tools.append(search_room_tool())
        
        if "BOOK_ROOM" in self.tools_used:
            from .tools.book_room import BookRoomTool
            book_room_tool = BookRoomTool(question=self.question, llm= self.llm, history=self.chat_history)
            tools.append(book_room_tool())
            
        if "missing_parameters" in self.tools_used:
            from .tools.agent_tool import AgentAnswerTool
            missing_param_tool = AgentAnswerTool(question=self.question)
            tools.append(missing_param_tool())
            
        if "CANCEL_ROOM" in self.tools_used:
            pass
        return tools
    
    def __parse_results(self, result) -> ChatResponse:
        result = result['output']
        print(type(result))
        final_result = ChatResponse(answer="Something went wrong can you ask question again.")
        if isinstance(result, str):
            openAILLM  = ChatOpenAI(
            openai_api_key=settings.OPENAI_KEY,
            model_name="gpt-3.5-turbo-0125",
        )
            # prompt = ChatPromptTemplate.from_template("give 2 or 3 responses only of customer at hotel booking counter in 1 or 2 words and separated by '/'. Response of hotel recieptent is {response} and today date is {date}")
            # output_parser = StrOutputParser()

            # chain = prompt | openAILLM | output_parser
            # followUpQuestions = chain.invoke({"response": result, "date": today.strftime("%d%t %B, %Y")})
            # followUpQuestions = followUpQuestions.split('/')
            # print(followUpQuestions)
            final_result = ChatResponse(answer=result)
            
        if isinstance(result, SearchRoomSchema):
            final_result = ChatResponse(
                answer=result.answer,
                followUpQuestions=result.followUpQuestions,
                additional=result.additional,
            ) 
        if isinstance(result, ShowRoomSchema):
            final_result = ChatResponse(
                answer=result.answer,
                followUpQuestions=result.followUpQuestions,
                additional=result.additional,
            ) 
        if isinstance(result, BookRoomSchema):
            final_result = ChatResponse(
                answer=result.answer,
                followUpQuestions=result.followUpQuestions,
                additional=result.additional,
            )
        if isinstance(result, dict):
            final_result = ChatResponse(
                answer=result["answer"],
                followUpQuestions=result['followUpResponses']
            ) 
        return final_result