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 GetPlacesTemplate(BaseModel):
    answer: List[str] = Field(description="List of places and their adresses separated by ','")


class Planner_Agent():
    def __init__(self):

        self.model_name = "gpt-4"
        self.model = OpenAI(model_name=self.model_name, temperature=0)

        self.output_parser_places = PydanticOutputParser(pydantic_object=GetPlacesTemplate)
        self.format_instructions_places = self.output_parser_places.get_format_instructions()

        self.prompt = 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. \n### User input to base itenerary on: \n{input}
                            ### Response:
                            """,
                    input_variables=["input"]
                    # partial_variables={"format_instructions": format_instructions_gether}
        )

        self.prompt_to_get_places = PromptTemplate(
                    template="""\
                            ### Instruction
                            You are a place retriever. From a given input you can creat a list of all the places referenced in it, as well as the adress of each location.
                            ### Input: {input}
                            ### Response: {format_instructions}
                            """,
                    input_variables=["input", "format_instructions"]
                    # partial_variables={"format_instructions": format_instructions_gether}
        )
    
    def format_prompt(self, input):
        return self.prompt.format_prompt(input=input)
    
    def get_itenerary(self, input):
        return self.model(input.to_string())
    
    def format_prompt_to_get_places(self, input):
        return self.prompt_to_get_places.format_prompt(input=input, format_instructions=self.format_instructions_places)
    
    def get_places_from_itenerary(self, itenerary):
        result = self.model(itenerary.to_string())
        parsed_result = self.output_parser_places.parse(result)
        return parsed_result.answer