src/export.py

from __future__ import annotations

import io
import tempfile
from datetime import datetime
from typing import Any

from PIL import Image as PILImage
from reportlab.lib.colors import HexColor, black, white
from reportlab.lib.enums import TA_CENTER, TA_JUSTIFY, TA_LEFT, TA_RIGHT
from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet
from reportlab.lib.units import cm
from reportlab.platypus import (
    HRFlowable,
    Image as RLImage,
    KeepTogether,
    PageBreak,
    Paragraph,
    SimpleDocTemplate,
    Spacer,
    Table,
    TableStyle,
)

# ---------------------------------------------------------------------------
# Kalbe brand colours
# ---------------------------------------------------------------------------
_C_DARK   = HexColor("#1A5C38")
_C_GREEN  = HexColor("#2E7D32")
_C_LIME   = HexColor("#6DB33F")
_C_LTGRN  = HexColor("#E8F5E9")
_C_BDGRN  = HexColor("#C8E6C9")
_C_DGRAY  = HexColor("#424242")
_C_MGRAY  = HexColor("#757575")
_C_LGRAY  = HexColor("#F5F5F5")

# ---------------------------------------------------------------------------
# Page geometry
# ---------------------------------------------------------------------------
PAGE_W, PAGE_H = A4
MARGIN    = 2.0 * cm
USABLE_W  = PAGE_W - 2 * MARGIN   # ≈ 17 cm in points

# ---------------------------------------------------------------------------
# Styles (built lazily so reportlab is only needed on import of this module)
# ---------------------------------------------------------------------------

def _build_styles() -> dict[str, ParagraphStyle]:
    base = getSampleStyleSheet()
    return {
        "title": ParagraphStyle(
            "rl_title",
            parent=base["Normal"],
            fontSize=20,
            textColor=_C_DARK,
            fontName="Helvetica-Bold",
            spaceAfter=2,
        ),
        "subtitle": ParagraphStyle(
            "rl_subtitle",
            parent=base["Normal"],
            fontSize=11,
            textColor=_C_GREEN,
            fontName="Helvetica",
            spaceAfter=4,
        ),
        "section": ParagraphStyle(
            "rl_section",
            parent=base["Normal"],
            fontSize=10,
            textColor=_C_DARK,
            fontName="Helvetica-Bold",
            spaceBefore=4,
            spaceAfter=3,
        ),
        "body": ParagraphStyle(
            "rl_body",
            parent=base["Normal"],
            fontSize=9,
            leading=13,
            textColor=_C_DGRAY,
            fontName="Helvetica",
        ),
        "caption": ParagraphStyle(
            "rl_caption",
            parent=base["Normal"],
            fontSize=7.5,
            textColor=_C_MGRAY,
            fontName="Helvetica-Oblique",
            alignment=TA_CENTER,
            spaceBefore=2,
        ),
        "report": ParagraphStyle(
            "rl_report",
            parent=base["Normal"],
            fontSize=9.5,
            leading=15,
            textColor=_C_DGRAY,
            fontName="Helvetica",
            alignment=TA_JUSTIFY,
            spaceAfter=6,
        ),
        "report_head": ParagraphStyle(
            "rl_report_head",
            parent=base["Normal"],
            fontSize=9.5,
            fontName="Helvetica-Bold",
            textColor=_C_DARK,
            spaceBefore=6,
            spaceAfter=2,
        ),
        "th": ParagraphStyle(
            "rl_th",
            parent=base["Normal"],
            fontSize=8,
            fontName="Helvetica-Bold",
            textColor=white,
        ),
        "td": ParagraphStyle(
            "rl_td",
            parent=base["Normal"],
            fontSize=8,
            fontName="Helvetica",
            textColor=_C_DGRAY,
        ),
        "td_bold": ParagraphStyle(
            "rl_td_bold",
            parent=base["Normal"],
            fontSize=8,
            fontName="Helvetica-Bold",
            textColor=_C_DARK,
        ),
        "info_key": ParagraphStyle(
            "rl_info_key",
            parent=base["Normal"],
            fontSize=8.5,
            fontName="Helvetica-Bold",
            textColor=_C_DARK,
        ),
        "info_val": ParagraphStyle(
            "rl_info_val",
            parent=base["Normal"],
            fontSize=8.5,
            fontName="Helvetica",
            textColor=_C_DGRAY,
        ),
    }


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _pil_to_rl(img: PILImage.Image, max_w: float, max_h: float) -> RLImage:
    buf = io.BytesIO()
    img.convert("RGB").save(buf, format="PNG")
    buf.seek(0)
    orig_w, orig_h = img.size
    scale = min(max_w / orig_w, max_h / orig_h)
    return RLImage(buf, width=orig_w * scale, height=orig_h * scale)


def _section_block(label: str, styles: dict) -> list:
    return [
        Spacer(1, 0.35 * cm),
        Paragraph(label.upper(), styles["section"]),
        HRFlowable(width="100%", thickness=1.5, color=_C_DARK, spaceAfter=0.25 * cm),
    ]


def _table_style(n_header_rows: int = 1) -> TableStyle:
    cmds = [
        # Header rows
        ("BACKGROUND",    (0, 0), (-1, n_header_rows - 1), _C_DARK),
        ("TEXTCOLOR",     (0, 0), (-1, n_header_rows - 1), white),
        ("LINEBELOW",     (0, n_header_rows - 1), (-1, n_header_rows - 1), 1.5, _C_DARK),
        # Data rows alternating
        ("ROWBACKGROUNDS",(0, n_header_rows), (-1, -1), [white, _C_LTGRN]),
        # Grid
        ("GRID",          (0, 0), (-1, -1), 0.4, _C_BDGRN),
        # Padding
        ("LEFTPADDING",   (0, 0), (-1, -1), 6),
        ("RIGHTPADDING",  (0, 0), (-1, -1), 6),
        ("TOPPADDING",    (0, 0), (-1, -1), 4),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 4),
        ("VALIGN",        (0, 0), (-1, -1), "TOP"),
    ]
    return TableStyle(cmds)


def _image_pair_table(
    left_img: PILImage.Image,
    right_img: PILImage.Image,
    left_cap: str,
    right_cap: str,
    styles: dict,
    col_w: float,
    max_h: float,
) -> Table:
    left_rl  = _pil_to_rl(left_img,  col_w, max_h)
    right_rl = _pil_to_rl(right_img, col_w, max_h)
    tbl = Table(
        [
            [left_rl, right_rl],
            [Paragraph(left_cap,  styles["caption"]),
             Paragraph(right_cap, styles["caption"])],
        ],
        colWidths=[col_w, col_w],
        hAlign="CENTER",
    )
    tbl.setStyle(TableStyle([
        ("ALIGN",         (0, 0), (-1, -1), "CENTER"),
        ("VALIGN",        (0, 0), (-1, -1), "MIDDLE"),
        ("LEFTPADDING",   (0, 0), (-1, -1), 3),
        ("RIGHTPADDING",  (0, 0), (-1, -1), 3),
        ("TOPPADDING",    (0, 0), (-1, -1), 2),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 2),
    ]))
    return tbl


# ---------------------------------------------------------------------------
# Page header / footer callback (runs on every page)
# ---------------------------------------------------------------------------

def _page_callback(canvas, doc):
    canvas.saveState()

    # ── Top green banner ──────────────────────────────────────────────────────
    banner_h = 1.15 * cm
    canvas.setFillColor(_C_DARK)
    canvas.rect(0, PAGE_H - banner_h, PAGE_W, banner_h, fill=1, stroke=0)

    canvas.setFillColor(white)
    canvas.setFont("Helvetica-Bold", 8.5)
    canvas.drawString(MARGIN, PAGE_H - 0.78 * cm, "THORAX REPORT GENERATION  ·  Kalbe Digital Lab")

    canvas.setFont("Helvetica", 7.5)
    canvas.drawRightString(PAGE_W - MARGIN, PAGE_H - 0.78 * cm,
                           datetime.now().strftime("%d %B %Y"))

    # ── Bottom footer ─────────────────────────────────────────────────────────
    canvas.setStrokeColor(_C_DARK)
    canvas.setLineWidth(0.5)
    canvas.line(MARGIN, 1.1 * cm, PAGE_W - MARGIN, 1.1 * cm)

    canvas.setFillColor(_C_MGRAY)
    canvas.setFont("Helvetica", 7)
    canvas.drawString(MARGIN, 0.65 * cm, "Kalbe Digital Lab  —  Confidential  —  For clinical use only")
    canvas.drawRightString(PAGE_W - MARGIN, 0.65 * cm, f"Page {doc.page}")

    canvas.restoreState()


# ---------------------------------------------------------------------------
# Public: build_pdf
# ---------------------------------------------------------------------------

def build_pdf(
    original_images: list[PILImage.Image],
    annotated_images: list[PILImage.Image],
    seg_images: list[PILImage.Image] | None,
    detection: Any,
    masks: list | None,
    report_text: str | None,
    structured_report: Any | None = None,
    case_id: str | None = None,
) -> bytes:
    """Assemble the PDF and return it as raw bytes."""

    styles = _build_styles()
    buf = io.BytesIO()

    doc = SimpleDocTemplate(
        buf,
        pagesize=A4,
        leftMargin=MARGIN,
        rightMargin=MARGIN,
        topMargin=1.7 * cm,     # below the top banner
        bottomMargin=1.8 * cm,  # above the footer
    )

    story: list = []

    # ── Cover block ───────────────────────────────────────────────────────────
    story.append(Spacer(1, 0.3 * cm))
    story.append(Paragraph("Thorax Report Generation", styles["title"]))
    story.append(Paragraph("Kalbe Digital Lab", styles["subtitle"]))
    story.append(HRFlowable(width="100%", thickness=1, color=_C_GREEN, spaceAfter=0.3 * cm))

    # ── Case info table ───────────────────────────────────────────────────────
    n_views    = len(original_images)
    n_findings = len(detection.findings) if detection else 0
    seg_status = "Completed" if (masks and any(m.status == "success" for m in masks)) else "Not run"
    rpt_status = "Completed" if report_text else "Not run"

    info_rows = [
        [Paragraph("Date",              styles["info_key"]),
         Paragraph(datetime.now().strftime("%d %B %Y"), styles["info_val"])],
        [Paragraph("Case ID",           styles["info_key"]),
         Paragraph(case_id or "—",      styles["info_val"])],
        [Paragraph("Views analysed",    styles["info_key"]),
         Paragraph(str(n_views),        styles["info_val"])],
        [Paragraph("Findings detected", styles["info_key"]),
         Paragraph(str(n_findings),     styles["info_val"])],
        [Paragraph("Stage 2 · Segmentation", styles["info_key"]),
         Paragraph(seg_status,          styles["info_val"])],
        [Paragraph("Stage 3 · Report",  styles["info_key"]),
         Paragraph(rpt_status,          styles["info_val"])],
    ]

    KEY_W = 5.5 * cm
    info_tbl = Table(info_rows, colWidths=[KEY_W, USABLE_W - KEY_W])
    info_tbl.setStyle(TableStyle([
        ("ROWBACKGROUNDS", (0, 0), (-1, -1), [_C_LTGRN, white]),
        ("GRID",           (0, 0), (-1, -1), 0.4, _C_BDGRN),
        ("LEFTPADDING",    (0, 0), (-1, -1), 8),
        ("RIGHTPADDING",   (0, 0), (-1, -1), 8),
        ("TOPPADDING",     (0, 0), (-1, -1), 4),
        ("BOTTOMPADDING",  (0, 0), (-1, -1), 4),
    ]))
    story.append(info_tbl)

    # =========================================================================
    # STAGE 1 – Detection & Localization
    # =========================================================================
    story.extend(_section_block("Stage 1 — Detection & Localization", styles))

    IMG_COL_W = (USABLE_W - 0.4 * cm) / 2
    IMG_MAX_H = 7.0 * cm

    for i, (orig, ann) in enumerate(zip(original_images, annotated_images)):
        view_tag = "PA" if i == 0 else "Lateral"
        tbl = _image_pair_table(
            orig, ann,
            f"View {i + 1} ({view_tag}) · Original",
            f"View {i + 1} ({view_tag}) · Annotated — detected findings",
            styles, IMG_COL_W, IMG_MAX_H,
        )
        story.append(tbl)
        story.append(Spacer(1, 0.25 * cm))

    # Findings table
    story.append(Spacer(1, 0.15 * cm))

    if detection and detection.findings:
        _CRT = {"positive": "Positive", "probable": "Probable", "questionable": "Questionable"}
        _STS = {
            "localized":               "Localized ✓",
            "abstained":               "Abstained",
            "rejected_by_quality_gate":"Rejected",
            "parser_error":            "Error",
        }

        hdr = [
            Paragraph("#",          styles["th"]),
            Paragraph("Finding",    styles["th"]),
            Paragraph("Anatomical Location", styles["th"]),
            Paragraph("Certainty",  styles["th"]),
            Paragraph("View Localizations",  styles["th"]),
        ]
        rows: list[list] = [hdr]

        for idx, f in enumerate(detection.findings, 1):
            view_lines = []
            for v in f.localizations:
                label = _STS.get(v.status, v.status)
                if v.boxes:
                    bbox = v.boxes[0].box_2d
                    label += f"  [{bbox[0]}, {bbox[1]}, {bbox[2]}, {bbox[3]}]"
                view_lines.append(f"View {v.image_index}: {label}")

            cert_text = _CRT.get(f.certainty, f.certainty.title())
            rows.append([
                Paragraph(str(idx),            styles["td_bold"]),
                Paragraph(f.finding,           styles["td"]),
                Paragraph(f.anatomical_location, styles["td"]),
                Paragraph(cert_text,           styles["td"]),
                Paragraph("<br/>".join(view_lines) if view_lines else "—", styles["td"]),
            ])

        # Column widths summing to USABLE_W (~17cm)
        f_cols = [0.8 * cm, 4.5 * cm, 4.0 * cm, 2.5 * cm, USABLE_W - 11.8 * cm]
        f_tbl = Table(rows, colWidths=f_cols, repeatRows=1)
        f_tbl.setStyle(_table_style())
        story.append(KeepTogether([
            Paragraph("Detected Findings", styles["report_head"]),
            Spacer(1, 0.1 * cm),
            f_tbl,
        ]))
    else:
        story.append(Paragraph("No positive findings detected.", styles["body"]))

    # =========================================================================
    # STAGE 2 – Segmentation
    # =========================================================================
    successful_masks = [m for m in masks if m.status == "success"] if masks else []

    if seg_images and successful_masks:
        story.append(PageBreak())
        story.extend(_section_block("Stage 2 — Segmentation", styles))

        SEG_COL_W = (USABLE_W - 0.4 * cm) / 2
        SEG_MAX_H = 9.0 * cm

        if len(seg_images) == 1:
            seg_rl = _pil_to_rl(seg_images[0], USABLE_W * 0.65, SEG_MAX_H * 1.2)
            story.append(Table(
                [[seg_rl], [Paragraph("View 1 · Polygon segmentation overlay", styles["caption"])]],
                colWidths=[USABLE_W],
                hAlign="CENTER",
            ))
        else:
            seg_rls = [_pil_to_rl(img, SEG_COL_W, SEG_MAX_H) for img in seg_images[:2]]
            caps    = [Paragraph(f"View {i+1} · Polygon segmentation overlay", styles["caption"])
                       for i in range(len(seg_rls))]
            # pad to 2 columns
            while len(seg_rls) < 2:
                seg_rls.append(Spacer(0, 0))
                caps.append(Spacer(0, 0))
            seg_tbl = Table(
                [seg_rls, caps],
                colWidths=[SEG_COL_W, SEG_COL_W],
                hAlign="CENTER",
            )
            seg_tbl.setStyle(TableStyle([
                ("ALIGN",        (0, 0), (-1, -1), "CENTER"),
                ("VALIGN",       (0, 0), (-1, -1), "MIDDLE"),
                ("LEFTPADDING",  (0, 0), (-1, -1), 3),
                ("RIGHTPADDING", (0, 0), (-1, -1), 3),
            ]))
            story.append(seg_tbl)

        story.append(Spacer(1, 0.4 * cm))

        # Segmentation summary table
        s_hdr = [
            Paragraph("Finding",        styles["th"]),
            Paragraph("View",           styles["th"]),
            Paragraph("Polygon Points", styles["th"]),
            Paragraph("Result",         styles["th"]),
        ]
        s_rows: list[list] = [s_hdr]

        for m in masks:  # type: ignore[union-attr]
            poly_pts = len(m.polygon) if (m.polygon and m.status == "success") else None
            result   = "Segmented" if m.status == "success" else f"Failed — {m.error or 'unknown'}"
            s_rows.append([
                Paragraph(m.finding_label,   styles["td"]),
                Paragraph(f"View {m.image_index}", styles["td"]),
                Paragraph(str(poly_pts) if poly_pts is not None else "—", styles["td"]),
                Paragraph(result,            styles["td"]),
            ])

        # 17cm total: 6.5 + 2.0 + 3.5 + 5.0
        s_cols = [6.5 * cm, 2.0 * cm, 3.5 * cm, USABLE_W - 12.0 * cm]
        s_tbl = Table(s_rows, colWidths=s_cols, repeatRows=1)
        s_tbl.setStyle(_table_style())
        story.append(KeepTogether([
            Paragraph("Segmentation Results", styles["report_head"]),
            Spacer(1, 0.1 * cm),
            s_tbl,
        ]))

    # =========================================================================
    # STAGE 3 – Radiology Report
    # =========================================================================
    has_report = (structured_report is not None) or (report_text and report_text.strip())
    if has_report:
        story.append(PageBreak())
        story.extend(_section_block("Stage 3 — Radiology Report", styles))
        story.append(Spacer(1, 0.15 * cm))

        if structured_report is not None:
            r = structured_report

            def _rpt_row(label: str, value: str) -> list:
                return [
                    Paragraph(label, styles["report_head"]),
                    Paragraph(value or "—", styles["report"]),
                ]

            def _rpt_list_rows(label: str, items: list[str]) -> list[list]:
                rows = [[Paragraph(label, styles["report_head"]), Spacer(1, 0)]]
                for item in items:
                    rows.append([Spacer(1, 0), Paragraph(f"• {item}", styles["report"])])
                return rows

            KEY_W2 = 4.5 * cm
            VAL_W  = USABLE_W - KEY_W2

            # Study type & summary block
            if r.study_type or r.summary:
                top_rows = []
                if r.study_type:
                    top_rows.append(_rpt_row("Study Type", r.study_type))
                if r.summary:
                    top_rows.append(_rpt_row("Summary", r.summary))
                top_tbl = Table(top_rows, colWidths=[KEY_W2, VAL_W])
                top_tbl.setStyle(TableStyle([
                    ("BACKGROUND",    (0, 0), (0, -1), _C_LTGRN),
                    ("GRID",          (0, 0), (-1, -1), 0.4, _C_BDGRN),
                    ("LEFTPADDING",   (0, 0), (-1, -1), 6),
                    ("RIGHTPADDING",  (0, 0), (-1, -1), 6),
                    ("TOPPADDING",    (0, 0), (-1, -1), 4),
                    ("BOTTOMPADDING", (0, 0), (-1, -1), 4),
                    ("VALIGN",        (0, 0), (-1, -1), "TOP"),
                ]))
                story.append(top_tbl)
                story.append(Spacer(1, 0.3 * cm))

            # Main findings list
            if r.main_findings:
                story.append(Paragraph("Main Findings", styles["report_head"]))
                story.append(Spacer(1, 0.05 * cm))
                for item in r.main_findings:
                    story.append(Paragraph(f"• {item}", styles["report"]))
                story.append(Spacer(1, 0.25 * cm))

            # Detailed findings list
            if r.detail_findings:
                story.append(Paragraph("Detailed Findings", styles["report_head"]))
                story.append(Spacer(1, 0.05 * cm))
                for item in r.detail_findings:
                    story.append(Paragraph(f"• {item}", styles["report"]))
                story.append(Spacer(1, 0.25 * cm))

            # Impression, Recommendations, Additional info
            bottom_rows = []
            if r.impression:
                bottom_rows.append(_rpt_row("Impression", r.impression))
            if r.recommendations:
                bottom_rows.append(_rpt_row("Recommendations", r.recommendations))
            ai = r.additional_informations or ""
            if ai.strip().lower() not in ("none", "n/a", "-", ""):
                bottom_rows.append(_rpt_row("Additional Information", ai))

            if bottom_rows:
                bot_tbl = Table(bottom_rows, colWidths=[KEY_W2, VAL_W])
                bot_tbl.setStyle(TableStyle([
                    ("BACKGROUND",    (0, 0), (0, -1), _C_LTGRN),
                    ("GRID",          (0, 0), (-1, -1), 0.4, _C_BDGRN),
                    ("LEFTPADDING",   (0, 0), (-1, -1), 6),
                    ("RIGHTPADDING",  (0, 0), (-1, -1), 6),
                    ("TOPPADDING",    (0, 0), (-1, -1), 4),
                    ("BOTTOMPADDING", (0, 0), (-1, -1), 4),
                    ("VALIGN",        (0, 0), (-1, -1), "TOP"),
                ]))
                story.append(bot_tbl)

        elif report_text:
            # Fallback: plain-text report (unstructured)
            for raw_line in report_text.split("\n"):
                line = raw_line.strip()
                if not line:
                    story.append(Spacer(1, 0.18 * cm))
                    continue
                if (line.endswith(":") and len(line) < 55) or (line.isupper() and len(line) < 55):
                    story.append(Paragraph(line, styles["report_head"]))
                else:
                    story.append(Paragraph(line, styles["report"]))

    # ── Build ─────────────────────────────────────────────────────────────────
    doc.build(story, onFirstPage=_page_callback, onLaterPages=_page_callback)
    return buf.getvalue()


# ---------------------------------------------------------------------------
# Public: export_to_tempfile
# ---------------------------------------------------------------------------

def export_to_tempfile(
    original_images: list[PILImage.Image],
    annotated_images: list[PILImage.Image],
    seg_images: list[PILImage.Image] | None,
    detection: Any,
    masks: list | None,
    report_text: str | None,
    structured_report: Any | None = None,
    case_id: str | None = None,
) -> str:
    """Write the PDF to a temp file and return its path (Gradio-compatible)."""
    pdf_bytes = build_pdf(
        original_images=original_images,
        annotated_images=annotated_images,
        seg_images=seg_images or None,
        detection=detection,
        masks=masks or None,
        report_text=report_text or None,
        structured_report=structured_report,
        case_id=case_id,
    )
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    with tempfile.NamedTemporaryFile(
        suffix=".pdf", delete=False, prefix=f"thorax_report_{ts}_"
    ) as f:
        f.write(pdf_bytes)
        return f.name