import json
import requests

from pydantic import BaseModel, Extra, Field
from typing import List, Optional, Union
from datetime import datetime, timedelta

from langchain.callbacks.manager import CallbackManager
from langchain.chains.openai_functions import create_structured_output_runnable
from langchain_core.runnables import Runnable
from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langchain.schema import SystemMessage, AIMessage, HumanMessage, BaseMessage
from langchain.prompts.chat import BaseMessagePromptTemplate

from app.configs import settings
from app.schemas.chat import ChatAdditonal
from . import MultiArgsTool

from app.utils.logger import CustomLogger
logger = CustomLogger.get_logger()

# class AgentToolSchema(BaseModel):
#     adult: int = Field(description="tells the number of adults", default=None)
#     child: int = Field(description="tells the number of children", default=None)
#     from_date: datetime = Field(description="Start date of booking", default=None)
#     to_date: datetime = Field(description="End date of booking", default=None)
    

class AgentAnswerSchema(BaseModel):
    answer: str = Field(description="Answer to user question with details of rooms")
    followUpResponses:List[str] = Field(description="Generate very short answers consisting of 2-3 words only. Example '1 adult and 1 child'")
    
    class Config:
        extra = Extra.allow 
        arbitrary_types_allowed = True
        
class AgentAnswerTool:
    def __init__(self, question: str) -> None:
        self.question = question
        
    def __call__(self):
        return MultiArgsTool(
            name= "missing_parameters",
            description="Use this tool to format your response for user. Only use this when you have any question of user.",
            func= self.find_missing_parameters,
            args_schema= AgentAnswerSchema,
            return_direct=True
        )
        
    def find_missing_parameters(self, answer, followUpResponses, **kwargs):
        return { 'answer':answer, 'followUpResponses':followUpResponses }