import gradio as gr
from docx import Document
from transformers import pipeline
import tempfile
import time

# Load Mistral model
chatbot = pipeline(
    "text-generation",
    model="mistralai/Mistral-7B-Instruct-v0.1",
    device_map="auto",
    trust_remote_code=True
)

# Read .docx file and extract questions
def read_questions_from_docx(document_path):
    doc = Document(document_path)
    return [p.text.strip() for p in doc.paragraphs if p.text.strip().endswith("?")]

# --- Load .docx once into memory ---
syntax_questions = read_questions_from_docx("Syntax Questions.docx")
morphology_questions = read_questions_from_docx("Morphology Questions.docx")

# Main function
def run_tool(passage, question_type):
    if not passage.strip():
        return "Please enter a passage to analyze.", None, None

    try:
        questions = syntax_questions if question_type.lower() == "syntax" else morphology_questions
        if not questions:
            return "No valid questions found.", None, None

        est_seconds = round(len(questions) * 3.5, 1)
        estimated_time = f"Estimated generation time: ~{est_seconds} seconds"

        prompt = f"You are a classical language expert.\n\nHere is the passage:\n{passage}\n\n"
        prompt += "Answer the following questions clearly and completely:\n\n"
        for idx, q in enumerate(questions):
            prompt += f"Q{idx+1}: {q}\n"

        result = chatbot(prompt, max_new_tokens=1600, do_sample=False)[0]["generated_text"].strip()

        # Save result to .txt
        with tempfile.NamedTemporaryFile(delete=False, suffix=".txt", mode="w", encoding="utf-8") as f:
            f.write(result)
            file_path = f.name

        return result, file_path, estimated_time

    except Exception as e:
        return f"An error occurred: {str(e)}", None, None

# --- Gradio UI ---
with gr.Blocks(theme="soft") as demo:
    gr.Markdown("""
    ## **Classical Language Query Assistant**
    Submit a Latin or Greek passage and select whether to analyze **syntax** or **morphology**.
    Answers are generated using **Mistral-7B-Instruct**, hosted via Hugging Face.
    """)

    with gr.Row():
        passage_input = gr.Textbox(label="Latin or Greek Passage", placeholder="Paste your passage here...", lines=4)
        question_type = gr.Radio(["Syntax", "Morphology"], label="Select Question Set")

    with gr.Row():
        output_text = gr.Textbox(label="Generated Answers", lines=25, interactive=False)
        output_file = gr.File(label="Download Answers (.txt)", interactive=False)

    est_time = gr.Textbox(label="Estimated Time", interactive=False)

    submit = gr.Button("Generate Answers")
    submit.click(fn=run_tool, inputs=[passage_input, question_type], outputs=[output_text, output_file, est_time])

demo.launch()