import asyncio
from datetime import date
from enum import Enum
import json
import time
from typing import Any, AsyncIterable, Awaitable
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field
import requests
from app.agents.hotel_agent import HotelAgent
from app.agents.tools.add_to_cart_tool import add_to_cart
from app.agents.tools.search_room_tool import get_available_dates, search_room
from app.utils.logger import CustomLogger
from app.schemas.chat import ChatRequest, ChatResponse
from sse_starlette.sse import EventSourceResponse
from langchain_openai.chat_models import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from app.configs.settings import settings
from langchain.callbacks.manager import collect_runs, atrace_as_chain_group
from langchain.chains.structured_output import create_structured_output_runnable
from app.utils.callback import MultipleLLMAsyncIteratorCallbackHandler
from langchain.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate, BaseMessagePromptTemplate, MessagesPlaceholder


router = APIRouter()
logger = CustomLogger.get_logger()


@router.post("/chat", tags=['Chat endpoint'])
async def chat(args: ChatRequest) -> ChatResponse:
    try:
        question = args.question
        history = args.history

        hotel_agent = HotelAgent(
            question=question,
            chat_history=history,    
        )
        
        chat_response = await hotel_agent.ainvoke()
        return chat_response
    except Exception as e:
        logger.error(f"Failed to upload file: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Failed to upload file: {str(e)}")
    
@router.post("/stream-chat", tags=['Streame chat endpoint'])
async def chat(args: ChatRequest):
    question = args.question
    history = args.history
    async def send_message() -> AsyncIterable[Any]:
        try:
            llm = ChatOpenAI(streaming=True,openai_api_key=settings.OPENAI_KEY)
            chain  = llm | StrOutputParser()
            
            for token in chain.stream(question):
                yield token
        
        except Exception as e:
            logger.error(f"Failed to upload file: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Failed to upload file: {str(e)}")
    return EventSourceResponse(send_message())


class Booking(BaseModel):
    question:str
    history:list

class toolEnum(Enum):
    modifyReservation = 'modify_reservation_tool'
    cancelReservation = "cancel_reservation_tool"
    search = "search_room_tool"
    booking = "add_to_cart"
    billing = "add_billing_details"
    default = "no_tool"
    
    
class params(BaseModel):
    check_in : str = Field(description="MM-DD-YYYY",)
    check_out: str = Field(description="MM-DD-YYYY",)
    adultCount: int = Field(default=2)
    childCount: int = Field(default=0)
    room_name: str = Field(default='')
    room_type_id: str = Field(description="If user query is add to cart or booking then find the room which user want to book and provide the room_type_id of that room, if not found the try to provide room_type_id of any room which is close to user requested room ",default='')

class sortEnum(Enum):
    ascending = "ascending"
    descending = "descending"
    default = "default"

class sortKey(Enum):
    rate = "rate"
    occupancy = "occupancy"
    default = "default"
class AgentSchema(BaseModel):
    tool:toolEnum = Field('tell which tool we should use according to user query, add_to_cart tool for room booking and add to cart, search_room_tool for search rooms')
    answer:str = Field(description="Always generate positive answer for user. If going to use tool then say 'Let me' then progress name we are going to do. ")
    params:params
    sort: sortEnum = Field(description="ascending for lowest price, descending for highest price else make it default if user has no specific requierment")
    key: sortKey = Field(description="Make it default if no user has no requierments")
    
    
def get_chat_prompt(history:list = []) -> ChatPromptTemplate:

    system_template = """You are a helpful AI desk manager at Bay occean Plazabeach hotel. politly decline to answer if question is irrelevant to our hotel business.
    
    search_room_tool provide real time data and prices and add_to_cart tool add the room to user cart and book the room.
    
    Use add_to_cart tool if user want to add the room to the cart or book the room.
    Use search_room_tool if user want to search the rooms.
    Use add_billing_details if user want to open the billing form or want to do billing or checkout.
    Use cancel_reservation_tool if user want to cancel their booked reservation.
    Use modify_reservation_tool if user want to modify the reservation.
    
    if user query is related to booking or add to cart, you will get room data in context given below.
    
    <CONTEXT>{context}</CONTEXT>
    
    Always use below currunt given date
    <DATE>{date}<DATE>
    
    Always use this default value for both tool given below if not provided by user
    <DEFAULT-VALUES>
        <CHECK-IN>{checkin}</CHECK-IN>
        <CHECK-OUT>{checkout}</CHECK-OUT>
    </DEFAULT-VALUES>
    """

    messages = [
        SystemMessagePromptTemplate.from_template(system_template),
        *history,
        HumanMessagePromptTemplate.from_template("{question}")
    ]
    return ChatPromptTemplate.from_messages(messages)

@router.post(
    "/test",
    summary="test",
    status_code=200,
)
async def chat(args:ChatRequest):
    
    async def send_message() -> AsyncIterable[Any]:
        try:
            with collect_runs() as runs_cb:
                async with atrace_as_chain_group("HotelBot", inputs={"question": args.question}) as group_manager:
                    
                    callback = MultipleLLMAsyncIteratorCallbackHandler()
                    llm = ChatOpenAI(api_key=settings.OPENAI_KEY,
                                     callbacks=[callback],
                                     streaming=True,
                                     model_name="gpt-4-turbo-preview")

                    chain = create_structured_output_runnable(
                        output_schema=AgentSchema,
                        prompt=get_chat_prompt(history=args.history),
                        llm=llm
                    )
                    
                    today = date.today()
                    checkin,checkout = await get_available_dates()
                    
                    async def wrap_done(fn: Awaitable, event: asyncio.Event):
                        """
                        Wrap an awaitable with a event to signal when it's done or an exception is raised.
                        """
                        try:
                            result = await fn
                        except Exception as e:
                            raise e
                        finally:
                            # Signal the aiter to stop.
                            print("stop")
                            event.set()
                        return result
            
                    
                    task = asyncio.create_task(wrap_done(chain.ainvoke(input={'question':args.question,'date':today.strftime("%d %B, %Y"), 'checkin':checkin,'checkout':checkout,'context':json.dumps(args.roomList)},config={'callbacks':group_manager,"run_name":"AgentExecutor"}),callback.done),)
                    
                    async for token in callback.aiter():
                        yield token
                        
                    result:AgentSchema = await task
                    
                    if result.tool.value == 'search_room_tool':
                        yield "#qwerasdf"
                        res = search_room(adultCount=result.params.adultCount,
                                          childCount=result.params.childCount,
                                          check_in=result.params.check_in,
                                          check_out=result.params.check_out,
                                          sortValue=result.sort.value,
                                          keyValue=result.key.value)
                        async for token in res:
                            yield token
                    
                    if result.tool.value == 'add_to_cart':
                        if args.roomList:
                            async for token in add_to_cart(args.question, args.roomList,llm,args.history,checkin,checkout,result.params.room_type_id,group_manager):
                                yield token
                    
                    if result.tool.value == "add_billing_details":
                        yield '#tEjdaY'
                        
                    if result.tool.value == "cancel_reservation_tool":
                        yield '#hlpwjw'
                        
                    if result.tool.value == "modify_reservation_tool":
                        yield '#mrt311'
                        
                    if result.tool.value == 'no_tool':
                        pass
                        # return result.__dict__
                        
                    await group_manager.on_chain_end({"output": result.__dict__})
                    run_id = runs_cb.traced_runs[0].id.__str__()
                    
                    return 
        except Exception as e:
            logger.error(f"Error: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Error: {str(e)}")
    return EventSourceResponse(send_message())