from fastapi import FastAPI, Form, UploadFile, File
from fastapi.responses import JSONResponse
from PIL import Image
import torch, clip, io, re, numpy as np, cv2
from pix2text import Pix2Text
import pytesseract
import time
import os
import logging

# -------------------------------
# Load models
# -------------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
clip_model, preprocess = clip.load("ViT-B/32", device=device)

# Prompts
math_labels = [
    "a math equation",
    "a handwritten equation",
    "scientific symbols",
    "algebraic formula",
]
text_labels = [
    "a page of text",
    "a handwritten note",
    "a printed paragraph",
    "a sentence of text",
]
all_labels = math_labels + text_labels
text_tokens = clip.tokenize(all_labels).to(device)

# Math-aware OCR
p2t = Pix2Text()

# Regex for math symbols
MATH_SYMBOLS = re.compile(r"[=+\-*/^√∫Σ∑π∞≤≥≠<>]")


# -------------------------------
# Helpers
# -------------------------------
def preprocess_text_image(image: Image.Image) -> Image.Image:
    """Enhance text images for better OCR."""
    img = np.array(image)
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    gray = cv2.convertScaleAbs(gray, alpha=1.5, beta=20)  # contrast
    _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return Image.fromarray(thresh)


def clean_text(text: str) -> str:
    """Remove noise like form feeds and empty lines."""
    return "\n".join(
        [line.strip() for line in text.splitlines() if line.strip()]
    ).replace("\f", "")


# -------------------------------
# Classifier
# -------------------------------
def classify_image_type(image: Image.Image, threshold: float = 0.6):
    image_tensor = preprocess(image).unsqueeze(0).to(device)
    with torch.no_grad():
        logits_per_image, _ = clip_model(image_tensor, text_tokens)
        probs = logits_per_image.softmax(dim=-1).cpu().numpy()[0]

    math_score = float(sum(probs[: len(math_labels)]))
    text_score = float(sum(probs[len(math_labels) :]))

    raw_text = pytesseract.image_to_string(image)
    math_chars = len(MATH_SYMBOLS.findall(raw_text))
    ratio = math_chars / max(1, len(raw_text))

    if math_score > threshold or ratio > 0.25:
        return "math", math_score
    elif text_score > threshold:
        return "text", text_score
    else:
        return "uncertain", max(math_score, text_score)


# -------------------------------
# FastAPI App
# -------------------------------
app = FastAPI()


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

        # make folder name with date
        folder_name = f"received_images/{time.strftime('%Y%m%d')}"

        # create folder if not exists
        os.makedirs(folder_name, exist_ok=True)

        # save uploaded image in that folder
        image.save(os.path.join(folder_name, file.filename))

        print("Image saved.")

        label, conf = classify_image_type(image)

        if (category_slug == "math" or (category_slug == "other" and label == "math")) or (category_slug == "physics"):
            # Use Pix2Text
            result = p2t.recognize(image)

            # Ensure string output (LaTeX)
            if isinstance(result, list):
                # Pix2Text sometimes returns list of dicts
                question = " ".join([r["text"] if isinstance(r, dict) else str(r) for r in result])
            else:
                question = str(result).strip()
        else:
            # Normal OCR
            question = pytesseract.image_to_string(image)
            question = question.strip()

        response = {
            "status": True,
            "message": "Image processed successfully",
            "data": {"text": question},
        }
        return JSONResponse(content=response)


        # contents = await file.read()
        # image = Image.open(io.BytesIO(contents)).convert("RGB")

        # label, conf = classify_image_type(image)

        # if label == "math":
        #     latex = p2t(image)
        #     return {
        #         "filename": file.filename,
        #         "type": "math",
        #         "confidence": conf,
        #         "latex": latex,
        #     }

        # elif label == "text":
        #     # question = pytesseract.image_to_string(image)
        #     # processed = preprocess_text_image(image)
        #     text = pytesseract.image_to_string(image)
        #     print(text)
        #     return {
        #         "filename": file.filename,
        #         "type": "text",
        #         "confidence": conf,
        #         "text": text,
        #     }

        # else:  # uncertain → return both
        #     processed = preprocess_text_image(image)
        #     text = clean_text(pytesseract.image_to_string(processed))
        #     latex = p2t(image)
        #     return {
        #         "filename": file.filename,
        #         "type": "uncertain",
        #         "confidence": conf,
        #         "text": text,
        #         "latex": latex,
        #     }

    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})