import os
import tempfile
import shutil
import logging

os.system('pip install paddlepaddle==2.4.2')
os.system('pip install paddleocr')

from paddleocr import PaddleOCR, draw_ocr
from PIL import Image
import gradio as gr
import torch
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import FileResponse
import uvicorn

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

CUSTOM_PATH = "/gradio"

app = FastAPI()

@app.get("/")
def read_main():
    return {"message": "This is your main app"}

io = gr.Interface(lambda x: "Hello, " + x + "!", "textbox", "textbox")

torch.hub.download_url_to_file('https://i.imgur.com/aqMBT0i.jpg', 'example.jpg')

@app.post("/ocr/")
async def ocr_endpoint(img: UploadFile = File(...), lang: str = Form(...)):
    logger.info("Processing OCR request")
    
    # Save the uploaded image to a temporary file
    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_img:
        shutil.copyfileobj(img.file, temp_img)
        img_path = temp_img.name

    # Perform OCR
    ocr = PaddleOCR(use_angle_cls=True, lang=lang, use_gpu=False)
    result = ocr.ocr(img_path, cls=True)[0]

    boxes = [line[0] for line in result]
    txts = [line[1][0] for line in result]
    scores = [line[1][1] for line in result]

    image = Image.open(img_path).convert('RGB')
    im_show = draw_ocr(image, boxes, txts=None, scores=None, font_path='simfang.ttf')
    im_show = Image.fromarray(im_show)
    result_img_path = 'result.jpg'
    im_show.save(result_img_path)

    # Prepare the response
    response_data = {
        "result_image": result_img_path,
        "ocr_result": result,
        "extracted_text": '\n'.join(txts)
    }

    logger.info("OCR request processed successfully")
    return response_data

app = gr.mount_gradio_app(app, io, path=CUSTOM_PATH)
uvicorn.run(app, host="0.0.0.0", port=7860)