from langchain.docstore.document import Document
from langchain.vectorstores import FAISS
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.memory.simple import SimpleMemory

from langchain.chains import ConversationChain, LLMChain, SequentialChain
from langchain.memory import ConversationBufferMemory

from langchain.prompts import ChatPromptTemplate, PromptTemplate
from langchain.document_loaders import UnstructuredFileLoader

from langchain.chat_models import ChatOpenAI
from langchain.llms import OpenAI
from langchain.memory import ConversationSummaryMemory

from langchain.callbacks import PromptLayerCallbackHandler
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

from langchain.schema import AIMessage, HumanMessage, SystemMessage
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.callbacks.base import BaseCallbackHandler
import gradio as gr

from threading import Thread
from queue import Queue, Empty
from threading import Thread
from collections.abc import Generator
from langchain.llms import OpenAI
from langchain.callbacks.base import BaseCallbackHandler

import itertools
import time
import os
import getpass
import json
import sys
from typing import Any, Dict, List, Union

import promptlayer
import openai
import gradio as gr

from pydantic import BaseModel, Field, validator

#Load the FAISS Model ( vector )
openai.api_key = os.environ["OPENAI_API_KEY"]
db = FAISS.load_local("db", OpenAIEmbeddings())

#API Keys 
promptlayer.api_key = os.environ["PROMPTLAYER"]

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

from langchain.callbacks import PromptLayerCallbackHandler
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.memory import ConversationSummaryMemory

# Defined a QueueCallback, which takes as a Queue object during initialization. Each new token is pushed to the queue.
class QueueCallback(BaseCallbackHandler):
    """Callback handler for streaming LLM responses to a queue."""

    def __init__(self, q):
        self.q = q

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        self.q.put(token)

    def on_llm_end(self, *args, **kwargs: Any) -> None:
        return self.q.empty()

MODEL = "gpt-3.5-turbo"

# Defined a QueueCallback, which takes as a Queue object during initialization. Each new token is pushed to the queue.
class QueueCallback(BaseCallbackHandler):
    """Callback handler for streaming LLM responses to a queue."""

    def __init__(self, q):
        self.q = q

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        self.q.put(token)

    def on_llm_end(self, *args, **kwargs: Any) -> None:
        return self.q.empty()

class UnitGenerator:

  def __init__(self, prompt_template='', model_name=MODEL, verbose=False, temp=0.2):
    self.verbose = verbose
    self.llm = ChatOpenAI(
      model_name=MODEL,
      temperature=temp
      )

    #The zero shot prompt provided at creation
    self.prompt_template = prompt_template

  def chain(self, prompt: PromptTemplate, llm: ChatOpenAI) -> LLMChain:
        return LLMChain(
            llm=llm,
            prompt=prompt,
            verbose=self.verbose,
        )

  def stream(self, input) -> Generator:

  # Create a Queue
    q = Queue()
    job_done = object()

    llm = ChatOpenAI(
        model_name=MODEL,
        callbacks=[QueueCallback(q),
                  PromptLayerCallbackHandler(pl_tags=["unit-generator"])],
        streaming=True,
        )

    prompt = PromptTemplate(
          input_variables=['input'], 
          template=self.prompt_template
        )

    # Create a funciton to call - this will run in a thread
    def task():
        resp = self.chain(prompt,llm).run(
            {'input':input})
        q.put(job_done)

    # Create a thread and start the function
    t = Thread(target=task)
    t.start()

    content = ""

    # Get each new token from the queue and yield for our generator
    while True:
        try:
            next_token = q.get(True, timeout=1)
            if next_token is job_done:
                break
            content += next_token
            yield next_token, content
        except Empty:
            continue

agent_prompt = """
            Take the following class overview (delimited by three dollar signs)
            $$$
            {input}
            $$$

            Make a list of 3 big and enduring ideas that students should walk away with.
            Make a list of 9 essential questions we want students to think about.These questions should be open-ended and provocative. Written in "kid friendly" language. Designed to focus instruction for uncovering the important ideas of the content.
            Make a list of 7 ideas, concepts, generalizations and principles we want students to know and understand about the unit or topic we are teaching?
            Make a list of critical skills describing what we want students to be able to do. Each item should begin with "Students will be able to..."
            Make a list of 7 potential assessments we can use to give students opportunities to demonstrate their skills. 
            Make a list of 7 ways that I might adapt the experience for different learning styles. Make a list of 5 materials and technologies that I can enhance the learning experience.
            Make a list of 4 learning activities that support and scaffold the learning experience. 
        """

unit_generator = UnitGenerator(prompt_template=agent_prompt)

def generate_unit(input):
    for next_token, content in unit_generator.stream(input):
        yield(content)

gr.Interface(generate_unit,
             [gr.Textbox(
                label="Enter your unit vision.",
                info="Provide high level details for the learning experience."
             )],
              [gr.Textbox(
                label="Unit",
              )],allow_flagging="never").queue().launch(debug=True, share=True)