import os
import platform

import openai
import chromadb
import langchain

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.text_splitter import TokenTextSplitter
from langchain.llms import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain.chains import ChatVectorDBChain
from langchain.document_loaders import GutenbergLoader

from langchain.embeddings import LlamaCppEmbeddings
from langchain.llms import LlamaCpp

from langchain.output_parsers import StructuredOutputParser, ResponseSchema
from langchain.prompts import PromptTemplate, ChatPromptTemplate, HumanMessagePromptTemplate
from langchain.llms import OpenAI
from langchain.chains import LLMChain
from langchain.chains import SimpleSequentialChain

from langchain.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field, validator
from typing import List, Dict

class AnswerTemplate(BaseModel):
    type: List[str] = Field(description="What is the type of the trip: business, family, vactions. And with whom are you travelling? If can't anwser then leave it empty")
    where: str = Field(description="Where is the user going? If can't anwser then leave it empty")
    start_date: str = Field(description="What is the start date? If can't anwser then leave it empty")
    end_date: str = Field(description="What is the end date? If can't anwser then leave it empty")
    time_constrains: str = Field(description="Is there any time constrains? If can't anwser then leave it empty")
    # dates: Dict[str, str] = Field(description="What are the importante dates and times? If can't anwser then leave it empty")
    preferences: List[str] = Field(description="What does the user want to visit? If can't anwser then leave it empty")
    conditions: str = Field(description="Does the user has any special medical condition? If can't anwser then leave it empty")
    dist_range: str = Field(description="Max distance from a place? If can't anwser then leave it empty")
    # missing: str = Field(description="Is any more information needed?")

class CheckAnswerTemplate(BaseModel):
    isComplete: bool = Field(description="Is the input complete?")
    # missing: str = Field(description="If the answer to the last question is false, then what is missing?")
    answer: str = Field(description="""If the answer to 'isComplete' is true leave this empty, else complete this by giving a nice compliment to the user's choices and asking the user for the missing information. Just one question.""")
    

# os.environ["OPENAI_API_KEY"] = "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"


model_name = "gpt-4"
# model_name = "gpt-3.5-turbo"

chat_history = ""

model = OpenAI(model_name=model_name, temperature=0)

output_parser_gather = PydanticOutputParser(pydantic_object=AnswerTemplate)
format_instructions_gather = output_parser_gather.get_format_instructions()

output_parser_check = PydanticOutputParser(pydantic_object=CheckAnswerTemplate)
format_instructions_check = output_parser_check.get_format_instructions()

helper_anwser = "Hello, can you tell me your trip details and constraints so I can give you great recomendations?"

user_input = input("Helper: " + helper_anwser + "\nUser: ")

# output_parser_2 = PydanticOutputParser(pydantic_object=AnswerTemplate_2)
# format_instructions_2 = output_parser_2.get_format_instructions()

# prompt_gather = PromptTemplate(
#     template="""\
#             ### Instruction
#             You are Trainline Mate an helpful assistant that plans tours for people at trainline.com.
#             As a smart itinerary planner with extensive knowledge of places around the
#             world, your task is to determine the user's travel destinations and any specific interests or preferences from
#             their message. Create an itinerary that caters to the user's needs, making sure to name all activities,
#             restaurants, and attractions specifically. When creating the itinerary, also consider factors such as time
#             constraints and transportation options. Additionally, all attractions and restaurants listed in the itinerary
#             must exist and be named specifically. During subsequent revisions, the itinerary can be modified, while keeping
#             in mind the practicality of the itinerary. New place for each day. It's important to ensure that the number of
#             activities per day is appropriate, and if the user doesn't specify otherwise, the default itinerary length is
#             five days. The itinerary length should remain the same unless there is a change by the user's message. Here is the history that you have so far: {history} \n### User: \n{input}
#             \n### Response: {format_instructions}
#             """,
#     input_variables=["input", "history", "format_instructions"]
#     # partial_variables={"format_instructions": format_instructions_gather}
# )

prompt_gather = PromptTemplate(
    template="""\
            ### Instruction
            You are Trainline Mate an helpful assistant that plans tours for people at trainline.com.
            As a smart itinerary planner with extensive knowledge of places around the
            world, your task is to determine the user's travel destinations and any specific interests or preferences from
            their message. Here is the history that you have so far: {history} \n### User: \n{input}
            \n### Response: {format_instructions}
            """,
    input_variables=["input", "history", "format_instructions"]
    # partial_variables={"format_instructions": format_instructions_gather}
)

prompt_check = PromptTemplate(
    template="""\
            ### Instruction
            Is this input complete? If not, what is missing? 
            ### Important: Give a nice compliment to the user's choices: {user_input} 
            ### Then ask for the missing information.
            Don't ask more then one question.
            Ask just one of the following:
            If 'type' is empty then ask the user what is the objective of the trip and with whom are you travelling?;
            If 'where' is empty then ask the user where is they going to travel to?;
            If 'start_date' is empty then ask the user what is the start date?;
            If 'end_date' is empty then ask the user what is the end date?;
            If 'time_constrains' is empty then ask the user if is there any time constrains that should be considered?;
            If 'preferences' is empty then ask the user if they have thought about any activities you want to do while you're there?;
            If 'conditions' is empty then ask the user if they have any special medical condition?;
            If 'dist_range' is empty then ask the user what is the distance range you prefer for your ativities? \n### Input: {input}
            \n### Response: {format_instructions}
            """,
    input_variables=["input", "user_input", "format_instructions"]
    # partial_variables={"format_instructions": format_instructions_check}
    # f 'dates' is empty then ask the user what are the importante dates and times?;
)

isComplete = False

while isComplete == False:

    _input_gather = prompt_gather.format_prompt(history=chat_history, input=user_input, format_instructions=format_instructions_gather)

    # chain_gather = LLMChain(llm=model, prompt=prompt_gather)
    # chain_check = LLMChain(llm=model, prompt=prompt_check)

    # overall_chain = SimpleSequentialChain(chains=[chain_gather, chain_check], verbose=True)

    result_gather = model(_input_gather.to_string())
    parsed_result_gather = output_parser_gather.parse(result_gather)
    print(parsed_result_gather)

    _input_check = prompt_check.format_prompt(input=parsed_result_gather, user_input="\nHelper: " + helper_anwser + "\nUser: " + user_input, format_instructions=format_instructions_check)
    result_check = model(_input_check.to_string())
    parsed_result_check = output_parser_check.parse(result_check)
    # print(parsed_result_check)

    isComplete = parsed_result_check.isComplete
    helper_anwser = parsed_result_check.answer

    if isComplete == False:
        chat_history += "User: " + user_input + "\nHelper: " + helper_anwser + "\n"
        user_input = input("Helper: " + helper_anwser + "\nUser: ")
        

    # print(overall_chain.run(input=user_input))

print(parsed_result_gather)