import gradio as gr
import time
from src.pipeline import generate_report

# Pre-load models
from src.tools_loader import get_tools
_ = get_tools()

def process_upload(image_path: str):
    if image_path is None:
        yield "Please upload a chest X-ray image."
        return

    start = time.time()
    yield "Analyzing image..."
    
    report = generate_report(image_path)
    elapsed = time.time() - start
    
    yield f"### Radiology Report\n{report}\n\n*Generated in {elapsed:.1f}s*"

with gr.Blocks(theme=gr.themes.Soft(), title="Radiology Assistant") as demo:
    gr.Markdown("#Multi-Agent Radiology Assistant")
    
    image_input = gr.Image(type="filepath", label="Upload Chest X-ray")
    generate_btn = gr.Button("Generate Report", variant="primary")
    report_output = gr.Markdown("Upload an image and click Generate Report.")
    
    generate_btn.click(process_upload, image_input, report_output)


gr.Markdown(
        "Download and use any frontal chest X-ray PNG/JPG file from the internet and click **Generate Report**.\n"
        "### NOTE: This is just a demo. It is not intended to diagnose or suggest treatment of any disease or condition, and should not be used for medical advice."
    )

if __name__ == "__main__":
    demo.launch()