import gradio as gr
import os
from utils import load_dicoms
from model import generate_report
from utils import load_dicoms, generate_pdf

def process_scan(dicom_files):
    if not dicom_files:
        return gr.Gallery(visible=False), gr.Button(interactive=False)
    
    images = load_dicoms(dicom_files)
    if not images:
        return gr.Gallery(visible=False), gr.Button(interactive=False)
        
    return gr.Gallery(value=images, visible=True), gr.Button(interactive=True)

def run_generate(dicom_files):
    if not dicom_files:
        return "", gr.DownloadButton(visible=False)
    
    images = load_dicoms(dicom_files)
    if not images:
        return gr.DownloadButton(visible=False)
    
    report = generate_report(images)
    
    return report, gr.DownloadButton(visible=True)

def download_report(report_text):
    if not report_text or not report_text.strip():
        gr.Warning("No report to download yet.")
        return None
    return generate_pdf(report_text)

def reset_on_clear():
    return (
        gr.Gallery(value=None, visible=False),  # scan_display
        "",                             # report_box
        gr.Button(interactive=False),   # generate_btn
        gr.DownloadButton(visible=False) # download_btn
    )

# Build the UI
with gr.Blocks(title="MRI Brain Report Generator") as app:
    
    gr.Markdown("# MRI Brain Scan Report Generator")
    gr.Markdown("### Radiologist Assistant")
    gr.Markdown("Upload DICOM brain scans to generate a report.")
    
    with gr.Row():
        # Left column - image input/display
        with gr.Column():
            upload_btn = gr.File(
                label="Upload MRI Sequences (T1, T2, FLAIR, DWI, ADC, GRE, T1+C)",
                file_types=[".dcm"],
                file_count="multiple"
            )
            scan_display = gr.Gallery(
                label="Scan Preview",
                columns=3,
                height="auto",
                visible=False
            )
        
        # Right column - report (empty for now)
        with gr.Column():
            report_box = gr.Textbox(
                label="Generated Report",
                placeholder="Report will appear here...",
                lines=15,
                interactive=True   # makes it editable
            )
    
    # Buttons row
    with gr.Row():
        generate_btn = gr.Button(
            "Generate Report", 
            variant="primary",
            interactive=False    # disabled by default
        )
        download_btn = gr.DownloadButton("Download PDF", visible=False)

    # When file is uploaded → show preview AND enable generate button
    upload_btn.change(
        fn=process_scan,
        inputs=upload_btn,
        outputs=[scan_display, generate_btn]
    )

    generate_btn.click(
        fn=run_generate,
        inputs=upload_btn,
        outputs=[report_box, download_btn]
    )

    download_btn.click(
        fn=download_report,
        inputs=report_box,
        outputs=download_btn
    )

    upload_btn.clear(
    fn=reset_on_clear,
    inputs=None,
    outputs=[scan_display, report_box, generate_btn, download_btn]
)

app.launch()