import gradio as gr
import requests
import time
import os
from dotenv import load_dotenv
from huggingface_hub import list_models

load_dotenv()
API_URL = os.getenv("API_URL")
class CreateProfile:
    def __init__(self, _id=None):
        self._id = _id
        self.name = ""
        self.username = ""
        
newprofile = CreateProfile()

def check_active_users():
    response = requests.get(f"{API_URL}/status")
    if response.status_code == 200:
        queue_length = response.json().get("queue_length", 0)
        if queue_length > 2:
            return False, "There's currently a high volume of requests. The server is busy processing these requests. Please try again later."
    return True, ""

def process_ui_image(image, user_task, progress=gr.Progress()):
    progress(0)
    
    # Check the number of active users
    can_proceed, message = check_active_users()
    if not can_proceed:
        return message
    
    # Upload image and user task to the API
    files = {'image': open(image, 'rb')}
    data = {'user_task': user_task, "name": newprofile.name, "username": newprofile.username}
    response = requests.post(f"{API_URL}/process", files=files, data=data)
    
    if response.status_code != 200:
        return f"Error: {response.text}"
    
    # Get the request ID from the response
    request_id = response.json().get("request_id")
    print("request_id",request_id)
    
    if not request_id:
        return "Error: Failed to get request ID from the server."
    
    # Poll the status of the request
    i = 0.01
    while True:
        status_response = requests.get(f"{API_URL}/status/{request_id}")
        status = status_response.json()
        
        if status.get("status") == "completed":
            progress(1)
            break
        elif status.get("status") == "failed":
            return "Error: Processing failed."

        progress(i)
        i += 0.01
        if i >= 0.9:
            i = 0.9
        time.sleep(1)
    
    # Download the final PDF
    final_pdf_response = requests.get(f"{API_URL}/get_file")
    
    if final_pdf_response.status_code == 200:
        with open("UXEva_report.pdf", "wb") as f:
            f.write(final_pdf_response.content)
        return "UXEva_report.pdf"
    else:
        return "Error: Failed to download the PDF."

def update_submit_button(image, user_task):
    if not image or not user_task:
        return gr.update(interactive=False), gr.update(visible=True)
    return gr.update(interactive=True), gr.update(visible=False)


with gr.Blocks(css="""
    .image-preview img {max-height: 400px; max-width: 400px;}
    .disabled {cursor: not-allowed; color: red;}
""") as demo:
    gr.Markdown("""
    <h1 style="text-align: center;">
        <span style="color: green;">UXEva</span>: UI/UX Evaluation Tool Powered by LLM Agent
    </h1>
    """)
    gr.LoginButton()
    
    with gr.Column(visible=False) as main_ui:
        image_input = gr.Image(type="filepath", label="Upload UI Image", elem_classes="image-preview")
        user_task_input = gr.Textbox(label="Enter User Task")
        pdf_output = gr.File(label="Download PDF Report")
        submit_btn = gr.Button("Submit", interactive=False)
        warning = gr.Markdown("<span class='disabled'>Upload image and enter user task</span>", visible=False)
    
    image_input.change(update_submit_button, inputs=[image_input, user_task_input], outputs=[submit_btn, warning])
    user_task_input.change(update_submit_button, inputs=[image_input, user_task_input], outputs=[submit_btn, warning])
    
    submit_btn.click(
        process_ui_image, 
        inputs=[image_input, user_task_input], 
        outputs=[pdf_output]
    )
    def get_profile(profile: gr.OAuthProfile):
        global newprofile
        if profile is None:
            print("unknow user.")
        else:
            newprofile.name = profile.name
            newprofile.username = profile.username
            print("profile name", profile.name)
            print("profile username", profile.username)
            return {main_ui: gr.Column(visible=True)}
            
    demo.load(get_profile, inputs=None, outputs=[main_ui])

demo.queue(default_concurrency_limit=1, max_size=7)
demo.launch()