import os
import tempfile
import time
import nibabel as nib
import numpy as np
from PIL import Image
import gradio as gr
import modal

try:
    Segmenter = modal.Cls.from_name("ct-summary-backend", "Segmenter")
    segmenter_instance = Segmenter()
except Exception as e:
    print(f"[LOCAL] Failed to connect to Modal backend: {e}")
    segmenter_instance = None

# Removed global ping that blocked HF Space startup


def slice_3d_volumetric_scan(nifti_path):
    try:
        img = nib.load(nifti_path)
        data = img.get_fdata()
        z_mid = data.shape[2] // 2
        slice_data = data[:, :, z_mid]
        slice_data = np.rot90(slice_data)
        data_min, data_max = np.min(slice_data), np.max(slice_data)
        if data_max - data_min > 0:
            normalized = 255.0 * (slice_data - data_min) / (data_max - data_min)
        else:
            normalized = np.zeros_like(slice_data)
        img_uint8 = normalized.astype(np.uint8)
        tmp_img = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        Image.fromarray(img_uint8).save(tmp_img.name)
        return tmp_img.name
    except Exception as e:
        print(f"Visualization Error: {e}")
        return None


def _validate_scan_local(nifti_path):
    """Minimal validation: 3D only, not a mask. Let TotalSegmentator handle the rest."""
    try:
        img = nib.load(nifti_path)
        data = img.get_fdata()
    except Exception as e:
        return False, f"Not supported file or wrong CT scan. Could not read volume: {e}"

    if len(data.shape) != 3:
        return False, f"Not supported file or wrong CT scan. Expected 3D volume, got {len(data.shape)}D shape {data.shape}."

    unique_count = len(np.unique(data))
    print(f"[LOCAL] Validation: shape={data.shape}, unique_values={unique_count}, min={np.min(data):.1f}, max={np.max(data):.1f}")

    if unique_count < 50:
        return False, "Not supported file or wrong CT scan. Uploaded file appears to be a segmentation mask (too few unique values)."

    return True, None


SECTION_ORDER = [
    "Solid Organs", "Gastrointestinal", "Thoracic", "Genitourinary", "Other Structures"
]


def build_preview_html(findings: dict) -> str:
    if findings.get("error"):
        return (
            '<div class="preview-alert preview-alert-error">'
            f'<strong>Processing issue:</strong> {findings["error"]}'
            '</div>'
        )

    alerts = findings.get("alerts", [])
    sections = findings.get("sections", {})
    total_structures = findings.get("total_structures", 0)

    if alerts:
        html = '<div class="preview-alert preview-alert-warning">'
        html += f'<div class="preview-alert-title">⚠ {len(alerts)} finding(s) outside expected range</div>'
        html += '<ul>'
        for a in alerts:
            vol_str = f" — {a['volume']:.1f} cm³" if a.get("volume") is not None else ""
            html += f'<li><strong>{a["name"]}</strong>{vol_str}<br><span class="preview-note">{a["note"]}</span></li>'
        html += '</ul></div>'
    else:
        html = (
            '<div class="preview-alert preview-alert-ok">'
            '✓ No findings outside expected range across measured structures.'
            '</div>'
        )

    html += '<div class="preview-metrics">'
    for section_name in SECTION_ORDER:
        entries = sections.get(section_name)
        if not entries:
            continue
        html += f'<div class="preview-section-title">{section_name}</div>'
        for e in entries:
            cls = "preview-metric-alert" if e["status"] == "alert" else "preview-metric"
            html += f'<div class="{cls}"><span>{e["name"]}</span><span>{e["volume"]:.1f} cm³</span></div>'
        html += '</div>'

    html += (
        f'<div class="preview-footnote">'
        f'{total_structures} structures measured. '
        f'Volumes are approximate (fast-mode segmentation) — screening only, not diagnostic.'
        f'</div>'
    )

    return html


def build_report_html(findings: dict, scan_label: str, for_pdf: bool = False) -> str:
    if findings.get("error"):
        body = f'<div class="alert-banner alert-error"><strong>Processing issue:</strong> {findings["error"]}</div>'
        return _wrap_html(body, scan_label, for_pdf)

    alerts = findings.get("alerts", [])
    sections = findings.get("sections", {})
    total_structures = findings.get("total_structures", 0)

    if alerts:
        body = '<div class="alert-banner alert-warning">'
        body += f'<div class="alert-title">⚠ {len(alerts)} finding(s) outside expected range</div>'
        body += '<ul class="alert-list">'
        for a in alerts:
            vol_str = f" ({a['volume']:.1f} cm³)" if a.get("volume") is not None else ""
            body += f'<li><span class="organ-name">{a["name"]}</span>{vol_str} — {a["note"]}</li>'
        body += '</ul></div>'
    else:
        body = '<div class="alert-banner alert-ok">'
        body += '<div class="alert-title">✓ No findings outside expected range</div>'
        body += '<p>All measured structures fall within typical adult volume ranges for the available reference set.</p>'
        body += '</div>'

    for section_name in SECTION_ORDER:
        entries = sections.get(section_name)
        if not entries:
            continue
        body += f'<div class="section-title">{section_name}</div><ul>'
        for e in entries:
            status_class = "status-alert" if e["status"] == "alert" else "status-normal"
            note_html = f'<div class="organ-note">{e["note"]}</div>' if e.get("note") else ""
            body += (
                f'<li class="{status_class}">'
                f'<span class="organ-name">{e["name"]}</span>: {e["volume"]:.1f} cm³'
                f'{note_html}</li>'
            )
        body += '</ul>'

    body += (
        f'<p class="meta-note">Total structures measured: {total_structures}. '
        f'Volumes are approximate, derived from a fast-mode segmentation pass and intended '
        f'for screening purposes only — not a substitute for radiologist review.</p>'
    )

    return _wrap_html(body, scan_label, for_pdf)


def _wrap_html(content_html: str, scan_label: str, for_pdf: bool) -> str:
    page_rule = """
        @page {
            size: A4;
            margin: 20mm 15mm 20mm 15mm;
            @bottom-right {
                content: "Page " counter(page) " of " counter(pages);
                font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif;
                font-size: 9pt;
                color: #64748b;
            }
        }
    """ if for_pdf else ""

    return f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <style>
        {page_rule}
        body {{
            font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif;
            color: #1e293b;
            margin: 0;
            padding: 0;
            line-height: 1.6;
            background-color: #ffffff;
        }}
        .header {{
            border-bottom: 2px solid #0f172a;
            padding-bottom: 12px;
            margin-bottom: 25px;
        }}
        .header h1 {{
            font-size: 22pt;
            color: #0f172a;
            margin: 0 0 6px 0;
            text-transform: uppercase;
            letter-spacing: 0.5px;
        }}
        .header .subtitle {{
            font-size: 11pt;
            color: #475569;
            margin: 0;
            font-weight: bold;
        }}
        .metadata-table {{
            width: 100%;
            margin-bottom: 25px;
            border-collapse: collapse;
            background-color: #f8fafc;
            border: 1px solid #e2e8f0;
        }}
        .metadata-table td {{
            padding: 10px 12px;
            font-size: 10pt;
            border: 1px solid #e2e8f0;
        }}
        .metadata-label {{
            font-weight: bold;
            color: #334155;
            background-color: #f1f5f9;
            width: 25%;
        }}
        .alert-banner {{
            border-radius: 4px;
            padding: 14px 16px;
            margin-bottom: 22px;
            border: 1px solid;
        }}
        .alert-warning {{
            background-color: #fef2f2;
            border-color: #fecaca;
            color: #991b1b;
        }}
        .alert-ok {{
            background-color: #f0fdf4;
            border-color: #bbf7d0;
            color: #166534;
        }}
        .alert-error {{
            background-color: #fef2f2;
            border-color: #fecaca;
            color: #991b1b;
        }}
        .alert-title {{
            font-size: 11.5pt;
            font-weight: bold;
            margin-bottom: 6px;
        }}
        .alert-list {{
            margin: 6px 0 0 0;
            padding-left: 20px;
        }}
        .alert-list li {{
            font-size: 10.5pt;
            margin-bottom: 4px;
        }}
        .section-title {{
            font-size: 12pt;
            color: #1e3a8a;
            background-color: #eff6ff;
            padding: 6px 10px;
            margin-top: 22px;
            margin-bottom: 12px;
            font-weight: bold;
            border-left: 4px solid #2563eb;
            text-transform: uppercase;
            letter-spacing: 0.5px;
            page-break-after: avoid;
        }}
        ul {{
            margin: 0 0 15px 0;
            padding-left: 20px;
        }}
        li {{
            font-size: 10.5pt;
            margin-bottom: 6px;
            page-break-inside: avoid;
        }}
        li.status-alert {{
            color: #991b1b;
        }}
        .organ-name {{
            font-weight: bold;
            color: #0f172a;
        }}
        li.status-alert .organ-name {{
            color: #991b1b;
        }}
        .organ-note {{
            font-size: 9.5pt;
            font-weight: normal;
            color: #7f1d1d;
            margin-top: 2px;
        }}
        .meta-note {{
            font-size: 9pt;
            color: #64748b;
            margin-top: 20px;
            font-style: italic;
        }}
    </style>
</head>
<body>
    <div class="header">
        <h1>Automated 3D Volumetric Report</h1>
        <div class="subtitle">Full-Body Clinical Quantification Pipeline Output</div>
    </div>

    <table class="metadata-table">
        <tr>
            <td class="metadata-label">Protocol Type</td>
            <td>{scan_label}</td>
            <td class="metadata-label">Analysis Target</td>
            <td>Full Volumetric Masking (Total Body)</td>
        </tr>
        <tr>
            <td class="metadata-label">Pipeline Engine</td>
            <td>TotalSegmentator 3D U-Net (fast mode)</td>
            <td class="metadata-label">Reporting Method</td>
            <td>Rule-Based Reference Range Analysis</td>
        </tr>
    </table>

    <div class="content-body">
        {content_html}
    </div>
</body>
</html>
"""


def run_pipeline(file_obj, progress=gr.Progress()):
    t_start = time.time()

    if file_obj is None:
        return None, '<div class="preview-alert preview-alert-error">Upload a NIfTI (.nii or .nii.gz) volume to begin.</div>', None

    scan_label = "Whole Body CT (Auto-Detected)"

    # --- Local validation ---
    progress(0.05, desc="Validating file...")
    is_valid, err_msg = _validate_scan_local(file_obj.name)
    if not is_valid:
        return None, f'<div class="preview-alert preview-alert-error"><strong>{err_msg}</strong></div>', None

    # --- Slice extraction ---
    progress(0.15, desc="Extracting preview slice...")
    slice_path = slice_3d_volumetric_scan(file_obj.name)
    if slice_path is None:
        return None, '<div class="preview-alert preview-alert-error">Failed to extract preview slice.</div>', None

    if segmenter_instance is None:
        err_html = (
            '<div class="preview-alert preview-alert-error">'
            "Could not connect to the Modal backend. Confirm the 'ct-summary-backend' app is deployed."
            '</div>'
        )
        return slice_path, err_html, None

    try:
        # --- Read file ---
        progress(0.25, desc="Reading file...")
        with open(file_obj.name, "rb") as f:
            file_bytes = f.read()

        # --- Pre-flight ping ---
        progress(0.30, desc="Connecting to backend...")
        try:
            segmenter_instance.ping.remote()
        except Exception as e:
            return slice_path, f'<div class="preview-alert preview-alert-error">Backend unreachable: {e}</div>', None

        # --- Modal remote call ---
        progress(0.35, desc="Uploading & running segmentation (~20-30s)...")
        t0 = time.time()
        findings = segmenter_instance.validate_and_report.remote(file_bytes)
        t_remote = time.time() - t0
        print(f"[frontend timing] Modal remote call: {t_remote:.1f}s")

        # --- Preview HTML ---
        progress(0.80, desc="Building report...")
        report_preview = build_preview_html(findings)

        # --- PDF generation ---
        progress(0.90, desc="Generating PDF...")
        from weasyprint import HTML
        pdf_html = build_report_html(findings, scan_label, for_pdf=True)
        pdf_dir = tempfile.mkdtemp()
        pdf_path = os.path.join(pdf_dir, "ct_report.pdf")
        HTML(string=pdf_html).write_pdf(pdf_path)

        progress(1.0, desc="Done")
        print(f"[frontend timing] TOTAL pipeline: {time.time() - t_start:.1f}s")

        return slice_path, report_preview, pdf_path

    except Exception as e:
        err_html = f'<div class="preview-alert preview-alert-error"><strong>Pipeline execution failed:</strong> {e}</div>'
        return slice_path, err_html, None


clinical_theme = gr.themes.Soft(
    primary_hue="blue",
    neutral_hue="slate",
).set(
    body_background_fill="#0f172a",
    block_background_fill="#1e293b",
    block_border_color="#334155",
    button_primary_background_fill="#2563eb",
    button_primary_text_color="#ffffff",
    body_text_color="#f1f5f9"
)

custom_css = """
.gradio-container { font-family: 'Helvetica Neue', Arial, sans-serif; }
h1, h2, h3, h4, h5, h6 { color: #ffffff !important; }

#main-heading { text-align: center; }

.full-height-image { height: 790px !important; }
.full-height-image img { height: 100% !important; object-fit: contain; }

.report-frame {
    background-color: #ffffff !important;
    border-radius: 6px;
    border: 1px solid #334155;
    min-height: 300px;
    max-height: 790px !important;
    padding: 16px;
    font-family: 'Helvetica Neue', Arial, sans-serif;
    overflow-y: auto !important;
}
.report-frame, .report-frame * {
    color: #1e293b !important;
}
.report-frame h1, .report-frame h2, .report-frame h3 { color: #0f172a !important; }

.preview-alert {
    border-radius: 4px;
    padding: 12px 14px;
    margin-bottom: 16px;
    border: 1px solid;
    font-size: 10.5pt;
}
.preview-alert-warning, .preview-alert-warning * { background-color: #fef2f2; border-color: #fecaca; color: #991b1b !important; }
.preview-alert-ok, .preview-alert-ok * { background-color: #f0fdf4; border-color: #bbf7d0; color: #166534 !important; }
.preview-alert-error, .preview-alert-error * { background-color: #fef2f2; border-color: #fecaca; color: #991b1b !important; }
.preview-alert-title { font-weight: bold; margin-bottom: 6px; }
.preview-alert ul { margin: 6px 0 0 0; padding-left: 18px; }
.preview-alert li { margin-bottom: 8px; }
.preview-note, .preview-note * { font-size: 9pt; color: #7f1d1d !important; }

.preview-section-title, .preview-section-title * {
    font-size: 10.5pt;
    font-weight: bold;
    color: #1e3a8a !important;
    background-color: #eff6ff;
    padding: 4px 8px;
    margin-top: 14px;
    margin-bottom: 6px;
    border-left: 3px solid #2563eb;
    text-transform: uppercase;
    letter-spacing: 0.5px;
}
.preview-metric, .preview-metric * {
    display: flex;
    justify-content: space-between;
    font-size: 10.5pt;
    padding: 3px 6px;
    border-bottom: 1px solid #f1f5f9;
    color: #1e293b !important;
}
.preview-metric-alert, .preview-metric-alert * {
    display: flex;
    justify-content: space-between;
    font-size: 10.5pt;
    padding: 3px 6px;
    border-bottom: 1px solid #f1f5f9;
    color: #991b1b !important;
    font-weight: bold;
    background-color: #fef2f2;
}
.preview-footnote, .preview-footnote * {
    font-size: 9pt;
    color: #64748b !important;
    font-style: italic;
    margin-top: 14px;
}
"""

PLACEHOLDER_HTML = """
<div style="padding: 40px 20px; text-align:center; color:#94a3b8; font-family: 'Helvetica Neue', Arial, sans-serif;">
    Upload a CT volume (.nii / .nii.gz) and run the analysis to see the metrics here.
</div>
"""

with gr.Blocks(theme=clinical_theme, css=custom_css, title="CT Report Generator") as demo:
    gr.Markdown("# Automated 3D Imaging Extraction & Reporting Pipeline", elem_id="main-heading")
    gr.Markdown(
        "Upload a 3D CT volume to generate a structured report with volume-based alerts.",
        elem_id="main-heading"
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. Cross-Section Visualization")
            image_output = gr.Image(
                label="Middle Z-Axis Cross-Section",
                type="filepath",
                height=790,
                elem_classes=["full-height-image"]
            )

        with gr.Column(scale=1):
            gr.Markdown("### 2. Upload & Analyze")
            file_input = gr.File(
                label="Upload 3D Volumetric Scan (.nii.gz / .nii)",
                file_types=[".gz", ".nii"]
            )

            submit_btn = gr.Button("Analyze Scan & Generate Report", variant="primary")

            gr.Markdown("#### Metrics & Alerts")
            report_output = gr.HTML(value=PLACEHOLDER_HTML, elem_classes=["report-frame"])

            pdf_download = gr.DownloadButton("Download Official PDF Report", variant="secondary")

    submit_btn.click(
        fn=run_pipeline,
        inputs=[file_input],
        outputs=[image_output, report_output, pdf_download]
    )

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