import gradio as gr
import os
from model import generate_report
from utils import convert_dicom_zip_to_nifti, 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 process_scan(zip_file):
    if zip_file is None:
        return gr.Gallery(visible=False), gr.Button(interactive=False), None

    sequence_images = convert_dicom_zip_to_nifti(zip_file)

    if not sequence_images:
        gr.Warning("No valid sequences found in ZIP.")
        return gr.Gallery(visible=False), gr.Button(interactive=False), None

    gallery_items = [(img, name) for name, img in sequence_images]

    return (
        gr.Gallery(value=gallery_items, visible=True),
        gr.Button(interactive=True),
        sequence_images        # ← saved into State
    )
def run_generate(sequence_images):
    if not sequence_images:
        return "", gr.DownloadButton(visible=False)

    PRIORITY_KEYWORDS = [
    "flair", "tirm",            # T2 FLAIR — highest priority
    "dwi", "diff",              # DWI
    "t1_se_tra_320+fs",         # T1 + Contrast (has 'C' in name)
    "t1",                       # T1
    "t2_tse_tra", "t2_tra",     # T2 axial
    ]

    # Score each sequence by priority
    def get_priority(name_img_tuple):
        name = name_img_tuple[0].lower()
        for i, keyword in enumerate(PRIORITY_KEYWORDS):
            if keyword in name:
                return i
        return 999  # unknown sequences go last

    # Sort by priority and take top 4
    sorted_sequences = sorted(sequence_images, key=get_priority)
    selected = sorted_sequences[:5]

    print("Sequences sent to model:")
    for name, _ in selected:
        print(f"  - {name}")

    images = [img for _, img in selected]
    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
        None
    )

# 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.")

    sequence_state = gr.State(value=None)
    
    with gr.Row():
        # Left column - image input/display
        with gr.Column():
            upload_btn = gr.File(
                label="Upload DICOM Study as ZIP",
                file_types=[".zip"],
                file_count="single"
            )
            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, sequence_state]
    )

    generate_btn.click(
        fn=run_generate,
        inputs=sequence_state,
        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, sequence_state]
)

app.launch()