ui.py

"""
ui.py - Gradio 6.x compatible UI - Extended timeout version
"""

import json
import logging
import os
from pathlib import Path

import gradio as gr
import httpx

logger = logging.getLogger(__name__)

BACKEND_URL = "http://127.0.0.1:8000"
REQUEST_TIMEOUT = 600  # 10 minutes — Phi-2 on CPU is slow

CUSTOM_CSS = """
.gr-markdown h1 { font-size: 2rem !important; color: #1a4a8a !important; font-weight: 700; border-bottom: 3px solid #1a4a8a; padding-bottom: 8px; }
.gr-markdown h3 { color: #1a4a8a !important; font-weight: 600; }
#report-box textarea { font-family: 'Courier New', monospace !important; font-size: 13px !important; line-height: 1.7 !important; background: #1e1e2e !important; color: #e0e0e0 !important; border: 1.5px solid #1a4a8a !important; border-radius: 8px !important; }
"""

def _check_backend() -> bool:
    try:
        r = httpx.get(f"{BACKEND_URL}/api/health", timeout=5)
        return r.status_code == 200
    except Exception:
        return False


def _call_generate_report(image_path, patient_name, patient_age, patient_gender,
                           symptoms, clinical_history, modality, body_region, view):
    with open(image_path, "rb") as f:
        image_bytes = f.read()
    filename = Path(image_path).name

    with httpx.Client(timeout=REQUEST_TIMEOUT) as client:
        response = client.post(
            f"{BACKEND_URL}/api/generate-report",
            files={"image": (filename, image_bytes, "application/octet-stream")},
            data={
                "patient_name": patient_name,
                "patient_age": str(patient_age),
                "patient_gender": patient_gender,
                "symptoms": symptoms,
                "clinical_history": clinical_history,
                "modality": modality,
                "body_region": body_region,
                "view": view,
            },
        )
    response.raise_for_status()
    return response.json()


def generate_report(image_file, patient_name, patient_age, patient_gender,
                    symptoms, clinical_history, modality, body_region, view):
    if image_file is None:
        return "", "", "⚠️ Please upload a radiology image.", "", "❌ No image uploaded."
    if not patient_name or not patient_name.strip():
        return "", "", "⚠️ Please enter the patient name.", "", "❌ Patient name required."
    if not patient_age or int(patient_age) <= 0:
        return "", "", "⚠️ Please enter a valid patient age.", "", "❌ Invalid age."
    if not symptoms.strip():
        return "", "", "⚠️ Please enter presenting symptoms.", "", "❌ Symptoms required."
    if not _check_backend():
        return "", "", "❌ Backend offline. Run: uvicorn backend.main:app --reload", "", "❌ Backend offline."

    try:
        data = _call_generate_report(
            image_path=image_file,
            patient_name=patient_name.strip(),
            patient_age=int(patient_age),
            patient_gender=patient_gender,
            symptoms=symptoms,
            clinical_history=clinical_history,
            modality=modality,
            body_region=body_region,
            view=view,
        )
    except httpx.ReadTimeout:
        err = "❌ Request timed out — model is still processing. Wait 2 min and try again, or restart backend."
        return "", "", err, "", err
    except httpx.HTTPStatusError as exc:
        err = f"❌ Server error {exc.response.status_code}: {exc.response.text[:200]}"
        return "", "", err, "", err
    except Exception as exc:
        err = f"❌ Request failed: {exc}"
        return "", "", err, "", err

    report = data.get("report", {})
    validation = data.get("validation", {})
    report_id = data.get("report_id", "")
    confidence = validation.get("confidence_score", 0.0)
    gen_time = report.get("generation_time_seconds", 0)

    report_text = data.get("report_text", "")
    if not report_text:
        r = data.get("report", {})
        report_text = f"""RADIOLOGY REPORT
{'='*60}

EXAMINATION DETAILS:
{r.get('examination_details','—')}

CLINICAL INFORMATION:
{r.get('clinical_information','—')}

FINDINGS:
{r.get('findings','—')}

IMPRESSION:
{r.get('impression','—')}

RECOMMENDATIONS:
{r.get('recommendations','—')}

CONFIDENCE SCORE: {confidence:.2f}
{'='*60}
END OF REPORT"""

    if confidence >= 0.80:
        conf_label, conf_color = "HIGH", "#138808"
    elif confidence >= 0.60:
        conf_label, conf_color = "MODERATE", "#e67e00"
    else:
        conf_label, conf_color = "LOW", "#c0392b"

    confidence_html = f"""
<div style="text-align:center;padding:16px;background:#f0f6ff;
            border:2px solid #1a4a8a;border-radius:12px;margin-top:8px;">
  <div style="font-size:0.85rem;color:#555;margin-bottom:4px;">CONFIDENCE SCORE</div>
  <div style="font-size:1.5rem;font-weight:bold;color:{conf_color};">
    {confidence:.2f} — {conf_label}
  </div>
  <div style="font-size:0.78rem;color:#888;margin-top:6px;">
    Generated in {gen_time:.1f}s
  </div>
</div>"""

    issues = validation.get("issues", [])
    warnings = validation.get("warnings", [])
    suggestions = validation.get("suggestions", [])
    section_scores = validation.get("section_scores", {})

    val_lines = ["**Validation Results**\n"]
    val_lines.append(f"{'✅ VALID' if validation.get('is_valid') else '⚠️ REVIEW NEEDED'}")
    if section_scores:
        val_lines.append("\n**Section Scores:**")
        for sec, score in section_scores.items():
            bar = "█" * int(score * 10) + "░" * (10 - int(score * 10))
            val_lines.append(f"  {sec.replace('_',' ').title()}: {bar} {score:.2f}")
    if issues:
        val_lines.append("\n**Issues:**")
        val_lines.extend(f"  ❌ {i}" for i in issues)
    if warnings:
        val_lines.append("\n**Warnings:**")
        val_lines.extend(f"  ⚠️ {w}" for w in warnings)
    if suggestions:
        val_lines.append("\n**Suggestions:**")
        val_lines.extend(f"  💡 {s}" for s in suggestions)

    # Medical Coding Block
    coding = data.get("medical_coding", {})
    val_lines.append("\n---")
    val_lines.append("## Medical Coding (UMLS / ICD-10 / SNOMED CT)")
    if coding and coding.get("total_findings", 0) > 0:
        primary = coding.get("primary_diagnosis", {})
        val_lines.append(f"**Primary Diagnosis:** {primary.get('label','—')} | CUI: {primary.get('cui','—')}")
        val_lines.append(f"**Overall Severity:** {coding.get('overall_severity','—')}")
        icd_list = coding.get("icd10_codes", [])
        if icd_list:
            val_lines.append(f"**ICD-10 Codes:** {' | '.join(icd_list)}")
        val_lines.append(f"**All Concepts Found ({coding.get('total_findings',0)}):**")
        for mc in coding.get("medical_codes", []):
            val_lines.append(f"  * {mc['label']} | CUI: {mc['cui']} | ICD-10: {mc['icd10']} | Severity: {mc['severity']}")
    else:
        val_lines.append("No significant pathological concepts detected (normal study).")

    validation_text = "\n".join(val_lines)
    status_final = f"✅ Report generated  |  ID: {report_id[:8]}"
    return report_text, confidence_html, validation_text, report_id, status_final


def _downloads_dir() -> Path:
    """Returns a temp download folder that works on both Windows and Linux."""
    d = Path(os.environ.get("TEMP", os.environ.get("TMP", "/tmp"))) / "radiology_reports"
    d.mkdir(parents=True, exist_ok=True)
    return d


def export_pdf(report_id):
    if not report_id:
        return None
    try:
        r = httpx.get(f"{BACKEND_URL}/api/export/pdf/{report_id}", timeout=60)
        r.raise_for_status()
        path = _downloads_dir() / f"radiology_report_{report_id[:8]}.pdf"
        with open(path, "wb") as f:
            f.write(r.content)
        logger.info(f"PDF saved: {path}")
        return str(path)
    except Exception as exc:
        logger.error(f"PDF export error: {exc}")
        return None


def export_json_file(report_id):
    if not report_id:
        return None
    try:
        r = httpx.get(f"{BACKEND_URL}/api/export/json/{report_id}", timeout=30)
        r.raise_for_status()
        path = _downloads_dir() / f"radiology_report_{report_id[:8]}.json"
        with open(path, "wb") as f:
            f.write(r.content)
        logger.info(f"JSON saved: {path}")
        return str(path)
    except Exception as exc:
        logger.error(f"JSON export error: {exc}")
        return None


def fill_demo():
    return (
        45, "Male",
        "Left abdominal pain, swelling of the abdominal wall, fever",
        "Suspected abdominal wall infection.",
        "CT", "Abdomen", "Axial",
    )


def check_backend_status():
    if _check_backend():
        return "🟢 Backend online — ready to generate reports."
    return "🔴 Backend offline — run: uvicorn backend.main:app --reload"


def fetch_report_history():
    """Fetch all reports from DB for history tab."""
    try:
        r = httpx.get(f"{BACKEND_URL}/api/reports?limit=50", timeout=10)
        r.raise_for_status()
        reports = r.json().get("reports", [])
        if not reports:
            return "No reports found in database yet."
        lines = ["| Patient Name | Age | Gender | Modality | Region | Confidence | Date |",
                 "|---|---|---|---|---|---|---|"]
        for rep in reports:
            conf  = rep.get("confidence_score", 0) or 0
            lines.append(
                f"| {rep.get('patient_name','—')} "
                f"| {rep.get('age','—')} "
                f"| {rep.get('gender','—')} "
                f"| {rep.get('modality','—')} "
                f"| {rep.get('body_region','—')} "
                f"| {conf:.2f} "
                f"| {str(rep.get('created_at','—'))[:16]} |"
            )
        return "\n".join(lines)
    except Exception as e:
        return f"❌ Could not load history: {e}"


def fetch_db_stats():
    """Fetch database statistics."""
    try:
        r = httpx.get(f"{BACKEND_URL}/api/stats", timeout=10)
        r.raise_for_status()
        s = r.json()
        mod_counts = s.get("modality_counts", {})
        mod_text   = " | ".join(f"{k}: {v}" for k,v in mod_counts.items())
        return (
            f"**Total Reports:** {s.get('total_reports', 0)}  \n"
            f"**Unique Patients:** {s.get('total_patients', 0)}  \n"
            f"**Avg Confidence:** {s.get('avg_confidence', 0):.2f}  \n"
            f"**By Modality:** {mod_text or 'None yet'}"
        )
    except Exception as e:
        return f"❌ Could not load stats: {e}"


def search_patient_reports(name: str):
    """Search reports by patient name."""
    if not name.strip():
        return "Enter a patient name to search."
    try:
        r = httpx.get(
            f"{BACKEND_URL}/api/reports/patient/{name.strip()}",
            timeout=10
        )
        r.raise_for_status()
        data    = r.json()
        reports = data.get("reports", [])
        if not reports:
            return f"No reports found for patient: {name}"
        lines = [f"**{len(reports)} report(s) found for: {name}**\n",
                 "| Date | Modality | Region | Confidence | Impression |",
                 "|---|---|---|---|---|"]
        for rep in reports:
            conf = rep.get("confidence_score", 0) or 0
            imp  = str(rep.get("impression","—"))[:60] + "..."
            lines.append(
                f"| {str(rep.get('created_at','—'))[:16]} "
                f"| {rep.get('modality','—')} "
                f"| {rep.get('body_region','—')} "
                f"| {conf:.2f} "
                f"| {imp} |"
            )
        return "\n".join(lines)
    except Exception as e:
        return f"❌ Search failed: {e}"


def get_report_uuids_for_patient(name: str):
    """Returns list of (label, uuid) tuples for a patient's reports."""
    if not name.strip():
        return []
    try:
        r = httpx.get(
            f"{BACKEND_URL}/api/reports/patient/{name.strip()}", timeout=10
        )
        r.raise_for_status()
        reports = r.json().get("reports", [])
        choices = []
        for rep in reports:
            label = (
                f"{str(rep.get('created_at',''))[:16]}  |  "
                f"{rep.get('modality','—')}  {rep.get('body_region','—')}  |  "
                f"Conf: {rep.get('confidence_score',0):.2f}"
            )
            choices.append((label, rep.get("report_uuid", "")))
        return choices
    except Exception:
        return []


def load_patient_report_choices(name: str):
    """Populate the dropdown with patient's reports."""
    choices = get_report_uuids_for_patient(name)
    if not choices:
        return gr.update(choices=[], value=None,
                         label="No reports found for this patient")
    return gr.update(choices=choices, value=choices[0][1],
                     label=f"{len(choices)} report(s) found — select one to download")


def download_report_pdf(report_uuid: str):
    """Download PDF for a specific report UUID from DB."""
    if not report_uuid:
        return None
    try:
        r = httpx.get(
            f"{BACKEND_URL}/api/export/pdf/{report_uuid}", timeout=60
        )
        r.raise_for_status()
        path = _downloads_dir() / f"report_{report_uuid[:8]}.pdf"
        with open(path, "wb") as f:
            f.write(r.content)
        return str(path)
    except Exception as exc:
        logger.error(f"PDF download failed: {exc}")
        return None


def download_report_json(report_uuid: str):
    """Download JSON for a specific report UUID from DB."""
    if not report_uuid:
        return None
    try:
        r = httpx.get(
            f"{BACKEND_URL}/api/export/json/{report_uuid}", timeout=30
        )
        r.raise_for_status()
        path = _downloads_dir() / f"report_{report_uuid[:8]}.json"
        with open(path, "wb") as f:
            f.write(r.content)
        return str(path)
    except Exception as exc:
        logger.error(f"JSON download failed: {exc}")
        return None


def preview_selected_report(report_uuid: str):
    """Show full report text for selected report UUID."""
    if not report_uuid:
        return "Select a report to preview."
    try:
        r = httpx.get(
            f"{BACKEND_URL}/api/report/{report_uuid}", timeout=10
        )
        r.raise_for_status()
        data   = r.json()
        rep    = data.get("report", {})
        pat    = data.get("patient", {})
        img    = data.get("imaging", {})
        valid  = data.get("validation", {})
        conf   = valid.get("confidence_score", rep.get("confidence_score", 0)) or 0
        name   = data.get("patient_name", pat.get("name", "Unknown"))
        return f"""RADIOLOGY REPORT
{'='*60}
Patient Name  : {name}
Age / Gender  : {pat.get('age','—')} / {pat.get('gender','—')}
Study         : {img.get('modality','—')} {img.get('body_region','—')} {img.get('view','') or ''}
Date          : {str(data.get('timestamp',''))[:16]}
{'='*60}

EXAMINATION DETAILS:
{rep.get('examination_details','—')}

CLINICAL INFORMATION:
{rep.get('clinical_information','—')}

FINDINGS:
{rep.get('findings','—')}

IMPRESSION:
{rep.get('impression','—')}

RECOMMENDATIONS:
{rep.get('recommendations','—')}

CONFIDENCE SCORE: {conf:.2f}
{'='*60}
END OF REPORT"""
    except Exception as exc:
        return f"❌ Could not load report: {exc}"


def build_interface() -> gr.Blocks:
    with gr.Blocks(title="Radiology AI — Report Generation System") as demo:

        report_id_state = gr.State("")

        gr.Markdown(
            "# 🏥 Radiology AI — Report Generation System\n"
            "### AI-Based Structured Radiology Report Generation with Clinical Validation & Patient Database"
        )

        backend_status = gr.Markdown("🔴 Checking backend…")

        with gr.Tabs():

            # ── TAB 1: Generate Report ─────────────────────────────────
            with gr.Tab("🚀 Generate Report"):
                gr.Markdown(
                    "> ⏳ **Note:** First report generation takes **30–60 seconds on GPU**. "
                    "Please click Generate Report and wait — do not click again."
                )

                with gr.Row(equal_height=False):

                    # LEFT: Inputs
                    with gr.Column(scale=1, min_width=340):

                        gr.Markdown("### 📤 Step 1: Upload Radiology Image")
                        image_input = gr.File(
                            label="Upload Image (.dcm / .png / .jpg)",
                            file_types=[".dcm", ".png", ".jpg", ".jpeg"],
                            type="filepath",
                        )
                        image_preview = gr.Image(
                            label="Preview",
                            interactive=False,
                            height=220,
                        )

                        gr.Markdown("### 👤 Step 2: Patient Information")
                        patient_name = gr.Textbox(
                            label="Patient Name",
                            placeholder="e.g. John Smith",
                            lines=1,
                        )
                        with gr.Row():
                            patient_age = gr.Number(
                                label="Age (years)", minimum=0, maximum=130, step=1, value=50
                            )
                            patient_gender = gr.Dropdown(
                                label="Gender",
                                choices=["Male", "Female", "Other"],
                                value="Male",
                            )
                        symptoms = gr.Textbox(
                            label="Presenting Symptoms",
                            placeholder="e.g. abdominal pain, fever, swelling",
                            lines=2,
                        )
                        clinical_history = gr.Textbox(
                            label="Clinical History",
                            placeholder="e.g. suspected infection, prior surgeries",
                            lines=2,
                        )

                gr.Markdown("### 🔬 Step 3: Imaging Information")
                with gr.Row():
                    modality = gr.Dropdown(
                        label="Modality",
                        choices=[
                            "X-ray", "CT", "MRI", "Ultrasound", "PET",
                            "Mammography", "Fluoroscopy", "Angiography",
                            "Nuclear Medicine",
                        ],
                        value="X-ray",
                    )
                    body_region = gr.Dropdown(
                        label="Body Region",
                        choices=[
                            "Chest", "Brain", "Abdomen", "Pelvis",
                            "Spine", "Knee", "Hip", "Shoulder", "Neck",
                            "Cardiac", "Breast", "Liver", "Kidney",
                            "Thyroid", "Hand", "Wrist", "Foot", "Ankle",
                            "Skull", "Prostate", "Obstetric", "Vascular",
                        ],
                        value="Chest",
                    )
                view = gr.Dropdown(
                    label="View / Projection",
                    choices=[
                        "PA", "AP", "Lateral", "Axial", "Coronal",
                        "Sagittal", "Oblique", "Left Lateral", "Right Lateral",
                        "Supine", "Prone", "Decubitus", "Cephalocaudal (CC)",
                        "Mediolateral Oblique (MLO)", "T1", "T2", "FLAIR",
                        "DWI", "STIR", "",
                    ],
                    value="PA",
                )

                with gr.Row():
                    demo_btn = gr.Button("📋 Fill Demo", variant="secondary", size="sm")
                    generate_btn = gr.Button("🚀 Generate Report", variant="primary", size="lg")

                status_label = gr.Markdown("")

                    # RIGHT: Outputs
                with gr.Column(scale=2, min_width=480):

                    gr.Markdown("### 📄 Generated Radiology Report")
                    report_output = gr.Textbox(
                        label="Structured Report",
                        lines=22,
                        interactive=False,
                        elem_id="report-box",
                    )

                    confidence_output = gr.HTML(
                        value="<div style='text-align:center;color:#aaa;padding:16px;'>"
                              "Confidence score will appear here after generation.</div>"
                    )

                    with gr.Accordion("📊 Validation Details", open=False):
                        validation_output = gr.Markdown("")

                    gr.Markdown("### 💾 Export Report")
                    with gr.Row():
                        export_pdf_btn = gr.Button("⬇️ Download PDF", variant="secondary")
                        export_json_btn = gr.Button("⬇️ Download JSON", variant="secondary")

                    pdf_file  = gr.File(label="PDF Report",  visible=False)
                    json_file = gr.File(label="JSON Report", visible=False)

            # ── TAB 2: Patient History ─────────────────────────────────
            with gr.Tab("📋 Patient History"):
                gr.Markdown("## 📋 Report History Database")

                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### 🔍 Search by Patient Name")
                        search_input  = gr.Textbox(
                            label="Patient Name", placeholder="e.g. John Smith"
                        )
                        search_btn    = gr.Button("🔍 Search", variant="primary")
                        search_output = gr.Markdown("")

                    with gr.Column(scale=1):
                        gr.Markdown("### 📊 Database Statistics")
                        stats_output = gr.Markdown("")
                        refresh_stats_btn = gr.Button("🔄 Refresh Stats", variant="secondary")

                # ── Download Previous Reports ──────────────────────────
                gr.Markdown("---")
                gr.Markdown("### ⬇️ Download Previous Reports")
                with gr.Row():
                    with gr.Column(scale=2):
                        dl_patient_input = gr.Textbox(
                            label="Enter Patient Name to Load Reports",
                            placeholder="e.g. Gopalamanikandan A",
                        )
                        dl_load_btn = gr.Button(
                            "📂 Load Reports for This Patient",
                            variant="primary"
                        )
                        report_selector = gr.Dropdown(
                            label="Select Report to Download",
                            choices=[],
                            value=None,
                            interactive=True,
                        )

                    with gr.Column(scale=2):
                        gr.Markdown("#### 👁️ Report Preview")
                        report_preview = gr.Textbox(
                            label="Selected Report",
                            lines=18,
                            interactive=False,
                            elem_id="report-box",
                        )

                with gr.Row():
                    dl_pdf_btn  = gr.Button("⬇️ Download PDF",  variant="secondary")
                    dl_json_btn = gr.Button("⬇️ Download JSON", variant="secondary")
                dl_pdf_file  = gr.File(label="PDF",  visible=False)
                dl_json_file = gr.File(label="JSON", visible=False)

                # ── All Reports Table ──────────────────────────────────
                gr.Markdown("---")
                gr.Markdown("### 📜 All Reports")
                history_btn    = gr.Button("🔄 Load All Reports", variant="secondary")
                history_output = gr.Markdown("")

        # ── Events ────────────────────────────────────────────────────
        image_input.change(fn=lambda f: f, inputs=image_input, outputs=image_preview)

        demo_btn.click(
            fn=fill_demo, inputs=[],
            outputs=[patient_age, patient_gender, symptoms, clinical_history,
                     modality, body_region, view],
        )

        generate_btn.click(
            fn=generate_report,
            inputs=[image_input, patient_name, patient_age, patient_gender,
                    symptoms, clinical_history, modality, body_region, view],
            outputs=[report_output, confidence_output, validation_output,
                     report_id_state, status_label],
        )

        export_pdf_btn.click(
            fn=export_pdf, inputs=[report_id_state], outputs=[pdf_file],
        ).then(
            fn=lambda f: gr.update(visible=f is not None),
            inputs=[pdf_file], outputs=[pdf_file],
        )

        export_json_btn.click(
            fn=export_json_file, inputs=[report_id_state], outputs=[json_file],
        ).then(
            fn=lambda f: gr.update(visible=f is not None),
            inputs=[json_file], outputs=[json_file],
        )

        # History tab events
        history_btn.click(
            fn=fetch_report_history, inputs=[], outputs=[history_output]
        )
        search_btn.click(
            fn=search_patient_reports, inputs=[search_input], outputs=[search_output]
        )
        refresh_stats_btn.click(
            fn=fetch_db_stats, inputs=[], outputs=[stats_output]
        )

        # Download previous report events
        dl_load_btn.click(
            fn=load_patient_report_choices,
            inputs=[dl_patient_input],
            outputs=[report_selector],
        )
        report_selector.change(
            fn=preview_selected_report,
            inputs=[report_selector],
            outputs=[report_preview],
        )
        dl_pdf_btn.click(
            fn=download_report_pdf,
            inputs=[report_selector],
            outputs=[dl_pdf_file],
        ).then(
            fn=lambda f: gr.update(visible=f is not None),
            inputs=[dl_pdf_file], outputs=[dl_pdf_file],
        )
        dl_json_btn.click(
            fn=download_report_json,
            inputs=[report_selector],
            outputs=[dl_json_file],
        ).then(
            fn=lambda f: gr.update(visible=f is not None),
            inputs=[dl_json_file], outputs=[dl_json_file],
        )

        demo.load(fn=check_backend_status, outputs=[backend_status])
        demo.load(fn=fetch_db_stats, outputs=[stats_output])

    return demo


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    interface = build_interface()
    interface.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("GRADIO_PORT", "7860")),
        share=False,
        show_error=True,
        css=CUSTOM_CSS,
        theme=gr.themes.Soft(primary_hue="blue", neutral_hue="slate"),
    )