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-16k"

# 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 Agent:

  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

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

            Do the following:

            Enduring Ideas: Make a list of 5 big and enduring ideas that students should walk away with.
            Essential Questions: Make a list of 5 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.
            Key Concepts: Make a list of 5 ideas, concepts, generalizations and principles we want students to know and understand about the unit or topic we are teaching?
            Skills: Make a list of 5 critical skills describing what we want students to be able to do. Each item should begin with "Students will be able to..."
            Assessment: Make a list of 5 example assessments we can use to give students opportunities to demonstrate their skills. Explain the assessment idea and which key concepts and skills they map to.
            Different Learners: Make a list of 5 creative ways that I might adapt the experience for different learning styles. Explain the way I might adapt the experience to the learning style.
            Character: Make a list of 5 opportunity that this unit can support the learners development towards the "portrait" of a graduate. Each opportunity should identify the trait and how it might be applied.
         """

unit_generator = Agent(prompt_template=unit_generator_prompt)

lesson_idea_generator_prompt = """
            Below is a curriculum unit expressed using the understanding by design methodology ( delimited by the triple dollar signs).
            $$$
            {input}
            $$$
            Use this unit definition to generate a list of 10 learning activity ideas inspired by different pedagogies. ( play-based, project-based, game-based, etc.)
            Each idea should include a title, a one sentence description of the activity, a one sentence description of how it addresses key concepts from the unit.
        """

lesson_idea_generator = Agent(prompt_template=lesson_idea_generator_prompt)

lesson_builder_prompt = """ 
The learning unit details are provided below (delimited by triple percent signs): 
%%% {input} %%%
Make the following items:

Objective/Goals: What specific outcomes do you want your students to achieve by the end of this lesson? Reference relevant unit details by section and number (example: Key Concepts #4, Skills #2, etc ). 
Materials/Resources: What tools, materials, or resources will you need to effectively deliver this lesson?
Introduction/Hook: How will you capture the students' interest and introduce the topic at the start of the lesson? Reference relevant unit details by section and number (example: Key Concepts #4, Skills #2, etc ).
Procedure/Steps: Build a step-by-step sequence and time schedule of activities and discussions for this lesson. Reference relevant unit details by section and number (example: Key Concepts #4, Skills #2, etc ). 
Differentiation / Different learners: Address how you'll modify or adapt the lesson to cater to the needs of diverse learners in your classroom (e.g., English language learners, students with special needs, gifted students). Reference relevant unit details by section and number (example: Key Concepts #4, Skills #2, etc ). 
Assessment/Evaluation: Create a plan to measure students' understanding and mastery of the lesson's objectives? Reference the example assessments in the unit details. Reference relevant unit details by section and number (example: Key Concepts #4, Skills #2, etc ). 
Closure: Summarize and wrap up the lesson's content to ensure clarity and retention.
Unit connection: Summarize the ways that this lesson plan supports the unit details. Reference relevant unit details by section and number (example: Key Concepts #4, Skills #2, etc ).
Unit Map: Make a list of the referenced unit details. Reference relevant unit details by section and number (example: Key Concepts #4, Skills #2, etc ).

        """

lesson_builder = Agent(prompt_template=lesson_builder_prompt)

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

def generate_lesson_ideas(unit):
    for next_token, content in lesson_idea_generator.stream(unit):
        yield(content)

def generate_lesson_plan(unit,lesson_description):
    input = unit + "Lesson description below (delimited by triple ampersand): @@@ " + lesson_description + " @@@"
    for next_token, content in lesson_builder.stream(input):
        yield(content)

app = gr.Blocks(theme=gr.themes.Soft())

with app:

  gr.Markdown(
    """    
    # Curriculum Co-Creator

    A suite of tools for designing and building course units and lessons. 

    questions? 
    email: errol.m.king@gmail.com

    """)

  with gr.Tab("Vision --> Unit"):
    gr.Markdown(
    """

    The Vision to Unit tool will take an idea of a learning experience and build out a UBD unit.

    Instructions:
    1) Describe the learning experience. Include a few key topics, age and integrations. 
    2) Press the 'Generate Unit' button.
    3) Review in the unit section.
    4) Go to the next section.
    """)

    unit_vision = gr.Textbox(label="Describe your idea for a learning experience. Provide as much detail as possible.")
    b1 = gr.Button("Generate Unit")
    unit = gr.Textbox(label="Generated unit:",show_copy_button=True)

  with gr.Tab("Unit --> Lesson Ideas"):
    gr.Markdown(
    """

    This tool will take the generated unit from the "Vision to Unit" tool to create a few relevant lesson idea. 

    Instructions:
    1) Press the 'Generate Lesson Ideas' button.
    2) Review the ideas.
    3) Identify the one you want to expand upon. Highlight and copy to clipboard.
    4) Go to the next section.
    """)
    b2 = gr.Button("Generate Lesson Ideas")
    lesson_ideas = gr.Textbox(label="Lesson Ideas:")

  with gr.Tab("Lesson Ideas --> Lesson Plan"):        
    gr.Markdown(
    """

    This tool takes an lesson idea and generates a complete lesson plan. 

    Instructions:
    1) Paste one of the generated ideas ( or your own original )
    2) Add additional considerations: time, number of kids, age, etc. 
    3) Press the 'Generate Lesson Plan' button.    
    4) Read the lesson plan.
    5) Do it again!
    """)
    lesson_description = gr.Textbox(label="Input lesson idea:")
    b3 = gr.Button("Generate Lesson Plan")
    lesson_plan = gr.Textbox(label="Lesson Plan:",show_copy_button=True)

  b1.click(generate_unit, inputs=unit_vision, outputs=unit)
  b2.click(generate_lesson_ideas, inputs=unit, outputs=lesson_ideas)
  b3.click(generate_lesson_plan, inputs=[unit,lesson_description], outputs=lesson_plan)

app.queue().launch()