from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from transformers import AutoModelForImageClassification, AutoProcessor
from PIL import Image
import io
import fitz  # PyMuPDF

app = FastAPI()

# Load the model and processor
model_name = "AsmaaElnagger/Diabetic_RetinoPathy_detection"
model = AutoModelForImageClassification.from_pretrained(model_name)
processor = AutoProcessor.from_pretrained(model_name)

# Convert PDF to images using PyMuPDF
def pdf_to_images_pymupdf(pdf_data):
    try:
        pdf_document = fitz.open(stream=pdf_data, filetype="pdf")
        images = []
        for page_num in range(pdf_document.page_count):
            page = pdf_document.load_page(page_num)
            pix = page.get_pixmap()
            img_data = pix.tobytes("jpeg")  # or "png"
            images.append(img_data)
        return images
    except Exception as e:
        print(f"Error converting PDF: {e}")
        return None

# Classification endpoint
@app.post("/classify")
async def classify_file(file: UploadFile = File(...)):
    file_type = file.filename.rsplit('.', 1)[-1].lower()
    file_data = await file.read()

    try:
        # Handle image file
        if file_type in ['jpg', 'jpeg', 'png', 'gif']:
            image = Image.open(io.BytesIO(file_data)).convert("RGB")

        # Handle PDF file
        elif file_type == 'pdf':
            images = pdf_to_images_pymupdf(file_data)
            if not images:
                return JSONResponse(status_code=500, content={"error": "PDF conversion failed."})
            image = Image.open(io.BytesIO(images[0])).convert("RGB")

        # Unsupported file
        else:
            return JSONResponse(status_code=400, content={"error": "Unsupported file type."})

        # Process and predict
        inputs = processor(images=image, return_tensors="pt")
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_class_idx = logits.argmax(-1).item()
        result = model.config.id2label[predicted_class_idx]

        return {"result": result}

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