from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi import FastAPI, Request, UploadFile, File, Form
from PIL import Image
from pix2text import Pix2Text
import io
import re
import threading
import pytesseract

app = FastAPI()

# ✅ Allow all origins (for development)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ✅ Load Pix2Text for OCR
ocr = Pix2Text()


# ✅ LaTeX wrapper (optional)
def wrap_latex(text: str) -> str:
    def wrap_block(match):
        content = match.group(0)
        if content.startswith("$$"):
            return content
        return f"$$ {content.strip()} $$"

    latex_block_pattern = r"(\\(?:frac|sqrt|sum|prod|int|[a-zA-Z]+)[^$]*)"
    text = re.sub(latex_block_pattern, wrap_block, text)

    power_pattern = r"([a-zA-Z0-9]+(?:\^\{[^}]+\}|\^\d+))"
    text = re.sub(power_pattern, wrap_block, text)

    return text


# ✅ OCR cleaning for Pix2Text output
def extract_clean_latex(page) -> str:
    latex_blocks = []
    for el in page.elements:
        text = el.text.strip()
        text = re.sub(r"^\${2}\s*|\s*\${2}$", "", text)
        text = text.replace("\\quad", " ").replace("\\qquad", " ")
        text = re.sub(r"\s+", " ", text).strip()
        latex_blocks.append(text)
    return "\n\n".join(latex_blocks)


@app.post("/image")
async def image_upload(file: UploadFile = File(...), category_slug: str = Form(1)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents)).convert("RGB")

    try:
        if (category_slug == "math") or (category_slug == "physics"):
            page = ocr(image)
            print("OCR Page:", page)

            question = extract_clean_latex(page)
            if not question.strip():
                question = "0"

        else:
            # Normal OCR
            question = pytesseract.image_to_string(image)
            question = question.strip()

        # Status
        status = bool(question.strip())

        response = {
            "status": status,
            "message": "Image processed successfully",
            "data": {"text": question},
        }
        return response

    except Exception as e:
        print("OCR failed:", str(e))
        return {
            "status": 500,
            "message": f"OCR failed: something went wrong",
            "data": None,
        }