import os
import json
import datetime
import gradio as gr
from gradio.components import Textbox

from langchain.prompts import PromptTemplate
from langchain.schema.runnable import RunnablePassthrough
from langchain_together import Together

os.environ["TOGETHER_API_KEY"] = '974d0f59c12b41d13e4f73500e3c035b9709df51ec2a67950542f1b24df0ee3b'


def generate_survey_questions(topic):
    template = """[INST] You are an expert in creating surveys. Generate a survey for the following topic: {question}. 
    Please provide 10 high-quality survey questions along with appropriate options. Avoid including open ended questions.
    Keep the formatting as follows:
    1. Questions 
    a. option 1 
    b. option 2 
    c. option 3
    d. ...
    [/INST]"""
    llm = Together(
        model="meta-llama/Llama-2-13b-chat-hf",
        temperature=0.7,
        top_k=50,
    )
    prompt = PromptTemplate.from_template(template)
    chain = (
        {"question": RunnablePassthrough()}
        | prompt
        | llm
    )

    answer_sub = chain.invoke(topic)
    return answer_sub


def generate_json_from_model_output(model_output):
    questions = []
    options = []

    lines = model_output.split('\n')
    lines = lines[1:-2]  # Remove the first and last line
    question = ''
    scale = []

    for line in lines:
        if line.strip():  # Non-empty line
            if line.startswith(('1.', '2.', '3.', '4.', '5.', '6.', '7.', '8.', '9.', '10.')):  # Question line
                # Add previous question and scale to lists
                if question and scale:
                    questions.append(question.strip())
                    options.append(scale)
                # Reset question and scale
                question = line.strip()
                scale = []
            else:  # Option line
                scale.append(line.strip().lstrip('*').strip())

    if question and scale:
        questions.append(question.strip())
        options.append(scale)

    json_data = []
    for q, o in zip(questions, options):
        json_data.append({'question': q, 'options': o})

    return json.dumps(json_data, indent=4)


# Create a Gradio interface
def survey_generator(topic):
    response = generate_survey_questions(topic)
    json_survey = generate_json_from_model_output(response)
    return json_survey


# Define the Gradio interface
inputs = Textbox(lines=5, label="Enter the survey topic")
outputs = Textbox(label="Generated Survey JSON", type="text")

# Create the Gradio app
title = "Survey Generator"
description = "Generate survey questions based on the provided topic."
examples = [["Hospitals cleanliness"]]
app = gr.Interface(fn=survey_generator, inputs=inputs, outputs=outputs, title=title, description=description, examples=examples)

# Launch the app
app.launch(share=True)