import time
import requests
import gradio as gr

def read_image_from_stream(image_path, subscription_key, endpoint):
    with open(image_path, 'rb') as image_file:
        image_data = image_file.read()

    headers = {
        'Ocp-Apim-Subscription-Key': subscription_key,
        'Content-Type': 'application/octet-stream'
    }

    response = requests.post(f"{endpoint}/vision/v3.2/read/analyze", headers=headers, data=image_data)

    if response.status_code == 202:
        read_response_headers = response.headers
        operation_location = read_response_headers["Operation-Location"]
    else:
        raise Exception(f"Unexpected response status: {response.status_code}")

    return operation_location

def get_read_result(operation_location, subscription_key):
    headers = {
        'Ocp-Apim-Subscription-Key': subscription_key
    }
    
    while True:
        response = requests.get(operation_location, headers=headers)
        if response.status_code == 200:  # HTTP 200 indicates success
            read_result = response.json()
            status = read_result['status']
            if status == 'succeeded':
                break
            elif status in ['failed', 'notStarted', 'running']:
                time.sleep(1)  # Wait before polling again
            else:
                raise Exception(f"Unexpected status: {status}")
        else:
            raise Exception(f"Unexpected response status: {response.status_code}")
    
    return read_result

def process_image(image_path, subscription_key, endpoint):
    operation_location = read_image_from_stream(image_path, subscription_key, endpoint)
    operation_id = operation_location.split("/")[-1]
    operation_location = f"{endpoint}/vision/v3.2/read/analyzeResults/{operation_id}"
    
    read_result = get_read_result(operation_location, subscription_key)
    
    if read_result['status'] == 'succeeded':
        output = []
        for text_result in read_result['analyzeResult']['readResults']:
            for line in text_result['lines']:
                output.append(line['text'])
        return " ".join(output).replace("\n", " ")  # Join lines and replace newlines with spaces
    else:
        return "Processing failed or did not succeed."

def main():
    interface = gr.Interface(
        fn=process_image,
        inputs=[
            gr.Image(type="filepath", label="Upload Image"),
            gr.Textbox(label="Subscription Key"),
            gr.Textbox(label="Endpoint URL")
        ],
        outputs=gr.Textbox(),
        title="Azure OCR with Gradio",
        description="Upload an image, provide your Azure subscription key and endpoint URL to perform OCR and get the text."
    )
    interface.launch()

if __name__ == "__main__":
    main()