#0.0.1
import openai

#Langchain Imports
from langchain.prompts import ChatPromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.schema import StrOutputParser
from langchain.callbacks.base import BaseCallbackHandler

#Type Hinting
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Type, Union

#Utils
import os
from operator import itemgetter
import json
from dataclasses import dataclass
import sys
import logging


from queue import Queue, Empty
from threading import Thread
import gradio as gr

openai.api_key = os.environ["OPENAI_API_KEY"]
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()

# 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=True, temp=0.2):
    self.verbose = verbose
    self.llm = ChatOpenAI(
      openai_api_key=api_key,
      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)],
        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?
            Misconceptions: Make a list of 5 common misconceptions that occur when learning about this topic.
            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..."
            Real-World Applications: Make a list of 5 ways that this skill can be used in "the real world"
            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 making it more universally accessible. 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}
            $$$
            Using this unit definition and STUDENT INTERESTS generate 10 unique learning activity idea inspired by different pedagogies. ( play-based, project-based, game-based, etc.)
            Each idea should include the randomly selected pedagogy, a title, a one sentence description of the activity, a one sentence description of how it addresses key concepts from the unit and one sentence explaining how it supports STUDENT INTERESTS.
            """

lesson_idea_generator = Agent(prompt_template=lesson_idea_generator_prompt)

lesson_builder_prompt = """
            The curriculum unit details are provided below (delimited by triple percent signs):
            %%% {input} %%%
            Use the curriculum unit and lesson description to generate a lesson plan using the universal by design methdology.
            Engagement: 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 ).
            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 ). 
            Alternative Representations: List 3 different ways the key concepts might be expressed or explained that could help different learners make the connection.
            Resources: What tools, materials, or resources will you need to effectively deliver this lesson? 
            Resource add-ons: List 3 additional technologies, tools or resources ( not listed above ) that might add more interest or enagagement for students. Explain why you chose them.   
            Assessment: Create a plan to measure students' understanding and mastery of the lesson's objectives. List at least 3 alternative ways for students to express their knowledge and skills. Reference the example assessments in the unit details. Reference relevant unit details by section and number (example: Key Concepts #4, Skills #2, etc ). 
            Support: List 5 ways to provide support and feedback to each student. 
            Expression: List 5 examples of ways they might personally express themselves throughout the lesson.
            Unit connections: 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 ).
            """

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, interests):
    input = unit + " \n\n STUDENT INTERESTS (delimited by triple ampersand): @@@ " + interests + " @@@"
    for next_token, content in lesson_idea_generator.stream(input):
        yield(content)

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

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

with app:

  gr.Markdown(
    """
    # Syllabo
    A suite of generative ai tools for designing and building learning experiences.
    """)

  with gr.Tab("Unit Builder"):
    gr.Markdown(
    """
    Use this tool to create a unit inspired by the Understanding by Design framework.
    """)

    unit_vision = gr.Textbox(label="Enter topic, activity or standard(s):")
    b1 = gr.Button("Make Unit")
    unit = gr.Textbox(label="Unit",show_copy_button=True)

  with gr.Tab("Lesson Generator"):
    gr.Markdown(
    """
    Use this tool to generate ideas for lesson activities. \n 
    The tool will select from a range of pedagogies.\n
    
    Note: It will use the unit generated from the "unit builder" and any notes on student interest you add below. \n
    If you like an idea copy and paste it into the next tool. 
    """)
    student_interests = gr.Textbox(label="General student interests, cultures & strengths")
    b2 = gr.Button("Generate Activities")
    lesson_ideas = gr.Textbox(label="Lesson Concepts")

  with gr.Tab("Lesson Builder"):
    gr.Markdown(
    """
    This tool takes a lesson idea and generates a lesson plan inspired by the Universal by Design framework.
    """)
    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,student_interests], outputs=lesson_ideas)
  b3.click(generate_lesson_plan, inputs=[unit,student_interests,lesson_description], outputs=lesson_plan)

app.queue().launch()