import gradio as gr
import requests
import os
import urllib
import time

document_types = {"Annual Financial Statement": "afs", "Invoice": "i", "Purchase Order": "po", "Bank Statement": "bs"}

API_KEY = os.getenv("API_KEY", None)
SERVER_URL = os.getenv("SERVER_URL", None)

with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.red, secondary_hue=gr.themes.colors.pink)) as demo:
    task_id = gr.State()
    if SERVER_URL is not None:
        file_input = gr.File(label="Upload a file", file_types=['.pdf'])
        document_type = gr.Dropdown(label="Select Document Type", choices=list(document_types.keys()))
        api_key = gr.Text(label="Enter API key", value=API_KEY if API_KEY else None, placeholder="Enter your API key here", visible=API_KEY is None)
        with gr.Row():
            clear_button = gr.Button("Clear")
            process_button = gr.Button("Submit")

        status_output = gr.Markdown(label="Current Status")
        file_upload_status = gr.Text(label="File Upload Status", interactive=False)
    else:
        gr.Markdown("Server URL is not configured. Please set the SERVER_URL environment variable.")
        
    
    def get_result(task_id, api_key, progress=gr.Progress()):
        if task_id is not None:
            current_status = ""
            total_steps = 1000
            progress(0, desc="Starting...")
            for i in progress.tqdm(range(total_steps)):
                response = requests.get(f"{SERVER_URL}/verify/get_results/{task_id}", headers={"X-Api-Key": api_key})
                current_status = response.json()['current_status']
                if current_status != "In Progress":
                    break
                
                time.sleep(1)
            return response.json()
    
    @gr.render(inputs=task_id)
    def show_results(progress_id):
        if progress_id is not None:
            view_results = gr.Button("View Results")
            view_results.click(fn=get_result, inputs=[task_id, api_key], outputs=gr.Json())

    def process_file(file, api_key, document_type, progress=gr.Progress()):
        if file is not None and document_type is not None:
            params = {"document-type": document_types[document_type]}
            url_encoded_file_name = urllib.parse.quote(os.path.basename(file.name))
            print(f"filename encoded {url_encoded_file_name}")
            response = requests.get(f"{SERVER_URL}/verify/get_upload_url/{url_encoded_file_name}", headers={"X-Api-Key": api_key, "Accept": "application/json"}, params=params)
            if response.status_code == 200:
                upload_url = response.json()['url']
                task_id = response.json()['task_id']
                print(f"presigned url: {upload_url}, Task Id: {task_id}")
                with open(file.name, 'rb') as f:
                    file_content = f.read()
                    file_size = len(file_content)
                    # Upload the entire file
                    upload_response = requests.put(upload_url, data=file_content, headers={'Content-Type': 'application/pdf'})
                    
                    # Update the progress bar
                    progress(1, desc="Uploading...")
                    
                    print(f"Uploaded {file_size} bytes")
                
                return task_id, "Upload successful"
        return None, "Upload failed"
    
    def clear_fields():
        return None, None, None

    if SERVER_URL is not None:
        process_button.click(fn=process_file, inputs=[file_input, api_key, document_type], outputs=[task_id, file_upload_status])
        clear_button.click(fn=clear_fields, inputs=None, outputs=[file_input, document_type, task_id])

demo.launch()