src/core/report_generator.py

"""
Human+ Branded PDF Report Generator
Creates a polished, branded PDF using reportlab.
Returns PDF as bytes for Streamlit download_button.
"""

import io
from datetime import datetime
from typing import Optional

from reportlab.lib.pagesizes import A4
from reportlab.lib.units import mm
from reportlab.lib.colors import HexColor, white, black
from reportlab.lib.enums import TA_CENTER, TA_LEFT, TA_RIGHT
from reportlab.lib.styles import ParagraphStyle
from reportlab.platypus import (
    SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle,
    HRFlowable, KeepTogether,
)
from reportlab.platypus import PageBreak


# ── Brand colors ─────────────────────────────────────────────────────
C_BG        = HexColor("#070710")
C_SURFACE   = HexColor("#0d0d1a")
C_ELEVATED  = HexColor("#131325")
C_LINE      = HexColor("#1f1f38")
C_TEXT      = HexColor("#f0f0fa")
C_MUTED     = HexColor("#8b8ba7")
C_VIOLET    = HexColor("#6d28d9")
C_VIOLET_LT = HexColor("#8b5cf6")
C_CYAN      = HexColor("#22d3ee")
C_OPTIMAL   = HexColor("#10b981")
C_WARNING   = HexColor("#f59e0b")
C_DANGER    = HexColor("#ef4444")
C_NEUTRAL   = HexColor("#6b7280")

STATUS_COLOR = {
    "optimal": C_OPTIMAL,
    "warning": C_WARNING,
    "danger":  C_DANGER,
    "neutral": C_NEUTRAL,
}


# ── Styles ────────────────────────────────────────────────────────────
def _make_styles():
    return {
        "hero_title": ParagraphStyle(
            "hero_title",
            fontName="Helvetica-Bold",
            fontSize=28,
            textColor=white,
            alignment=TA_CENTER,
            spaceAfter=4,
            leading=32,
        ),
        "hero_sub": ParagraphStyle(
            "hero_sub",
            fontName="Helvetica",
            fontSize=10,
            textColor=C_MUTED,
            alignment=TA_CENTER,
            spaceAfter=2,
            leading=14,
            letterSpacing=2,
        ),
        "section_label": ParagraphStyle(
            "section_label",
            fontName="Helvetica-Bold",
            fontSize=7,
            textColor=C_MUTED,
            spaceBefore=20,
            spaceAfter=8,
            leading=10,
            letterSpacing=2,
        ),
        "section_title": ParagraphStyle(
            "section_title",
            fontName="Helvetica-Bold",
            fontSize=14,
            textColor=white,
            spaceBefore=4,
            spaceAfter=10,
            leading=18,
        ),
        "body": ParagraphStyle(
            "body",
            fontName="Helvetica",
            fontSize=9,
            textColor=C_MUTED,
            spaceAfter=4,
            leading=14,
        ),
        "body_white": ParagraphStyle(
            "body_white",
            fontName="Helvetica",
            fontSize=9,
            textColor=C_TEXT,
            spaceAfter=4,
            leading=14,
        ),
        "biomarker_name": ParagraphStyle(
            "biomarker_name",
            fontName="Helvetica-Bold",
            fontSize=9,
            textColor=white,
            leading=12,
        ),
        "biomarker_ref": ParagraphStyle(
            "biomarker_ref",
            fontName="Helvetica",
            fontSize=7,
            textColor=C_MUTED,
            leading=10,
        ),
        "biomarker_value": ParagraphStyle(
            "biomarker_value",
            fontName="Helvetica-Bold",
            fontSize=11,
            alignment=TA_RIGHT,
            leading=14,
        ),
        "caption": ParagraphStyle(
            "caption",
            fontName="Helvetica",
            fontSize=7,
            textColor=C_MUTED,
            alignment=TA_CENTER,
            leading=10,
        ),
        "quote": ParagraphStyle(
            "quote",
            fontName="Helvetica-Oblique",
            fontSize=10,
            textColor=C_TEXT,
            alignment=TA_CENTER,
            spaceBefore=8,
            spaceAfter=4,
            leading=15,
        ),
        "footer": ParagraphStyle(
            "footer",
            fontName="Helvetica",
            fontSize=7,
            textColor=C_MUTED,
            alignment=TA_CENTER,
            leading=10,
        ),
    }


# ── Page template with dark background ───────────────────────────────

def _draw_background(canvas, doc):
    """Draw dark background on every page."""
    canvas.saveState()
    canvas.setFillColor(C_BG)
    canvas.rect(0, 0, A4[0], A4[1], fill=1, stroke=0)

    # Subtle top gradient bar (violet → cyan)
    canvas.setFillColor(C_VIOLET)
    canvas.rect(0, A4[1] - 4*mm, A4[0] * 0.5, 4*mm, fill=1, stroke=0)
    canvas.setFillColor(C_CYAN)
    canvas.rect(A4[0] * 0.5, A4[1] - 4*mm, A4[0] * 0.5, 4*mm, fill=1, stroke=0)

    # Bottom bar
    canvas.setFillColor(C_LINE)
    canvas.rect(0, 0, A4[0], 1*mm, fill=1, stroke=0)

    canvas.restoreState()


def _draw_footer(canvas, doc):
    """Draw page footer."""
    canvas.saveState()
    canvas.setFont("Helvetica", 7)
    canvas.setFillColor(C_MUTED)
    canvas.drawCentredString(
        A4[0] / 2,
        12*mm,
        f"Human+ Bali  ·  From Injured → Optimized  ·  @bennywhitethatsright  ·  Halaman {doc.page}"
    )
    canvas.restoreState()


def _on_page(canvas, doc):
    _draw_background(canvas, doc)
    _draw_footer(canvas, doc)


# ── Main generator ────────────────────────────────────────────────────

def generate_pdf_report(
    biomarkers: list[dict],
    protocol: dict,
    patient_name: str = "",
    generated_at: Optional[str] = None,
) -> bytes:
    """
    Generate branded PDF report.

    Args:
        biomarkers: list of enriched biomarker dicts
        protocol: dict with keys summary, iv_drip, supplements, lifestyle, retest
        patient_name: optional patient name for personalization
        generated_at: optional datetime string, defaults to now

    Returns:
        PDF as bytes — ready for st.download_button
    """
    buf = io.BytesIO()

    doc = SimpleDocTemplate(
        buf,
        pagesize=A4,
        leftMargin=20*mm,
        rightMargin=20*mm,
        topMargin=18*mm,
        bottomMargin=22*mm,
        title="Human+ Lab-to-Protocol Report",
        author="Human+ Bali",
        subject="Personalized Longevity Protocol",
    )

    styles = _make_styles()
    story  = []
    W      = A4[0] - 40*mm  # usable width

    if generated_at is None:
        generated_at = datetime.now().strftime("%d %B %Y, %H:%M")

    # ── Cover block ──────────────────────────────────────────────────
    story += _build_cover(styles, W, patient_name, generated_at)

    # ── Biomarker summary ────────────────────────────────────────────
    story += _build_biomarker_section(styles, W, biomarkers)

    # ── Protocol sections ────────────────────────────────────────────
    story += _build_protocol_section(
        styles, W,
        icon="IV DRIP",
        title="Rekomendasi IV Drip",
        color=C_VIOLET_LT,
        content=protocol.get("iv_drip", ""),
    )
    story += _build_protocol_section(
        styles, W,
        icon="SUPLEMEN",
        title="Suplemen Harian",
        color=C_CYAN,
        content=protocol.get("supplements", ""),
    )
    story += _build_protocol_section(
        styles, W,
        icon="LIFESTYLE",
        title="Lifestyle & Recovery Protocol",
        color=C_OPTIMAL,
        content=protocol.get("lifestyle", ""),
    )
    story += _build_protocol_section(
        styles, W,
        icon="RETEST",
        title="Jadwal Retest & Next Step",
        color=C_WARNING,
        content=protocol.get("retest", ""),
    )

    # ── Closing quote & disclaimer ───────────────────────────────────
    story += _build_footer_block(styles, W)

    doc.build(story, onFirstPage=_on_page, onLaterPages=_on_page)
    return buf.getvalue()


# ── Section builders ─────────────────────────────────────────────────

def _build_cover(styles, W, patient_name: str, generated_at: str) -> list:
    s = styles
    items = []

    # Spacer for top bar
    items.append(Spacer(1, 8*mm))

    # Logo / brand name
    items.append(Paragraph("Human+", s["hero_title"]))
    items.append(Paragraph("LAB-TO-PROTOCOL REPORT", s["hero_sub"]))
    items.append(Spacer(1, 2*mm))

    # Divider
    items.append(HRFlowable(
        width=W, thickness=1,
        color=C_LINE, spaceAfter=6*mm, spaceBefore=4*mm,
    ))

    # Meta info row
    meta_rows = []
    if patient_name:
        meta_rows.append(["Pasien", patient_name])
    meta_rows.append(["Dibuat", generated_at])
    meta_rows.append(["Klinik", "Human+ Bali"])
    meta_rows.append(["Founder", "Benjamin White @bennywhitethatsright"])

    meta_data = [[
        Paragraph(k, ParagraphStyle("ml", fontName="Helvetica-Bold", fontSize=8,
                                     textColor=C_MUTED, leading=12)),
        Paragraph(v, ParagraphStyle("mv", fontName="Helvetica", fontSize=8,
                                     textColor=C_TEXT, leading=12)),
    ] for k, v in meta_rows]

    meta_table = Table(meta_data, colWidths=[35*mm, W - 35*mm])
    meta_table.setStyle(TableStyle([
        ("BACKGROUND",  (0, 0), (-1, -1), C_SURFACE),
        ("ROWBACKGROUNDS", (0, 0), (-1, -1), [C_SURFACE, C_ELEVATED]),
        ("TOPPADDING",  (0, 0), (-1, -1), 5),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 5),
        ("LEFTPADDING",  (0, 0), (-1, -1), 8),
        ("RIGHTPADDING", (0, 0), (-1, -1), 8),
        ("ROUNDEDCORNERS", [6]),
    ]))
    items.append(meta_table)
    items.append(Spacer(1, 6*mm))

    # Protocol summary
    summary = styles  # shorthand
    if True:  # placeholder — summary handled in protocol section
        pass

    return items


def _build_biomarker_section(styles, W, biomarkers: list[dict]) -> list:
    if not biomarkers:
        return []

    items = []

    # Section label
    items.append(Paragraph("HASIL BIOMARKER", styles["section_label"]))
    items.append(Spacer(1, 2*mm))

    # Summary counts
    optimal = sum(1 for b in biomarkers if b["status"] == "optimal")
    warning  = sum(1 for b in biomarkers if b["status"] == "warning")
    danger   = sum(1 for b in biomarkers if b["status"] == "danger")
    total    = len(biomarkers)

    summary_data = [[
        _stat_cell(str(danger),  "PERLU PERHATIAN", C_DANGER),
        _stat_cell(str(warning), "SUB-OPTIMAL",     C_WARNING),
        _stat_cell(str(optimal), "OPTIMAL",         C_OPTIMAL),
        _stat_cell(str(total),   "TOTAL MARKER",    C_MUTED),
    ]]
    summary_table = Table(summary_data, colWidths=[W/4]*4)
    summary_table.setStyle(TableStyle([
        ("BACKGROUND",   (0, 0), (-1, -1), C_SURFACE),
        ("TOPPADDING",   (0, 0), (-1, -1), 8),
        ("BOTTOMPADDING",(0, 0), (-1, -1), 8),
        ("LEFTPADDING",  (0, 0), (-1, -1), 4),
        ("RIGHTPADDING", (0, 0), (-1, -1), 4),
        ("LINEAFTER",    (0, 0), (2, 0),   0.5, C_LINE),
    ]))
    items.append(summary_table)
    items.append(Spacer(1, 4*mm))

    # Biomarker rows — sorted danger → warning → optimal
    priority = {"danger": 0, "warning": 1, "optimal": 2, "neutral": 3}
    sorted_bm = sorted(biomarkers, key=lambda b: priority.get(b["status"], 3))

    # Table header
    header = [
        Paragraph("BIOMARKER", ParagraphStyle("th", fontName="Helvetica-Bold", fontSize=7,
                                               textColor=C_MUTED, leading=10)),
        Paragraph("NILAI", ParagraphStyle("th", fontName="Helvetica-Bold", fontSize=7,
                                           textColor=C_MUTED, alignment=TA_RIGHT, leading=10)),
        Paragraph("TARGET HUMAN+", ParagraphStyle("th", fontName="Helvetica-Bold", fontSize=7,
                                                   textColor=C_MUTED, alignment=TA_RIGHT, leading=10)),
        Paragraph("STATUS", ParagraphStyle("th", fontName="Helvetica-Bold", fontSize=7,
                                            textColor=C_MUTED, alignment=TA_CENTER, leading=10)),
    ]

    bm_rows = [header]
    row_colors = [C_ELEVATED]

    for b in sorted_bm:
        sc = STATUS_COLOR.get(b["status"], C_NEUTRAL)
        status_label = {
            "optimal": "OPTIMAL",
            "warning": "SUB-OPTIMAL",
            "danger":  "KRITIS",
            "neutral": "-",
        }.get(b["status"], "-")

        row = [
            Paragraph(b["name"], styles["biomarker_name"]),
            Paragraph(
                f'<font color="#{sc.hexval()[1:]}">{b["value"]} {b.get("unit","")}</font>',
                ParagraphStyle("bv", fontName="Helvetica-Bold", fontSize=10,
                               alignment=TA_RIGHT, textColor=sc, leading=13)
            ),
            Paragraph(b.get("reference", "—"),
                      ParagraphStyle("br", fontName="Helvetica", fontSize=8,
                                     textColor=C_MUTED, alignment=TA_RIGHT, leading=11)),
            Paragraph(
                f'<font color="#{sc.hexval()[1:]}">{status_label}</font>',
                ParagraphStyle("bs", fontName="Helvetica-Bold", fontSize=7,
                               alignment=TA_CENTER, textColor=sc, leading=10)
            ),
        ]
        bm_rows.append(row)
        row_colors.append(C_SURFACE)

    col_w = [W * 0.38, W * 0.22, W * 0.25, W * 0.15]
    bm_table = Table(bm_rows, colWidths=col_w, repeatRows=1)

    bm_style = [
        ("BACKGROUND",     (0, 0), (-1, 0),  C_ELEVATED),
        ("TOPPADDING",     (0, 0), (-1, -1), 7),
        ("BOTTOMPADDING",  (0, 0), (-1, -1), 7),
        ("LEFTPADDING",    (0, 0), (-1, -1), 6),
        ("RIGHTPADDING",   (0, 0), (-1, -1), 6),
        ("LINEBELOW",      (0, 0), (-1, -1), 0.4, C_LINE),
        ("ROWBACKGROUNDS", (0, 1), (-1, -1), [C_SURFACE, C_ELEVATED]),
        ("VALIGN",         (0, 0), (-1, -1), "MIDDLE"),
    ]

    # Color left border per status
    for i, b in enumerate(sorted_bm, 1):
        sc = STATUS_COLOR.get(b["status"], C_NEUTRAL)
        bm_style.append(("LINESBEFORE", (0, i), (0, i), 3, sc))

    bm_table.setStyle(TableStyle(bm_style))
    items.append(bm_table)
    items.append(Spacer(1, 6*mm))

    return items


def _build_protocol_section(
    styles, W,
    icon: str,
    title: str,
    color: HexColor,
    content: str,
) -> list:
    if not content or not content.strip():
        return []

    items = []

    # Section header card
    header_data = [[
        Paragraph(icon, ParagraphStyle(
            "icon", fontName="Helvetica-Bold", fontSize=7,
            textColor=color, leading=10, letterSpacing=1.5,
        )),
        Paragraph(title, ParagraphStyle(
            "st", fontName="Helvetica-Bold", fontSize=12,
            textColor=white, leading=15,
        )),
    ]]
    header_table = Table(header_data, colWidths=[25*mm, W - 25*mm])
    header_table.setStyle(TableStyle([
        ("BACKGROUND",    (0, 0), (-1, -1), C_ELEVATED),
        ("TOPPADDING",    (0, 0), (-1, -1), 9),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 9),
        ("LEFTPADDING",   (0, 0), (-1, -1), 10),
        ("RIGHTPADDING",  (0, 0), (-1, -1), 10),
        ("LINESBEFORE",   (0, 0), (0, 0),   3, color),
        ("VALIGN",        (0, 0), (-1, -1), "MIDDLE"),
    ]))
    items.append(KeepTogether([header_table]))

    # Content — strip markdown and render as styled paragraphs
    content_lines = _parse_markdown_to_paragraphs(content, styles, W)
    items += content_lines
    items.append(Spacer(1, 4*mm))
    items.append(HRFlowable(width=W, thickness=0.5, color=C_LINE, spaceAfter=4*mm))

    return items


def _build_footer_block(styles, W) -> list:
    items = []
    items.append(Spacer(1, 4*mm))

    items.append(Paragraph(
        "Tubuh Anda bukan sekadar 'normal' — Anda layak mendapatkan yang optimal.",
        styles["quote"],
    ))
    items.append(Paragraph(
        "— Benjamin White, Founder Human+ Bali",
        ParagraphStyle("attr", fontName="Helvetica", fontSize=8,
                       textColor=C_MUTED, alignment=TA_CENTER, leading=11),
    ))
    items.append(Spacer(1, 6*mm))

    # Gradient-style divider (simulated with two colored cells)
    div_data = [[" ", " "]]
    div_table = Table(div_data, colWidths=[W/2, W/2])
    div_table.setStyle(TableStyle([
        ("BACKGROUND",    (0, 0), (0, 0), C_VIOLET),
        ("BACKGROUND",    (1, 0), (1, 0), C_CYAN),
        ("TOPPADDING",    (0, 0), (-1, -1), 1.5),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 1.5),
    ]))
    items.append(div_table)
    items.append(Spacer(1, 4*mm))

    disclaimer = (
        "Disclaimer: Laporan ini bukan pengganti konsultasi medis. "
        "Final protocol selalu direview oleh practitioner Human+ sebelum diberikan. "
        "Untuk konsultasi personal, hubungi @bennywhitethatsright di Instagram."
    )
    items.append(Paragraph(disclaimer, styles["caption"]))

    return items


# ── Helpers ───────────────────────────────────────────────────────────

def _stat_cell(value: str, label: str, color: HexColor):
    """Build a stat count cell for the summary table."""
    return [
        Paragraph(value, ParagraphStyle(
            "sv", fontName="Helvetica-Bold", fontSize=20,
            textColor=color, alignment=TA_CENTER, leading=24,
        )),
        Paragraph(label, ParagraphStyle(
            "sl", fontName="Helvetica-Bold", fontSize=6,
            textColor=C_MUTED, alignment=TA_CENTER, leading=8, letterSpacing=1.2,
        )),
    ]


def _parse_markdown_to_paragraphs(text: str, styles: dict, W: float) -> list:
    """
    Simple markdown → ReportLab paragraph converter.
    Handles: **bold**, # headers, - bullets, tables (basic), plain text.
    """
    import re

    items = []
    lines = text.strip().split("\n")

    for line in lines:
        stripped = line.strip()
        if not stripped:
            items.append(Spacer(1, 2*mm))
            continue

        # Header lines
        if stripped.startswith("### "):
            text_content = stripped[4:]
            items.append(Paragraph(
                _md_inline(text_content),
                ParagraphStyle("h3", fontName="Helvetica-Bold", fontSize=10,
                               textColor=C_TEXT, spaceBefore=6, spaceAfter=3, leading=14),
            ))
        elif stripped.startswith("## "):
            text_content = stripped[3:]
            items.append(Paragraph(
                _md_inline(text_content),
                ParagraphStyle("h2", fontName="Helvetica-Bold", fontSize=11,
                               textColor=white, spaceBefore=8, spaceAfter=4, leading=15),
            ))
        elif stripped.startswith("**") and stripped.endswith("**") and stripped.count("**") == 2:
            # Standalone bold line — treat as sub-header
            inner = stripped.strip("*")
            items.append(Paragraph(
                inner,
                ParagraphStyle("bold_line", fontName="Helvetica-Bold", fontSize=9,
                               textColor=C_TEXT, spaceBefore=4, spaceAfter=2, leading=13),
            ))
        elif stripped.startswith("- ") or stripped.startswith("* "):
            # Bullet
            content = stripped[2:]
            items.append(Paragraph(
                f"• {_md_inline(content)}",
                ParagraphStyle("bullet", fontName="Helvetica", fontSize=9,
                               textColor=C_MUTED, leftIndent=8, spaceAfter=2,
                               leading=13),
            ))
        elif re.match(r"^\|.+\|$", stripped):
            # Skip markdown table rows — complex tables handled separately
            pass
        elif stripped.startswith("|---") or stripped == "":
            pass
        else:
            items.append(Paragraph(
                _md_inline(stripped),
                styles["body"],
            ))

    return items


def _md_inline(text: str) -> str:
    """Convert inline markdown (**bold**, *italic*) to ReportLab XML."""
    import re
    # Bold
    text = re.sub(r"\*\*(.+?)\*\*", r"<b>\1</b>", text)
    # Italic
    text = re.sub(r"\*(.+?)\*", r"<i>\1</i>", text)
    # Inline code
    text = re.sub(r"`(.+?)`", r"<font name='Courier'>\1</font>", text)
    return text