"""
Gradio web interface for the Thyroid Nodule Analysis Pipeline.

Supports two analysis modes:
  Single image  - upload one ultrasound image and run the full pipeline.
  Patient mode  - upload multiple images of the same patient; results are
                  shown per-image and then aggregated using WMV + TI-RADS
                  feature aggregation.

File validation:
  Only images are accepted (.png, .jpg, .jpeg).
  Invalid formats produce a clear error message before any model inference.

Partial-failure handling:
  Each pipeline step reports its status; if an upstream step fails, the
  interface still shows results from all steps that succeeded.
"""

import os

# ZeroGPU
_ON_HF = os.environ.get("SPACE_ID") is not None
if _ON_HF:
    import spaces # type: ignore

import time
import numpy as np
from PIL import Image
import gradio as gr


from config import (
    DEVICE, ACCEPTED_EXTENSIONS,
    DISPLAY_FEATURES, FEATURE_LABELS, FEATURE_CATEGORY,
    TIRADS_RISK, TIRADS_FNAB,
)
from pipeline import ThyroidPipeline, ImageResult, PatientResult

# Initialise pipeline (models loaded once at startup)
print(f"[App] Starting Thyroid Pipeline — device: {DEVICE}")
pipeline = ThyroidPipeline()
print("[App] Pipeline ready.")

# Input validation

ACCEPTED_EXT_STR = ", ".join(sorted(ACCEPTED_EXTENSIONS))


def validate_images(files: list) -> tuple[bool, str]:
    """
    Check that all uploaded files are valid image formats.

    Parameters
    files : list - Gradio file objects (have a .name attribute with the path)

    Returns
    (is_valid, error_message)
    """
    if not files:
        return False, "No files uploaded. Please upload at least one ultrasound image."

    for f in files:
        path = f.name if hasattr(f, "name") else str(f)
        ext  = os.path.splitext(path)[-1].lower()
        if ext not in ACCEPTED_EXTENSIONS:
            return False, (
                f"Unsupported file format: '{os.path.basename(path)}' ({ext}). "
                f"Accepted formats: {ACCEPTED_EXT_STR}. "
                "Please upload a valid ultrasound image."
            )
        # Try to open with PIL to catch corrupted files early
        try:
            img = Image.open(path)
            img.verify()
        except Exception as e:
            return False, (
                f"File '{os.path.basename(path)}' could not be read as an image. "
                f"It may be corrupted or have an incorrect extension. ({e})"
            )

    return True, ""


# HTML result builders

def _error_html(title: str, message: str) -> str:
    return f"""
<div style="font-family:sans-serif; padding:12px; border-left:4px solid #e74c3c;
            background:#fdf2f2; border-radius:4px; margin-bottom:8px">
  <h3 style="color:#c0392b; margin:0 0 6px 0">{title}</h3>
  <p style="margin:0; color:#555; line-height:1.5">{message}</p>
</div>"""


def _warning_html(title: str, message: str) -> str:
    return f"""
<div style="font-family:sans-serif; padding:12px; border-left:4px solid #e67e22;
            background:#fef9f0; border-radius:4px; margin-bottom:8px">
  <h3 style="color:#d35400; margin:0 0 6px 0">{title}</h3>
  <p style="margin:0; color:#555; line-height:1.5">{message}</p>
</div>"""


def _status_section_html(result: ImageResult) -> str:
    """Generate a warning/error banner if the pipeline did not complete fully."""
    if result.status == "OK":
        return ""
    if result.status == "NO_DETECTION":
        return _error_html("No Nodule Detected", result.status_message)
    if result.status == "SEG_FAILED":
        return _warning_html("Segmentation Failed", result.status_message)
    if result.status == "CLASSIF_FAILED":
        return _warning_html("Classification Failed", result.status_message)
    return _warning_html("Pipeline Warning", result.status_message)


def _image_result_html(result: ImageResult, idx: int | None = None) -> str:
    """Build the full HTML summary for a single ImageResult."""
    header = f"Image {idx}" if idx is not None else "Result"
    status_block = _status_section_html(result)

    if result.status == "NO_DETECTION":
        return status_block

    # Detection section (always present if boxes were found)
    detection_html = ""
    if result.n_boxes > 0:
        detection_html = f"""
  <h3 style="margin:16px 0 6px 0">Detection (YOLOv8s)</h3>
  <ul style="margin:0 0 8px 0">
    <li>Nodules detected: <b>{result.n_boxes}</b></li>
    <li>Best box confidence: <b>{result.best_conf:.1%}</b></li>
  </ul>"""

    if result.status in ("SEG_FAILED", "CLASSIF_FAILED"):
        seg_html = ""
        if result.status == "CLASSIF_FAILED" and result.nodule_area_px is not None:
            seg_html = f"""
    <h3 style="margin:16px 0 6px 0">Segmentation (UNet++ · ResNet50 encoder · SCSE)</h3>
    <ul style="margin:0 0 8px 0">
        <li>Nodule area: <b>{result.nodule_area_px} px²</b></li>
    </ul>"""
        return f"""
    <div style="font-family:sans-serif; padding:8px">
    <h3 style="margin:0 0 8px 0">{header}</h3>
    {status_block}
    {detection_html}
    {seg_html}
    <p style="color:#888; font-size:13px">
        Elapsed: {result.elapsed_s:.2f}s
    </p>
    </div>"""

    # Full result (status == "OK")
    # NOTE: ACR TI-RADS score is intentionally NOT shown here.
    # TI-RADS is a patient-level score computed after aggregating radiomic
    # features across ALL images. It is displayed only in the Patient Result
    # section above. Showing a per-image TR score would be clinically incorrect.
    label_color = "#c0392b" if result.label == "MALIGNANT" else "#27ae60"

    return f"""
<div style="font-family:sans-serif; padding:8px">
  <h2 style="color:{label_color}; margin:0 0 4px 0">{result.label}</h2>
  <hr style="margin:8px 0"/>
  {detection_html}
  <h3 style="margin:16px 0 6px 0">Segmentation (UNet++ · ResNet50 encoder · SCSE)</h3>
  <ul style="margin:0 0 8px 0">
    <li>Nodule area: <b>{result.nodule_area_px} px²</b></li>
  </ul>
  <h3 style="margin:16px 0 6px 0">Classification (ResNet50)</h3>
  <ul style="margin:0 0 8px 0">
    <li>Malignancy probability: <b>{result.prob_malign:.1%}</b></li>
    <li>Prediction: <b style="color:{label_color}">{result.label}</b></li>
  </ul>
  <p style="color:#888; font-size:13px; margin-top:12px">
    Processing time: {result.elapsed_s:.2f}s
    &nbsp;|&nbsp; TI-RADS score is shown in the Patient Result above (aggregated across all images).
  </p>
</div>"""


def _patient_summary_html(patient: PatientResult) -> str:
    """Build the aggregated patient-level summary HTML block."""

    has_classification = patient.final_label in ("MALIGNANT", "BENIGN")

    if patient.final_label == "MALIGNANT":
        label_color = "#c0392b"
        result_title = "Patient Result: MALIGNANT"
    elif patient.final_label == "BENIGN":
        label_color = "#27ae60"
        result_title = "Patient Result: BENIGN"
    else:
        label_color = "#6b7280"
        result_title = "Patient Result: Classification unavailable"

    prob_text = (
        f"{patient.aggregated_prob:.1%}"
        if has_classification and patient.aggregated_prob is not None
        else "N/A"
    )

    tr_class = patient.tirads_class

    tirads_block = f"""
    <tr><td style="padding:5px 12px; font-weight:600">Class</td>
        <td style="padding:5px 12px"><b>TR{tr_class}</b></td></tr>
    <tr style="background:#f0f4ff">
        <td style="padding:5px 12px">Risk interpretation</td>
        <td style="padding:5px 12px">{patient.tirads_risk}</td></tr>
    <tr><td style="padding:5px 12px">FNA biopsy recommendation</td>
        <td style="padding:5px 12px">{patient.fnab_recommendation}</td></tr>
""" if tr_class is not None else """
    <tr><td colspan="2" style="padding:5px 12px; color:#888">
        Not enough data for TI-RADS scoring</td></tr>
"""

    return f"""
<div style="font-family:sans-serif; padding:12px; background:#f8f9fa;
            border-radius:8px; border:2px solid {label_color}">
  <h2 style="color:{label_color}; margin:0 0 8px 0">
    {result_title}
  </h2>
  <p style="margin:0 0 12px 0; color:#555">
    <b>Aggregated malignancy probability:</b> {prob_text}
    &nbsp;|&nbsp; {patient.n_images_processed} / {patient.n_images_submitted} images classified
    &nbsp;|&nbsp; {patient.elapsed_total_s:.2f}s total
  </p>
  <hr style="margin:8px 0"/>
  <h3 style="margin:12px 0 6px 0">ACR TI-RADS Score (patient-level, aggregated features)</h3>
  <table style="border-collapse:collapse; width:100%; font-size:14px">
    {tirads_block}
  </table>
</div>"""


def _features_html(features: dict) -> str:
    """
    Build the top-10 radiomic features HTML table.

    Parameters
    features : dict  - feature values (aggregated patient-level vector)
    """
    rows = ""
    for i, feat in enumerate(DISPLAY_FEATURES):
        val     = features.get(feat, float("nan"))
        val_str = f"{val:.4f}" if not (isinstance(val, float) and np.isnan(val)) else "N/A"
        bg      = "#f5f5f5" if i % 2 == 0 else "#ffffff"
        rows += f"""
        <tr style="background:{bg}">
          <td style="padding:8px 14px; font-weight:500">{FEATURE_LABELS[feat]}</td>
          <td style="padding:8px 14px; color:#555">{FEATURE_CATEGORY[feat]}</td>
          <td style="padding:8px 14px; background:#eaf4fb;
                     font-family:monospace; text-align:center">{val_str}</td>
        </tr>"""

    return f"""
<table style="width:100%; border-collapse:collapse; font-size:14px; font-family:sans-serif">
  <thead>
    <tr style="background:#4a6fa5; color:white">
      <th style="padding:10px 14px; text-align:left">Feature</th>
      <th style="padding:10px 14px; text-align:left">TI-RADS Category</th>
      <th style="padding:10px 14px; text-align:center">Value</th>
    </tr>
  </thead>
  <tbody>{rows}</tbody>
</table>"""


# Main pipeline callback (called by Gradio)

def analyze_images(files: list) -> tuple:
    """
    Entry point called by the Gradio Analyze button.

    Parameters
    files : list of Gradio file objects

    Returns
    7-tuple matching Gradio output components:
      (patient_summary_html,
       image_gallery,         # list of PIL images for gr.Gallery
       per_image_html,        # per-image result accordion HTML
       features_html,         # radiomic features table (best image or aggregated)
       status_html,           # top-level status / error banner
       timer_html)            # processing time summary
    """
    EMPTY = (None, [], "", "", "", "")

    is_valid, err_msg = validate_images(files)
    if not is_valid:
        error_block = _error_html("Invalid Input", err_msg)
        return error_block, [], "", "", error_block, ""

    pil_images = []
    for f in files:
        path = f.name if hasattr(f, "name") else str(f)
        try:
            img = Image.open(path).convert("RGB")
            pil_images.append(img)
        except Exception as e:
            error_block = _error_html(
                "Image Load Error",
                f"Could not load '{os.path.basename(path)}': {e}"
            )
            return error_block, [], "", "", error_block, ""

    # Run patient pipeline
    patient: PatientResult = pipeline.run_patient(pil_images)

    # Build gallery images
    # Gallery contains: for each image - original, detection, segmentation, crop, gradcam (None slots are skipped)
    gallery_images = []
    for idx, res in enumerate(patient.per_image_results, start=1):
        label = f"Img {idx}"
        if res.image_original    is not None:
            gallery_images.append((res.image_original,    f"{label} - Original"))
        if res.image_detection   is not None:
            gallery_images.append((res.image_detection,   f"{label} - Detection (YOLO)"))
        if res.image_segmentation is not None:
            gallery_images.append((res.image_segmentation, f"{label} - Segmentation (UNet++)"))
        if res.image_crop        is not None:
            gallery_images.append((res.image_crop,         f"{label} - Nodule ROI (ResNet50)"))
        if res.image_gradcam     is not None:
            gallery_images.append((res.image_gradcam,      f"{label} - Grad-CAM"))

    # Per-image results HTML
    per_img_parts = []
    for idx, res in enumerate(patient.per_image_results, start=1):
        per_img_parts.append(
            f"<h3 style='margin:16px 0 4px 0'>Image {idx}</h3>"
            + _image_result_html(res)
        )
    per_image_html = "".join(per_img_parts) if per_img_parts else ""

    # Patient summary HTML
    all_statuses = [res.status for res in patient.per_image_results]

    if patient.n_images_processed == 0 and not patient.aggregated_features:
        has_classif_failed = any(s == "CLASSIF_FAILED" for s in all_statuses)
        has_seg_failed     = any(s == "SEG_FAILED" for s in all_statuses)
        has_any_detection  = any(s != "NO_DETECTION" for s in all_statuses)

        if has_classif_failed:
            patient_summary = _error_html(
                "Classification and TI-RADS Scoring Unavailable",
                "The nodule was detected and segmented, but the segmented mask is too "
                "small to produce a valid ResNet50 classification or TI-RADS score. "
                "This typically occurs when the nodule occupies less than 0.5% of the "
                "image area. Detection and segmentation outputs are shown in the gallery below."
            )
        elif has_seg_failed:
            patient_summary = _error_html(
                "Segmentation Failed",
                "A nodule was detected by YOLO but UNet++ could not produce a valid "
                "segmentation mask. Classification and TI-RADS scoring were not performed. "
                "The detection result is shown in the gallery below."
            )
        elif has_any_detection:
            patient_summary = _error_html(
                "Pipeline Incomplete",
                "Detection succeeded but the pipeline could not complete classification "
                "or TI-RADS scoring. See the per-image details below for the specific error."
            )
        else:
            patient_summary = _error_html(
                "No Nodule Detected",
                "None of the uploaded images produced a valid detection. "
                "Please verify that the images are thyroid ultrasound scans."
            )
    else:
        patient_summary = _patient_summary_html(patient)

    # Radiomic features table
    # Show the AGGREGATED feature vector that was passed to the Random Forest,
    # not raw features from a single image. This is the patient-level vector
    # produced by PatientAggregator (MAX/MIN/MEAN across all images), which
    # is exactly what determined the TI-RADS score shown above.
    if patient.aggregated_features:
        feats_html = _features_html(patient.aggregated_features)
    else:
        feats_html = ""

    # Timer banner
    timer_html = (
        f"<p style='font-family:sans-serif; color:#888; font-size:13px; "
        f"margin:4px 0'>Total processing time: "
        f"<b>{patient.elapsed_total_s:.2f}s</b> "
        f"| Device: <b>{str(DEVICE).upper()}</b> "
        f"| Images classified: <b>{patient.n_images_processed}"
        f"/{patient.n_images_submitted}</b></p>"
    )

    return (
        patient_summary,
        gallery_images,
        per_image_html,
        feats_html,
        "",
        timer_html,
    )


# Gradio interface

css = """
.main-title { text-align: center !important; }
h1          { text-align: center !important; }
footer      { display: none !important; }

.thumbnails button[aria-label="Share"],
.gallery button[aria-label="Share"],
.icon-button-wrapper button[aria-label="Share"],
button[aria-label="Share"] { display: none !important; }

#file-upload-box .wrap.svelte-1vmd51o > .svelte-8prmba:not(svg):not(button) {
    display: none !important;
}

.gallery caption,
.gallery .caption,
[class*="caption"],
[class*="thumbnail"] span,
.thumbnails span,
.gallery-item span,
figure figcaption {
    font-family: ui-sans-serif, system-ui, sans-serif !important;
    font-style: normal !important;
    font-weight: 400 !important;
    font-size: 13px !important;
    border-top: none !important;
    padding-top: 4px !important;
}

#gallery-no-scroll,
#gallery-no-scroll .grid-wrap,
#gallery-no-scroll .grid-container {
    max-height: none !important;
    height: auto !important;
    overflow: visible !important;
}
"""

with gr.Blocks(title="Thyroid Nodule Analysis Pipeline") as demo:

    # Header
    gr.Markdown("""
    <div class="main-title">

    # Thyroid Nodule Analysis Pipeline
    **Automated detection · segmentation · classification · ACR TI-RADS scoring**

    </div>
    """)

    # Input row
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(
                "Upload one or more ultrasound images of the same patient, "
                "then click **Analyze**."
            )
            input_files = gr.File(
                label       = "Upload ultrasound image(s) - single image or multiple views of the same patient",
                file_count  = "multiple",
                file_types  = [".png", ".jpg", ".jpeg"],
                elem_id     = "file-upload-box",
            )
            run_btn = gr.Button("Analyze", variant="primary", size="lg")
            timer_out = gr.HTML(label="")

        with gr.Column(scale=1):
            status_out  = gr.HTML(label="")
            patient_out = gr.HTML(label="Patient Result")

    gr.Markdown("---")

    # Image gallery
    gr.Markdown("### Visual Pipeline Outputs")
    gr.HTML(
        "<p style='font-family:sans-serif; font-size:14px; margin:0 0 8px 0; color:#374151'>"
        "Gallery shows: Original &rarr; Detection &rarr; Segmentation &rarr; "
        "Nodule ROI &rarr; Grad-CAM, for each uploaded image.</p>"
    )
    gallery_out = gr.Gallery(
        label       = "Pipeline outputs",
        show_label  = True,
        columns     = 5,
        height      = "auto",
        object_fit  = "contain",
        elem_id     = "gallery-no-scroll",
    )

    gr.Markdown("---")

    gr.Markdown("#### Per-image pipeline details")
    per_image_out = gr.HTML(label="")

    gr.Markdown("---")

    # TI-RADS features table
    gr.Markdown("### TI-RADS Radiomic Features")
    features_out = gr.HTML(label="")

    # Button wiring
    run_btn.click(
        fn      = analyze_images,
        inputs  = [input_files],
        outputs = [
            patient_out,
            gallery_out,
            per_image_out,
            features_out,
            status_out,
            timer_out,
        ],
    )

# Launch
if __name__ == "__main__":
    demo.launch(
        theme       = gr.themes.Soft(),
        css         = css,
        server_name = "0.0.0.0",
        share       = False,
    )