import gradio as gr
import requests
from io import BytesIO
import os

def call_predict_api(image, private_key=None):
    if private_key != os.environ["password"]:
        return "Invalid API key. Please check your key and try again."
    # Convert PIL images to base64 string
    src_buffer = BytesIO()
    image.save(src_buffer, format='PNG')
    src_buffer.seek(0)
    
    # Prepare files for upload
    files = {
        "image": ("src_image.png", src_buffer, "image/png"),
    }
    headers = {"X-API-Key": os.environ["api_key"]}
    response = requests.get(os.environ["endpoint"],
        files=files,
        headers=headers,
        stream=True
    )
    accumulated = ""
    for line in response:
        if line:
            accumulated += line.decode("utf-8")#.replace("\n", "\n\n")
            yield accumulated

    return accumulated

if __name__ == "__main__":
    iface = gr.Interface(
        fn=call_predict_api,
        inputs=[
            gr.Image(
                    sources=["upload"],
                    type="pil",
                    label="Garment Image",
                    width=512,
                    height=512,
                ),
            gr.Textbox(
                label="private-key",
                placeholder="Enter your private key",
                type="password",
                lines=1,
            ),
        ],
        outputs=[
                gr.Markdown(label="Result")
                 ],
        title="BBC-OCR",
        description="A simple OCR tool to extract text from images.",
    )
    iface.launch(show_api=False)