import os
import gradio as gr
import pytesseract
import pdfplumber
import docx
import requests
from PIL import Image

# API Configuration
api_token = os.getenv("API_TOKEN")
if not api_token:
    raise ValueError("API token not found. Set 'API_TOKEN' in your environment variables.")

API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3"
HEADERS = {"Authorization": f"Bearer {api_token}"}

def extract_text_from_pdf(pdf_file):
    text = ""
    with pdfplumber.open(pdf_file) as pdf:
        for page in pdf.pages:
            text += page.extract_text() + "\n"
    return text.strip()

def extract_text_from_docx(docx_file):
    doc = docx.Document(docx_file)
    return "\n".join([para.text for para in doc.paragraphs])

def extract_text_from_image(image_file):
    image = Image.open(image_file)
    return pytesseract.image_to_string(image)

def generate_quiz(text):
    if not text.strip():
        return "No content extracted. Please upload a valid file."
    
    prompt = f"""<s>[INST] Generate a quiz with multiple-choice questions based on the following content:
    {text} [/INST] Quiz:</s>"""
    
    payload = {"inputs": prompt}
    
    try:
        response = requests.post(API_URL, headers=HEADERS, json=payload)
        response.raise_for_status()
        result = response.json()
        
        if isinstance(result, list) and len(result) > 0:
            return result[0].get("generated_text", "Quiz generation failed.")
        else:
            return "Error: Unexpected response format."
    except requests.exceptions.RequestException as e:
        return f"API Error: {e}"

def quiz_app(uploaded_file):
    file_type = uploaded_file.name.split('.')[-1]
    extracted_text = ""
    
    if file_type == "pdf":
        extracted_text = extract_text_from_pdf(uploaded_file)
    elif file_type == "docx":
        extracted_text = extract_text_from_docx(uploaded_file)
    elif file_type in ["png", "jpg", "jpeg"]:
        extracted_text = extract_text_from_image(uploaded_file)
    
    if not extracted_text:
        return "No text extracted. Please upload a valid file."
    
    return generate_quiz(extracted_text)

iface = gr.Interface(
    fn=quiz_app,
    inputs=gr.File(label="Upload PDF, DOCX, or Image"),
    outputs=gr.Textbox(label="Generated Quiz"),
    title="AI Quiz Generator",
    description="Upload a file to generate a quiz based on its content."
)

if __name__ == "__main__":
    iface.launch()