"""
MedPanel - HuggingFace Spaces Gradio App
"""

import gradio as gr
import torch
import json
import os
from PIL import Image

# Import your MedPanel functions
from medpanel import initialize_models, run_medpanel

# Initialize models
print("🚀 Initializing MedPanel...")
HF_TOKEN = os.environ.get('HF_TOKEN')
initialize_models(HF_TOKEN)
print("✅ Ready!")


def analyze_case(image, clinical_notes):
    """
    Main analysis function for Gradio interface
    
    Args:
        image: PIL Image or None
        clinical_notes: str
    
    Returns:
        JSON string with results
    """
    try:
        if not clinical_notes or len(clinical_notes.strip()) < 10:
            return json.dumps({
                "success": False,
                "error": "Please provide clinical notes (at least 10 characters)"
            }, indent=2)
        
        # Run MedPanel
        result = run_medpanel(image, clinical_notes)
        
        # Parse report
        report = result["final_report"]
        if isinstance(report, dict) and "raw_response" in report:
            try:
                raw = report["raw_response"]
                if not raw.strip().endswith('}'):
                    last_complete = raw.rfind('",')
                    if last_complete > 0:
                        raw = raw[:last_complete+2] + '\n}'
                report = json.loads(raw)
            except:
                pass
        
        # Return formatted response
        response = {
            "success": True,
            "report": report,
            "trace": result["panel_trace"]
        }
        
        return json.dumps(response, indent=2)
        
    except Exception as e:
        return json.dumps({
            "success": False,
            "error": str(e)
        }, indent=2)


# Create Gradio Interface
with gr.Blocks(theme=gr.themes.Soft(), title="MedPanel API") as demo:
    gr.Markdown("""
    # 🏥 MedPanel - Multi-Agent Clinical AI
    
    **Multi-specialist AI system for clinical decision support**
    
    *Built for Google MedGemma Impact Challenge 2025*
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Input")
            image_input = gr.Image(
                type="pil",
                label="📷 Medical Image (Optional)",
                height=300
            )
            notes_input = gr.Textbox(
                lines=8,
                label="📝 Clinical Notes & Symptoms (Required)",
                placeholder="""Example:
65 year old male. Persistent cough for 6 weeks.
Night sweats, 8kg weight loss over 2 months.
Low grade fever. Recently moved from high TB prevalence region.
No prior TB diagnosis. Mild fatigue."""
            )
            submit_btn = gr.Button("▶ Run Panel Review", variant="primary", size="lg")
        
        with gr.Column(scale=2):
            gr.Markdown("### Results")
            output = gr.JSON(label="📋 MedPanel Report")
    
    gr.Markdown("""
    ---
    ### About
    - 🩻 **Radiologist Agent** - Analyzes medical images
    - 🩺 **Internist Agent** - Analyzes symptoms
    - 📚 **Evidence Reviewer** - Searches PubMed
    - 😈 **Devil's Advocate** - Challenges diagnoses
    - 🎯 **Orchestrator** - Synthesizes final report
    
    **⚠️ Disclaimer:** This is a proof-of-concept for research purposes only. Not for actual medical use.
    
    **API Access:** You can call this Space programmatically via the API endpoint shown below.
    """)
    
    # Examples
    gr.Examples(
        examples=[
            [
                None,
                """45 year old female. Severe headache for 3 days.
Fever, stiff neck, photophobia.
No recent travel. No known sick contacts."""
            ],
            [
                None,
                """65 year old male. Persistent cough for 6 weeks.
Night sweats, 8kg weight loss over 2 months.
Low grade fever. Recently moved from high TB prevalence region."""
            ]
        ],
        inputs=[image_input, notes_input],
        label="Try Sample Cases"
    )
    
    submit_btn.click(
        fn=analyze_case,
        inputs=[image_input, notes_input],
        outputs=output
    )

# Launch
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )