app.py

import os
from dotenv import load_dotenv
import google.generativeai as genai
from PIL import Image
import cloudconvert
import requests
import streamlit as st

# Load environment variables
load_dotenv()
genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
cloudconvert.configure(api_key=os.environ["CLOUDCONVERT_API_KEY"])

# Gemini model
model = genai.GenerativeModel('gemini-1.5-flash-latest')

def generate_latex_from_image(image_path, prompt):
    image = Image.open(image_path)
    response = model.generate_content([prompt, image])
    return response.text

def clean_latex_output(text):
    text = text.strip()
    if text.startswith("```latex"):
        first_newline = text.find("\n")
        text = text[first_newline+1:] if first_newline != -1 else text
    elif text.startswith("```"):
        first_newline = text.find("\n")
        text = text[first_newline+1:] if first_newline != -1 else text
    if text.endswith("```"):
        text = text[:text.rfind("```")]
    return text.strip()

def convert_latex_to_pdf(tex_file):
    job = cloudconvert.Job.create(payload={
        "tasks": {
            "upload-my-file": {"operation": "import/upload"},
            "convert-my-file": {
                "operation": "convert",
                "input": "upload-my-file",
                "input_format": "tex",
                "output_format": "pdf"
            },
            "export-my-file": {
                "operation": "export/url",
                "input": "convert-my-file"
            }
        }
    })
    upload_task_id = [t["id"] for t in job["tasks"] if t["name"] == "upload-my-file"][0]
    upload_task = cloudconvert.Task.find(id=upload_task_id)
    cloudconvert.Task.upload(file_name=tex_file, task=upload_task)

    job = cloudconvert.Job.wait(id=job["id"])

    for task in job["tasks"]:
        if task.get("status") == "error":
            raise RuntimeError(f"Task '{task['name']}' failed: {task.get('message', 'Unknown error')}")

    export_task = next((t for t in job["tasks"] if t["name"] == "export-my-file"), None)
    if not export_task or not export_task.get("result"):
        raise RuntimeError("Export task did not produce a file. Possibly LaTeX compile error.")

    file_url = export_task["result"]["files"][0]["url"]

    output_path = os.path.splitext(tex_file)[0] + ".pdf"
    response = requests.get(file_url)
    with open(output_path, "wb") as f:
        f.write(response.content)

    return output_path

# ---- Prompt to LLM ----
user_prompt = r"""
    Analyze the content of the whiteboard (may or may not be a whiteboard, could be a differnt board color) image. Act as a meticulous math tutor creating a set of LaTeX notes based on what is in the image. The goal is to produce a clean, professional-looking AND easy to understand (explain things in simple understandable terms) document to be used for studying.

    Here are the specific instructions for the output:
    1.  **Document Format:** Generate a single block of raw, compilable LaTeX code.
    2.  **Preamble:** Use `\documentclass{article}`, `\usepackage{amsmath}`, `\usepackage{amssymb}`, `\usepackage{graphicx}`. Do NOT include an `\author` command.
    3.  **Title:** Create a concise and accurate title for the document based on what is in the image"
    4.  **Document Body:**
        * Start with `\begin{document}` and `\maketitle`.
        * Use `\section` and `\subsection` commands to logically organize the content.
        * The first `\section` should be "Problem Statement," or "Topic," or whatever fits what is in the provided image.
        * The second `\section` should be "Step-by-Step Solution," or "Solving Steps," or whatever fits what is in the provided image, breaking down the image content into logical `\subsection`s.
        * **Logical Flow:** The subsections should follow the natural progression of the content shown on the whiteboard. For example:
            * Identify and state the core mathematical principle being used.
            * Break down complex steps into their own subsections.
            * Clearly show all intermediate steps and definitions as they appear on the board.
        * **Clarity and Aesthetics:**
            * Use `$$...$$` for key equations and results to ensure they are on their own lines.
            * Use the `align*` environment for multi-line derivations to ensure proper alignment and readability, preventing equations from being crammed on a single line.
            * Add concise, prose-based explanations to connect the steps, but do not write long paragraphs. The focus is on the math.
            * If the whiteboard shows a logical error or a messy final step, correct it and present the final answer clearly and correctly. The output should be a refined version of the notes on the board.

    5.  **Final Output:** The entire output must be a single block of raw, compilable LaTeX code. Do not include any explanatory text before or after the code block.
    
    The student uploading the image may or may not provide a custom prompt. If there is a prompt, it will be here: 
    """


# ---- Streamlit UI ----
st.title("📄 NoteBoard AI")

uploaded_image = st.file_uploader("Upload your board image", type=["jpg", "jpeg", "png"])
custom_prompt = st.text_area("Enter a custom prompt or leave the space below blank and click the 'Generate PDF' button", value="")

final_prompt = user_prompt + custom_prompt

if uploaded_image and st.button("Generate PDF"):
    with st.spinner("Processing..."):
        tex_path = "temp.tex"
        img_path = "temp_img.jpg"
        
        # Save uploaded image to file
        with open(img_path, "wb") as f:
            f.write(uploaded_image.read())

        # Generate LaTeX
        latex_notes = generate_latex_from_image(img_path, final_prompt)
        latex_notes = clean_latex_output(latex_notes)
        if "\\begin{document}" not in latex_notes:
            latex_notes = "\\documentclass{article}\n\\usepackage{amsmath}\n\\usepackage{amssymb}\n\\usepackage{graphicx}\n\\begin{document}\n" + latex_notes + "\n\\end{document}"


        with open(tex_path, "w", encoding="utf-8") as f:
            f.write(latex_notes)

        # Convert to PDF
        pdf_path = convert_latex_to_pdf(tex_path)

    st.success("✅ PDF generated!")
    with open(pdf_path, "rb") as f:
        st.download_button("📥 Download PDF", f, file_name="notes.pdf", mime="application/pdf")