import os
import sys
import cv2
import numpy as np
import gradio as gr
from PIL import Image

# Ensure backend directory and root directory are in sys.path
base_dir = os.path.dirname(os.path.abspath(__file__))
backend_dir = os.path.join(base_dir, "backend")

if os.path.exists(backend_dir) and backend_dir not in sys.path:
    sys.path.insert(0, backend_dir)
if base_dir not in sys.path:
    sys.path.insert(0, base_dir)

try:
    import spaces
except ImportError:
    class spaces:
        @staticmethod
        def GPU(func):
            return func

try:
    from app.services.ai_pipeline import ai_pipeline_service
    from app.kb import kb_service
except ModuleNotFoundError:
    from backend.app.services.ai_pipeline import ai_pipeline_service
    from backend.app.kb import kb_service

# Custom Dark Theme CSS for Gradio
custom_css = """
body, .gradio-container {
    background-color: #0b0f19 !important;
    color: #e2e8f0 !important;
    font-family: 'Inter', system-ui, sans-serif !important;
}
.header-box {
    text-align: center;
    padding: 24px;
    background: linear-gradient(135deg, rgba(30, 41, 59, 0.8), rgba(15, 23, 42, 0.9));
    border: 1px solid rgba(255, 255, 255, 0.1);
    border-radius: 16px;
    margin-bottom: 24px;
    box-shadow: 0 10px 30px rgba(0,0,0,0.5);
}
.header-box h1 {
    font-size: 2.2rem;
    font-weight: 800;
    background: linear-gradient(90deg, #38bdf8, #818cf8, #c084fc);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 8px;
}
.badge-normal {
    background-color: #10b981; color: white; padding: 4px 12px; border-radius: 9999px; font-weight: bold;
}
.badge-severe {
    background-color: #ef4444; color: white; padding: 4px 12px; border-radius: 9999px; font-weight: bold;
}
.badge-moderate {
    background-color: #f59e0b; color: white; padding: 4px 12px; border-radius: 9999px; font-weight: bold;
}
"""

@spaces.GPU
def analyze_scan(image_input):
    if image_input is None:
        return None, None, "<h3 style='color:#ef4444;'>Please upload a medical scan image to analyze.</h3>", {}

    # Save temporary file if image is numpy array
    if isinstance(image_input, np.ndarray):
        temp_path = os.path.abspath("temp_gradio_upload.jpg")
        cv2.imwrite(temp_path, cv2.cvtColor(image_input, cv2.COLOR_RGB2BGR))
        image_path = temp_path
    else:
        image_path = str(image_input)

    filename = os.path.basename(image_path)
    
    # Run full clinical AI pipeline
    res = ai_pipeline_service.run_pipeline(image_path, filename)
    
    if not res.get("success", False):
        error_msg = f"<h3 style='color:#ef4444;'>Analysis Failed</h3><p>{res.get('error', 'Unknown error')}</p>"
        return None, None, error_msg, res

    modality = res.get("image_type", "Unknown")
    disease = res.get("disease", "Normal")
    confidence = res.get("confidence", 0.0)
    severity = res.get("severity", "Normal")
    findings = res.get("findings", {})
    evidence = res.get("evidence", {})
    
    heatmap_path = res.get("heatmap_filepath")
    mask_path = os.path.join(os.path.dirname(heatmap_path), f"mask_{res.get('report_id')}.png") if heatmap_path else None
    
    heatmap_img = Image.open(heatmap_path) if heatmap_path and os.path.exists(heatmap_path) else None
    mask_img = Image.open(mask_path) if mask_path and os.path.exists(mask_path) else None
    
    # Format HTML / Markdown Diagnostic Report
    conf_pct = f"{confidence * 100:.1f}%"
    severity_badge = f"<span class='badge-{severity.lower()}'>{severity}</span>"
    
    report_md = f"""
## 🩺 Clinical Diagnostic Report `#${res.get('report_id')}`

| Metric | Value |
|---|---|
| **Detected Modality** | **`{modality}`** *(Auto-Detected)* |
| **Primary Diagnosis** | **`<span style="color:#38bdf8; font-size:1.3rem;">{disease}</span>`** |
| **Certainty Confidence** | **`{conf_pct}`** |
| **Severity Index** | {severity_badge} |

---

### 🔬 **Clinical Findings**
> **Summary**: {findings.get('summary', 'No findings summary available.')}

**Observation Details**:
"""
    for detail in findings.get("details", []):
        report_md += f"- {detail}\n"

    if evidence and disease != "Normal":
        report_md += f"""
---
### 📚 **Evidence-Based Guidelines & Recommendations**
- **Overview**: {evidence.get('overview', 'N/A')}
- **Key Manifestations**: {", ".join(evidence.get('clinical_manifestations', []))}
- **Recommended Action**:
"""
        for rec in evidence.get("recommendations", []):
            report_md += f"  1. {rec}\n"

    report_md += "\n\n*Disclaimer: MedVision AI output is intended for clinical decision support and researcher validation.*"

    return heatmap_img, mask_img, report_md, res


def get_kb_info(disease_name):
    if not disease_name:
        return "Select a disease to view clinical evidence guidelines."
    info = kb_service.retrieve_info(disease_name)
    md = f"## 📚 Medical Knowledge Base: **{disease_name}**\n\n"
    md += f"**Overview**: {info.get('overview', 'N/A')}\n\n"
    md += f"**ICD-10 Code**: `{info.get('icd_10', 'N/A')}` | **Modality**: `{info.get('modality', 'N/A')}`\n\n"
    md += "### Clinical Manifestations:\n"
    for item in info.get("clinical_manifestations", []):
        md += f"- {item}\n"
    md += "\n### Recommended Clinical Workup:\n"
    for item in info.get("recommendations", []):
        md += f"1. {item}\n"
    return md


# Build Gradio UI Blocks
with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo"), css=custom_css, title="MedVision AI Diagnostic Engine") as demo:
    gr.HTML("""
    <div class="header-box">
        <h1>🩺 MedVision AI — Multi-Modality Diagnostic Engine</h1>
        <p style="color: #94a3b8; font-size: 1.05rem;">
            Instant Auto-Modality Routing (Chest X-Ray • Brain MRI • Eye Fundus) with GradCAM & Contour Explainability
        </p>
    </div>
    """)

    with gr.Tabs():
        with gr.TabItem("🔬 Live Scan Analyzer"):
            with gr.Row():
                with gr.Column(scale=1):
                    scan_input = gr.Image(
                        type="filepath",
                        label="Upload Medical Scan Image",
                        sources=["upload", "clipboard"]
                    )
                    analyze_btn = gr.Button("🚀 Run AI Clinical Analysis", variant="primary", size="lg")
                    
                    # Example scans
                    example_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "test_scans"))
                    ex_files = []
                    if os.path.exists(example_dir):
                        for f in ["sample_normal_xray.jpg", "sample_pneumonia_xray.jpg", "sample_brain_mri.jpg", "sample_glaucoma_fundus.jpg"]:
                            fp = os.path.join(example_dir, f)
                            if os.path.exists(fp):
                                ex_files.append([fp])
                    if ex_files:
                        gr.Examples(examples=ex_files, inputs=scan_input, label="Sample Clinical Scans")

                with gr.Column(scale=2):
                    with gr.Row():
                        out_heatmap = gr.Image(label="🔥 GradCAM Heatmap Overlay", type="pil")
                        out_mask = gr.Image(label="🎯 Contour Segmentation Mask", type="pil")
                    
                    out_report = gr.Markdown(label="Clinical Report")
                    out_json = gr.JSON(label="Structured AI Pipeline JSON", visible=False)

            analyze_btn.click(
                fn=analyze_scan,
                inputs=[scan_input],
                outputs=[out_heatmap, out_mask, out_report, out_json]
            )

        with gr.TabItem("📚 Medical Knowledge Base"):
            with gr.Row():
                disease_dropdown = gr.Dropdown(
                    choices=[
                        "Pneumonia", "Tuberculosis", "Pleural Effusion",
                        "Glioma", "Meningioma", "Pituitary Tumor",
                        "Diabetic Retinopathy", "Glaucoma", "Age Related Macular Degeneration",
                        "Hypertensive Retinopathy", "Normal"
                    ],
                    value="Pneumonia",
                    label="Select Disease / Pathology"
                )
            kb_output = gr.Markdown()
            disease_dropdown.change(fn=get_kb_info, inputs=[disease_dropdown], outputs=[kb_output])
            demo.load(fn=get_kb_info, inputs=[disease_dropdown], outputs=[kb_output])

        with gr.TabItem("ℹ️ About MedVision"):
            gr.Markdown("""
            ### About MedVision AI
            MedVision AI is an ensemble deep-learning medical imaging platform engineered for high-precision diagnostic support across 12 pathologies:
            - **Chest X-ray**: Pneumonia, Tuberculosis, Pleural Effusion, Normal
            - **Brain MRI**: Glioma, Meningioma, Pituitary Tumor, Normal
            - **Eye Fundus**: Diabetic Retinopathy, Glaucoma, AMD, Hypertensive Retinopathy, Normal
            
            Built with PyTorch, TorchXRayVision, MONAI, EfficientNet, ConvNeXt-Base, and GradCAM explainability.
            """)

if __name__ == "__main__":
    demo.launch()