import gradio as gr
from medfusion_pipeline import MedFusionPipeline

pipe = MedFusionPipeline.from_pretrained(".", mode="pro")

def switch_mode(mode):
    pipe.set_mode(mode)
    return f"Mode set to: {mode}"

def analyze(image, mode, max_tokens):
    if mode != pipe.mode:
        pipe.set_mode(mode)
    report = pipe.analyze(image, max_new_tokens=int(max_tokens))
    return report

with gr.Blocks(title="MedFusion-AI") as demo:
    gr.Markdown("# 🩺 MedFusion-AI — Pro & Lite in one")
    with gr.Row():
        with gr.Column(scale=1):
            mode = gr.Radio(choices=["pro","lite"], value="pro", label="Mode")
            max_tokens = gr.Slider(64, 512, value=256, step=32, label="Max tokens")
            set_btn = gr.Button("Apply mode")
            set_msg = gr.Markdown("")
            img = gr.Image(type="filepath", label="Upload X-ray / DICOM")
            run = gr.Button("Analyze")
        with gr.Column(scale=1):
            out = gr.Textbox(label="AI Report", lines=16)

    set_btn.click(fn=switch_mode, inputs=[mode], outputs=[set_msg])
    run.click(fn=analyze, inputs=[img, mode, max_tokens], outputs=[out])

demo.launch(server_name="0.0.0.0", server_port=7860)