# -----------------------------------------------
# app.py
# Gradio UI for Medical Report Analyzer
# Supports PDF and image uploads
# Hindi and English output
# -----------------------------------------------

import gradio as gr
import os
import shutil
from graph import run_pipeline
from config import SUPPORTED_FORMATS, SUPPORTED_LANGUAGES, RISK_LEVELS, DISCLAIMER

# -----------------------------------------------
# Process uploaded file through pipeline
# -----------------------------------------------
def analyze_report(file, language):
    # Validate file uploaded
    if file is None:
        return (
            "⚠️ Please upload a medical document.",
            "", "", "", "", "", None
        )

    # Validate file format
    ext = os.path.splitext(file.name)[1].lower()
    if ext not in SUPPORTED_FORMATS:
        return (
            f"⚠️ Unsupported format: {ext}\nSupported: {', '.join(SUPPORTED_FORMATS)}",
            "", "", "", "", "", None
        )

    # Run pipeline
    try:
        final_state = run_pipeline(file.name, language)
    except Exception as e:
        return (
            f"❌ Pipeline error: {e}",
            "", "", "", "", "", None
        )

    # Check for errors
    if final_state.get("error") and final_state.get("document_type") == "unknown":
        return (
            f"❌ {final_state['error']}",
            "", "", "", "", "", None
        )

    # -----------------------------------------------
    # Format output sections
    # -----------------------------------------------

    # 1. Severity banner
    severity      = final_state.get("overall_severity", "UNKNOWN")
    severity_text = RISK_LEVELS.get(severity, severity)
    urgency       = final_state.get("urgency_timing", "")
    severity_out  = f"{severity_text}\n{urgency}"

    # 2. Summary
    summary = final_state.get("summary", "")
    if not summary:
        summary = "Analysis complete. Please see detailed findings below."

    # 3. Abnormal values
    abnormal_values = final_state.get("abnormal_values", [])
    urgent_flags    = final_state.get("urgent_flags", [])
    abnormal_out    = ""

    if urgent_flags:
        abnormal_out += "🚨 URGENT FLAGS:\n"
        for flag in urgent_flags:
            abnormal_out += f"  • {flag}\n"
        abnormal_out += "\n"

    if abnormal_values:
        abnormal_out += "⚠️ ABNORMAL VALUES:\n"
        for val in abnormal_values:
            abnormal_out += f"  • {val}\n"
    else:
        abnormal_out = "✅ No significant abnormalities detected."

    # 4. Conditions + Explanation
    matched_conditions  = final_state.get("matched_conditions", [])
    explanation_parsed  = final_state.get("explanation_parsed", {})
    possible_conditions = explanation_parsed.get("possible_conditions", "")
    what_is_abnormal    = explanation_parsed.get("what_is_abnormal", [])

    explanation_out = ""
    if matched_conditions:
        explanation_out += "🔍 POSSIBLY RELATED CONDITIONS:\n"
        for cond in matched_conditions:
            explanation_out += f"  • {cond}\n"
        explanation_out += "\n"

    if possible_conditions:
        explanation_out += f"📋 EXPLANATION:\n{possible_conditions}\n\n"

    if what_is_abnormal:
        explanation_out += "🔬 WHAT YOUR VALUES MEAN:\n"
        for item in what_is_abnormal:
            explanation_out += f"  • {item}\n"

    # 5. Specialist advice
    primary_specialist    = final_state.get("primary_specialist", "")
    secondary_specialists = final_state.get("secondary_specialists", [])
    where_to_go           = final_state.get("where_to_go", "")
    questions             = final_state.get("questions_for_doctor", [])
    what_to_bring         = final_state.get("what_to_bring", [])
    what_to_do            = final_state.get("what_to_do", "")

    specialist_out = ""
    if primary_specialist:
        specialist_out += f"👨‍⚕️ PRIMARY SPECIALIST:\n  {primary_specialist}\n\n"

    if secondary_specialists:
        specialist_out += "👩‍⚕️ OTHER SPECIALISTS:\n"
        for spec in secondary_specialists:
            specialist_out += f"  • {spec}\n"
        specialist_out += "\n"

    if where_to_go:
        specialist_out += f"🏥 WHERE TO GO IN INDIA:\n{where_to_go}\n"

    if what_to_do:
        specialist_out += f"📌 WHAT TO DO NEXT:\n{what_to_do}\n\n"

    if questions:
        specialist_out += "❓ QUESTIONS TO ASK YOUR DOCTOR:\n"
        for q in questions:
            specialist_out += f"  • {q}\n"
        specialist_out += "\n"

    if what_to_bring:
        specialist_out += "🎒 WHAT TO BRING:\n"
        for item in what_to_bring:
            specialist_out += f"  • {item}\n"

    # 6. Report file
    report_path = final_state.get("report_path", "")
    report_file = report_path if report_path and os.path.exists(report_path) else None

    return (
        severity_out,
        summary,
        abnormal_out,
        explanation_out,
        specialist_out,
        DISCLAIMER,
        report_file
    )

# -----------------------------------------------
# Build Gradio UI
# -----------------------------------------------
def build_ui():
    with gr.Blocks(
        title="Medical Report Analyzer",
    ) as app:

        # Header
        gr.Markdown("""
        # 🏥 AI Medical Report Analyzer
        ### Understand your medical reports in simple language
        Supports blood reports, X-rays, MRI reports, prescriptions, and discharge summaries.
        Upload your report and get instant AI-powered analysis in Hindi or English.
        """)

        with gr.Row():
            # Left column — inputs
            with gr.Column(scale=1):
                gr.Markdown("### 📤 Upload Your Report")

                file_input = gr.File(
                    label="Upload Medical Document",
                    file_types=[".pdf", ".jpg", ".jpeg", ".png"],
                    type="filepath"
                )

                language_input = gr.Radio(
                    choices=list(SUPPORTED_LANGUAGES.keys()),
                    value="English",
                    label="Select Output Language"
                )

                analyze_btn = gr.Button(
                    "🔍 Analyze Report",
                    variant="primary",
                    size="lg"
                )

                gr.Markdown("""
                **📸 Photo Tips:**
                - Good lighting
                - Place report flat
                - Take from directly above
                - Text must be clearly visible
                """)

            # Right column — severity output
            with gr.Column(scale=1):
                gr.Markdown("### 🚦 Health Status")
                severity_out = gr.Textbox(
                    label="Overall Risk Level",
                    lines=3,
                    interactive=False
                )

                gr.Markdown("### 📝 Summary")
                summary_out = gr.Textbox(
                    label="Report Summary",
                    lines=5,
                    interactive=False
                )

        # Full width outputs
        gr.Markdown("---")

        with gr.Row():
            with gr.Column():
                gr.Markdown("### ⚠️ Abnormal Values")
                abnormal_out = gr.Textbox(
                    label="Detected Abnormalities",
                    lines=10,
                    interactive=False
                )

            with gr.Column():
                gr.Markdown("### 🔬 Explanation")
                explanation_out = gr.Textbox(
                    label="What Your Results Mean",
                    lines=10,
                    interactive=False
                )

        with gr.Row():
            with gr.Column():
                gr.Markdown("### 👨‍⚕️ Specialist Advice")
                specialist_out = gr.Textbox(
                    label="Doctor Recommendations",
                    lines=12,
                    interactive=False
                )

            with gr.Column():
                gr.Markdown("### 📄 Download Report")
                report_out = gr.File(
                    label="Download Simplified Report (.docx)"
                )

                gr.Markdown("### ⚠️ Disclaimer")
                disclaimer_out = gr.Textbox(
                    label="Medical Disclaimer",
                    lines=6,
                    interactive=False
                )

        # Connect button to function
        analyze_btn.click(
            fn=analyze_report,
            inputs=[file_input, language_input],
            outputs=[
                severity_out,
                summary_out,
                abnormal_out,
                explanation_out,
                specialist_out,
                disclaimer_out,
                report_out
            ]
        )

    return app

# -----------------------------------------------
# Main
# -----------------------------------------------
if __name__ == "__main__":
    print("=== Medical Report Analyzer — Starting UI ===")
    app = build_ui()
    
    # HuggingFace sets SPACE_ID env variable automatically
    # Use it to detect deployment environment
    import os
    is_huggingface = os.getenv("SPACE_ID") is not None
    
    if is_huggingface:
        # HuggingFace Docker deployment
        app.launch(server_name="0.0.0.0", server_port=7860)
    else:
        # Local development
        app.launch(server_name="127.0.0.1", server_port=7861)