import gradio as gr
from transformers import pipeline
import os

# Load model (replace with any instruct-style model hosted on Hugging Face Hub)
model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"  # You can change this
generator = pipeline("text-generation", model=model_name, device=-1)  # CPU only

# Prompt template
def build_prompt(essay_text):
    return f"""You are an experienced English teacher.

Your task is to evaluate the following student essay in 6 categories, each scored from 0 to 100.

Categories:
1. Grammar & Mechanics
2. Coherence & Flow
3. Clarity & Style
4. Argument Strength & Evidence
5. Structure & Organization
6. Teacher-Specific Style

Then, calculate the average of the 6 scores as the Total Grade.

Output format:
Return only a JSON object like this:
{{
  "Grammar & Mechanics": score,
  "Coherence & Flow": score,
  "Clarity & Style": score,
  "Argument Strength & Evidence": score,
  "Structure & Organization": score,
  "Teacher-Specific Style": score,
  "Total Grade": average
}}

Here is the essay:
{essay_text}
"""

# Text extractor for uploaded files
def extract_text(file):
    if file.name.endswith(".txt"):
        return file.read().decode("utf-8")
    return "Unsupported file type. Please upload a .txt file."

# Main function
def grade_essay(essay, file):
    if not essay and not file:
        return "Please provide either text or a file."
    
    if file:
        essay = extract_text(file)
        if "Unsupported" in essay:
            return essay

    prompt = build_prompt(essay)
    result = generator(prompt, max_new_tokens=400, temperature=0.5)[0]["generated_text"]
    
    # Extract JSON from result
    start = result.find("{")
    end = result.rfind("}") + 1
    try:
        return result[start:end]
    except:
        return "Model response couldn't be parsed."

# UI
with gr.Blocks() as demo:
    gr.Markdown("# ✏️ MimicMark - Essay Evaluator")
    gr.Markdown("Upload or paste your essay. The AI will rate it in 6 categories and return a final score.")

    with gr.Row():
        essay_input = gr.Textbox(label="Paste Your Essay Here", lines=12)
        file_input = gr.File(label="Or Upload a .txt File", file_types=[".txt"])
    
    output = gr.Textbox(label="Evaluation Output", lines=15)

    btn = gr.Button("Evaluate Essay")
    btn.click(fn=grade_essay, inputs=[essay_input, file_input], outputs=output)

demo.launch()