from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse
from PIL import Image
import numpy as np
import cv2
import io
import requests
from typing import Optional
from pydantic import BaseModel
from rapidocr_onnxruntime import RapidOCR

app = FastAPI(title="OCR API", description="RapidOCR API Service")
engine = RapidOCR()

def process_ocr(image: np.ndarray, use_det: bool, use_cls: bool, use_rec: bool) -> list:
    """Xử lý OCR và trả về kết quả"""
    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    result, _ = engine(img_bgr, use_det=use_det, use_cls=use_cls, use_rec=use_rec)
    
    if not result:
        return []
    
    texts = []
    for item in result:
        if len(item) == 3:
            box, text, score = item
            texts.append({
                "text": text,
                "confidence": float(score),
                "bbox": box.tolist() if hasattr(box, 'tolist') else box
            })
        elif len(item) == 2:
            _, text = item
            texts.append({
                "text": str(text),
                "confidence": None,
                "bbox": None
            })
    
    return texts

@app.post("/ocr")
async def ocr_from_file(
    file: UploadFile = File(..., description="Image file to process"),
    use_det: bool = Form(True, description="Use text detection"),
    use_cls: bool = Form(True, description="Use classification"),
    use_rec: bool = Form(True, description="Use recognition")
):
    if not file.content_type or not file.content_type.startswith('image/'):
        raise HTTPException(status_code=400, detail="File must be an image")
    
    try:
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))
        
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        img_np = np.array(image)
        results = process_ocr(img_np, use_det, use_cls, use_rec)
        
        return JSONResponse(content={
            "success": True,
            "texts": [item["text"] for item in results],
            "details": results,
            "num_texts": len(results)
        })
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {"status": "healthy", "service": "XOCR AI"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)