from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import gradio as gr

# Load the Hugging Face model
model_name = "fbellame/llama2-pdf-to-quizz-13b"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
quiz_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)

# Example quiz data (replace with dynamic generation logic)
quiz_data = [
    {
        "question": "What is the purpose of transformers in power systems?",
        "options": [
            "Increase voltage",
            "Decrease voltage",
            "Both increase and decrease voltage",
            "None of the above",
        ],
        "answer": "Both increase and decrease voltage",
    }
]

# Feedback collection
feedback_list = []
current_index = [0]

def get_next_question():
    if current_index[0] >= len(quiz_data):
        return "No more questions to review.", []
    question = quiz_data[current_index[0]]
    return f"Question: {question['question']}", question["options"]

def evaluate_question(rating, feedback, email, discipline, experience):
    if rating not in ["Exactly Accurate and Relevant", "I Have No Idea (Not in My Area of Expertise)"] and not feedback.strip():
        return "Feedback required for this rating. Please provide corrections."
    
    question_data = quiz_data[current_index[0]]
    feedback_list.append({
        "question": question_data["question"],
        "options": question_data["options"],
        "rating": rating,
        "feedback": feedback,
        "email": email,
        "discipline": discipline,
        "experience": experience,
    })
    current_index[0] += 1
    if current_index[0] >= len(quiz_data):
        return "Feedback submitted. No more questions to review."
    return "Feedback submitted. Next question ready."

def generate_mailto_link():
    if not feedback_list:
        return "No feedback to send."
    
    subject = "Feedback for Train the Trainer App"
    body_lines = ["Feedback for Train the Trainer App:\n"]
    for entry in feedback_list:
        body_lines.append(f"Question: {entry['question']}\n")
        body_lines.append(f"Options: {', '.join(entry['options'])}\n")
        body_lines.append(f"Rating: {entry['rating']}\n")
        body_lines.append(f"Feedback: {entry['feedback']}\n")
        body_lines.append(f"Email: {entry['email']}\n")
        body_lines.append(f"Discipline: {', '.join(entry['discipline'])}\n")
        body_lines.append(f"Experience: {entry['experience']} years\n")
        body_lines.append("\n---\n")
    body = "\n".join(body_lines)
    body_encoded = body.replace(" ", "%20").replace("\n", "%0A")
    mailto_link = f"mailto:van.woods@usace.army.mil?subject={subject}&body={body_encoded}"
    return mailto_link

# Gradio Interface
with gr.Blocks() as trainer_app:
    gr.Markdown("# Train the Trainer: Refine AI-Generated Questions")
    question_box = gr.Textbox(lines=5, label="Generated Question", interactive=False)
    options_box = gr.CheckboxGroup([], label="Options (for reference)")
    rating = gr.Radio(
        [
            "Exactly Accurate and Relevant",
            "Somewhat",
            "Kinda Sorta",
            "Not Even Close",
            "I Have No Idea (But I Should)!",
            "I Have No Idea (Not in My Area of Expertise)"
        ],
        label="Rate the Question",
        required=True
    )
    feedback = gr.Textbox(label="Explain corrections or feedback (optional)", lines=3)
    email = gr.Textbox(label="Email Address (optional)")
    discipline = gr.CheckboxGroup(
        ["Electrical Power Systems", "Lighting Design", "Controls", "Other"],
        label="Select your discipline(s)"
    )
    experience = gr.Number(label="Years of Experience", value=0)
    submit_btn = gr.Button("Submit Feedback")
    next_btn = gr.Button("Next Question")
    mailto_btn = gr.Button("Generate Mailto Link")
    feedback_output = gr.Textbox(label="Feedback Status", interactive=False)
    mailto_output = gr.Textbox(label="Mailto Link", interactive=False)

    next_btn.click(get_next_question, inputs=None, outputs=[question_box, options_box])
    submit_btn.click(
        evaluate_question,
        inputs=[rating, feedback, email, discipline, experience],
        outputs=feedback_output,
    )
    mailto_btn.click(generate_mailto_link, inputs=None, outputs=mailto_output)

trainer_app.launch()