import gradio as gr
import os
from loader import loader
from tutor import run_council_deliberation
from inference import generate_universal_heatmap

# NerdMedica Dark Tech CSS
CSS = """
body, .gradio-container { background-color: #0A0E1A !important; color: #D1D5DB !important; }
.gr-button-primary { background: linear-gradient(90deg, #E5534B, #A855F7) !important; border: none !important; border-radius: 8px !important; color: white !important; font-weight: bold !important; }
#header { border-bottom: 2px solid #E5534B; padding-bottom: 15px; margin-bottom: 25px; }
.message.user { background-color: #1F2937 !important; border-left: 4px solid #E5534B !important; }
.message.bot { background-color: #111827 !important; border-left: 4px solid #3B82F6 !important; }
"""

def medical_audit_pipeline(image, modality, question, history):
    loader.clear_vram()
    
    # 1. Visual Evidence (If image exists)
    heatmap = None
    if image is not None:
        clip_model, preprocess = loader.load_biomed_clip()
        heatmap = generate_universal_heatmap(image, question, clip_model, preprocess)
    
    # 2. Council Deliberation
    response = run_council_deliberation(image, modality, question)
    
    # Append to Chat History
    history.append({"role": "user", "content": question})
    history.append({"role": "assistant", "content": response})
    
    return history, heatmap, response

# Gradio 6.0: We build the blocks without passing theme/css here
with gr.Blocks() as demo:
    gr.HTML("<div id='header'><h1 style='color:white; font-family:Consolas;'>🧬 MediVance LLC | NerdMedica Council</h1></div>")
    
    with gr.Row():
        with gr.Column(scale=1):
            img_in = gr.Image(type="pil", label="Medical Scan (Optional)")
            modality_drop = gr.Dropdown(
                ["General Inquiry", "Chest X-Ray", "CT Scan", "MRI", "Pathology"], 
                value="General Inquiry", 
                label="Modality Focus"
            )
            heatmap_out = gr.Image(label="Visual Evidence Audit")
            
        with gr.Column(scale=2):
            chat = gr.Chatbot(label="Council Deliberation", type="messages", height=500)
            query = gr.Textbox(placeholder="Ask about 'Isolated Hypertension' or specific scan findings...", label="Doctor's Inquiry")
            btn = gr.Button("Consult the Council", variant="primary")

    with gr.Accordion("🩺 Official Clinical Audit Report", open=False):
        audit_box = gr.Textbox(label="Final Finding (Editable)", interactive=True, lines=4)
        status = gr.Radio(["Approved", "Corrected", "Flagged"], label="Clinical Signature")
        gr.Button("Archive to MediVance API", variant="secondary")

    # Wire up the button
    btn.click(
        fn=medical_audit_pipeline, 
        inputs=[img_in, modality_drop, query, chat], 
        outputs=[chat, heatmap_out, audit_box]
    )

# Gradio 6.0: Theme and CSS MUST be passed in the launch method
demo.launch(
    ssr_mode=False,
    theme=gr.themes.Base(),
    css=CSS
)