olivas/backend/app/services/report_generator.py

import io
import os
from datetime import datetime

from PIL import Image as PILImage
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import ParagraphStyle
from reportlab.lib.units import inch, mm
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.platypus import (
    Image,
    PageBreak,
    Paragraph,
    SimpleDocTemplate,
    Spacer,
    Table,
    TableStyle,
)

# ─── Colors ──────────────────────────────────────────
NAVY = colors.HexColor("#1a1a2e")
GOLD = colors.HexColor("#ffc407")
GOLD_LIGHT = colors.HexColor("#fff8e0")
LIGHT_GRAY = colors.HexColor("#f5f5f5")
MEDIUM_GRAY = colors.HexColor("#666666")
DARK_TEXT = colors.HexColor("#333333")
WHITE = colors.white
GREEN = colors.HexColor("#16a34a")
BLUE = colors.HexColor("#2563eb")
AMBER = colors.HexColor("#d97706")
PURPLE = colors.HexColor("#7c3aed")

# ─── Register Montserrat fonts ──────────────────────
FONT_DIR = os.path.join(os.path.dirname(__file__), "fonts")

pdfmetrics.registerFont(TTFont("Montserrat", os.path.join(FONT_DIR, "Montserrat-Regular.ttf")))
pdfmetrics.registerFont(TTFont("Montserrat-Bold", os.path.join(FONT_DIR, "Montserrat-Bold.ttf")))
pdfmetrics.registerFont(TTFont("Montserrat-SemiBold", os.path.join(FONT_DIR, "Montserrat-SemiBold.ttf")))

pdfmetrics.registerFontFamily(
    "Montserrat",
    normal="Montserrat",
    bold="Montserrat-Bold",
)


def _bytes_to_image(data: bytes) -> io.BytesIO:
    buf = io.BytesIO(data)
    buf.seek(0)
    return buf


def _make_styles():
    title = ParagraphStyle(
        "Title",
        fontName="Montserrat-Bold",
        fontSize=32,
        textColor=NAVY,
        spaceAfter=6,
        leading=38,
    )
    subtitle = ParagraphStyle(
        "Subtitle",
        fontName="Montserrat",
        fontSize=14,
        textColor=GOLD,
        spaceAfter=20,
        leading=18,
    )
    heading = ParagraphStyle(
        "Heading",
        fontName="Montserrat-Bold",
        fontSize=18,
        textColor=NAVY,
        spaceBefore=12,
        spaceAfter=8,
        leading=22,
    )
    subheading = ParagraphStyle(
        "Subheading",
        fontName="Montserrat-SemiBold",
        fontSize=13,
        textColor=NAVY,
        spaceBefore=8,
        spaceAfter=4,
        leading=16,
    )
    body = ParagraphStyle(
        "Body",
        fontName="Montserrat",
        fontSize=10,
        textColor=DARK_TEXT,
        spaceAfter=6,
        leading=14,
    )
    body_small = ParagraphStyle(
        "BodySmall",
        fontName="Montserrat",
        fontSize=9,
        textColor=MEDIUM_GRAY,
        spaceAfter=4,
        leading=12,
    )
    meta_label = ParagraphStyle(
        "MetaLabel",
        fontName="Montserrat-SemiBold",
        fontSize=10,
        textColor=MEDIUM_GRAY,
        spaceAfter=2,
    )
    meta_value = ParagraphStyle(
        "MetaValue",
        fontName="Montserrat-Bold",
        fontSize=14,
        textColor=NAVY,
        spaceAfter=8,
    )
    footer = ParagraphStyle(
        "Footer",
        fontName="Montserrat",
        fontSize=8,
        textColor=MEDIUM_GRAY,
        alignment=1,  # center
    )
    return {
        "title": title,
        "subtitle": subtitle,
        "heading": heading,
        "subheading": subheading,
        "body": body,
        "body_small": body_small,
        "meta_label": meta_label,
        "meta_value": meta_value,
        "footer": footer,
    }


def _insight_type_label(t: str) -> tuple[str, colors.Color]:
    return {
        "success": ("STRENGTH", GREEN),
        "warning": ("ATTENTION", AMBER),
        "info": ("INSIGHT", BLUE),
    }.get(t, ("INSIGHT", BLUE))


def _build_insight_table(insights: list[dict], styles: dict, is_ai: bool = False) -> list:
    """Build styled insight rows as ReportLab elements."""
    elements = []
    for insight in insights:
        label_text, label_color = _insight_type_label(insight["type"])
        accent = PURPLE if is_ai else label_color

        # Build a mini table for each insight card
        badge = f'<font color="#{accent.hexval()[2:]}" size="7"><b>{label_text}</b></font>'
        title = f'<font name="Montserrat-SemiBold" size="10" color="#1a1a2e">{insight["title"]}</font>'
        desc = f'<font name="Montserrat" size="9" color="#555555">{insight["description"]}</font>'

        content = Paragraph(f"{badge}<br/>{title}<br/>{desc}", styles["body"])

        t = Table(
            [[content]],
            colWidths=[170 * mm],
        )
        t.setStyle(TableStyle([
            ("BACKGROUND", (0, 0), (-1, -1), GOLD_LIGHT if is_ai else LIGHT_GRAY),
            ("LEFTPADDING", (0, 0), (-1, -1), 10),
            ("RIGHTPADDING", (0, 0), (-1, -1), 10),
            ("TOPPADDING", (0, 0), (-1, -1), 8),
            ("BOTTOMPADDING", (0, 0), (-1, -1), 8),
            ("ROUNDEDCORNERS", [4, 4, 4, 4]),
            ("LINEBEFOREDECORCOLOR", (0, 0), (0, -1), accent),
            # Left accent border
            ("LINEBEFORE", (0, 0), (0, -1), 3, accent),
        ]))
        elements.append(t)
        elements.append(Spacer(1, 4))

    return elements


def generate_report(
    analysis,
    original_image: bytes,
    heatmap_image: bytes,
    gaze_image: bytes,
    aois: list,
    rule_insights: list[dict] | None = None,
    ai_insights: list[dict] | None = None,
    ai_cost_usd: float | None = None,
) -> bytes:
    """Generate a professional PDF report for an analysis."""
    buffer = io.BytesIO()
    doc = SimpleDocTemplate(
        buffer,
        pagesize=A4,
        topMargin=20 * mm,
        bottomMargin=20 * mm,
        leftMargin=15 * mm,
        rightMargin=15 * mm,
    )

    styles = _make_styles()
    elements = []

    # ─── Cover Page ──────────────────────────────────────────
    elements.append(Spacer(1, 50))
    elements.append(Paragraph("OliVAS", styles["title"]))
    elements.append(Paragraph("Visual Attention Analysis Report", styles["subtitle"]))
    elements.append(Spacer(1, 8))

    # Gold divider line
    divider = Table([[""]], colWidths=[60 * mm], rowHeights=[2])
    divider.setStyle(TableStyle([
        ("BACKGROUND", (0, 0), (-1, -1), GOLD),
        ("LEFTPADDING", (0, 0), (-1, -1), 0),
        ("RIGHTPADDING", (0, 0), (-1, -1), 0),
        ("TOPPADDING", (0, 0), (-1, -1), 0),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 0),
    ]))
    elements.append(divider)
    elements.append(Spacer(1, 20))

    # Metadata grid
    meta_data = [
        ["ANALYSIS", "MODEL", "DATE", "IMAGE SIZE"],
        [
            analysis.name,
            analysis.model_used.replace("_", " ").title(),
            datetime.now().strftime("%B %d, %Y"),
            f"{analysis.image_width} x {analysis.image_height}",
        ],
    ]
    meta_table = Table(meta_data, colWidths=[45 * mm] * 4)
    meta_table.setStyle(TableStyle([
        ("FONTNAME", (0, 0), (-1, 0), "Montserrat-SemiBold"),
        ("FONTSIZE", (0, 0), (-1, 0), 8),
        ("TEXTCOLOR", (0, 0), (-1, 0), MEDIUM_GRAY),
        ("FONTNAME", (0, 1), (-1, 1), "Montserrat-Bold"),
        ("FONTSIZE", (0, 1), (-1, 1), 11),
        ("TEXTCOLOR", (0, 1), (-1, 1), NAVY),
        ("TOPPADDING", (0, 0), (-1, -1), 4),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 4),
        ("VALIGN", (0, 0), (-1, -1), "TOP"),
    ]))
    elements.append(meta_table)
    elements.append(Spacer(1, 10))

    # Score badge
    if analysis.overall_score is not None:
        score = analysis.overall_score
        score_color = GREEN if score >= 60 else (AMBER if score >= 30 else colors.HexColor("#ef4444"))
        score_text = f'<font name="Montserrat-SemiBold" size="9" color="#666666">ATTENTION SCORE</font><br/><font name="Montserrat-Bold" size="28" color="#{score_color.hexval()[2:]}">{score:.0f}</font><font name="Montserrat" size="12" color="#999999"> / 100</font>'
        elements.append(Paragraph(score_text, styles["body"]))
        elements.append(Spacer(1, 10))

    # Thumbnail on cover
    max_width = 160 * mm
    img = PILImage.open(io.BytesIO(original_image))
    aspect = img.width / img.height
    img_width = min(max_width, 130 * mm)
    img_height = img_width / aspect
    if img_height > 90 * mm:
        img_height = 90 * mm
        img_width = img_height * aspect
    elements.append(Image(_bytes_to_image(original_image), width=img_width, height=img_height))

    elements.append(PageBreak())

    # ─── Heatmap Page ────────────────────────────────────────
    elements.append(Paragraph("Attention Heatmap", styles["heading"]))
    elements.append(
        Paragraph(
            "Areas highlighted in warm colors (red/yellow) indicate high predicted "
            "attention during the first 3-5 seconds of viewing. Cool colors (blue) "
            "indicate lower attention probability.",
            styles["body"],
        )
    )
    elements.append(Spacer(1, 8))
    elements.append(Image(_bytes_to_image(heatmap_image), width=img_width, height=img_height))

    elements.append(PageBreak())

    # ─── Gaze Sequence Page ──────────────────────────────────
    elements.append(Paragraph("Predicted Gaze Sequence", styles["heading"]))
    elements.append(
        Paragraph(
            "Numbered points show the predicted order in which viewers will "
            "fixate on different areas of the design. Point #1 is where the eye "
            "lands first.",
            styles["body"],
        )
    )
    elements.append(Spacer(1, 8))
    elements.append(Image(_bytes_to_image(gaze_image), width=img_width, height=img_height))

    if analysis.gaze_sequence:
        elements.append(Spacer(1, 10))
        elements.append(Paragraph("Fixation Details", styles["subheading"]))
        table_data = [["#", "Position", "From Left", "From Top", "Probability"]]
        for point in analysis.gaze_sequence:
            table_data.append([
                str(point["rank"]),
                f"({point['x']}, {point['y']})",
                f"{point['x_pct']:.1f}%",
                f"{point['y_pct']:.1f}%",
                f"{point['probability']:.1%}",
            ])
        t = Table(table_data, colWidths=[25, 80, 60, 60, 70])
        t.setStyle(TableStyle([
            ("BACKGROUND", (0, 0), (-1, 0), NAVY),
            ("TEXTCOLOR", (0, 0), (-1, 0), WHITE),
            ("FONTNAME", (0, 0), (-1, 0), "Montserrat-SemiBold"),
            ("FONTNAME", (0, 1), (-1, -1), "Montserrat"),
            ("FONTSIZE", (0, 0), (-1, -1), 9),
            ("GRID", (0, 0), (-1, -1), 0.5, colors.HexColor("#dddddd")),
            ("ROWBACKGROUNDS", (0, 1), (-1, -1), [WHITE, LIGHT_GRAY]),
            ("ALIGN", (0, 0), (-1, -1), "CENTER"),
            ("TOPPADDING", (0, 0), (-1, -1), 5),
            ("BOTTOMPADDING", (0, 0), (-1, -1), 5),
        ]))
        elements.append(t)

    # ─── Hotspots Page ───────────────────────────────────────
    if analysis.hotspots:
        elements.append(PageBreak())
        elements.append(Paragraph("Attention Hotspots", styles["heading"]))
        elements.append(
            Paragraph(
                "The top regions ranked by predicted attention intensity. Higher rank "
                "means more predicted visual attention.",
                styles["body"],
            )
        )
        elements.append(Spacer(1, 8))

        hs_data = [["Rank", "Position (x, y)", "Intensity"]]
        for hs in analysis.hotspots:
            cx = hs.get("center_x", hs["x"])
            cy = hs.get("center_y", hs["y"])
            hs_data.append([
                f"#{hs['rank']}",
                f"({cx}, {cy})",
                f"{hs['intensity']:.2%}",
            ])
        t = Table(hs_data, colWidths=[50, 120, 80])
        t.setStyle(TableStyle([
            ("BACKGROUND", (0, 0), (-1, 0), NAVY),
            ("TEXTCOLOR", (0, 0), (-1, 0), WHITE),
            ("FONTNAME", (0, 0), (-1, 0), "Montserrat-SemiBold"),
            ("FONTNAME", (0, 1), (-1, -1), "Montserrat"),
            ("FONTSIZE", (0, 0), (-1, -1), 9),
            ("GRID", (0, 0), (-1, -1), 0.5, colors.HexColor("#dddddd")),
            ("ROWBACKGROUNDS", (0, 1), (-1, -1), [WHITE, LIGHT_GRAY]),
            ("ALIGN", (0, 0), (-1, -1), "CENTER"),
            ("TOPPADDING", (0, 0), (-1, -1), 5),
            ("BOTTOMPADDING", (0, 0), (-1, -1), 5),
        ]))
        elements.append(t)

    # ─── Insights Page ───────────────────────────────────────
    has_rule = rule_insights and len(rule_insights) > 0
    has_ai = ai_insights and len(ai_insights) > 0

    if has_rule or has_ai:
        elements.append(PageBreak())
        elements.append(Paragraph("Analysis Insights", styles["heading"]))

        if has_rule:
            elements.append(Paragraph("Rule-Based Insights", styles["subheading"]))
            elements.append(
                Paragraph(
                    "Automatically generated observations based on the attention metrics.",
                    styles["body_small"],
                )
            )
            elements.append(Spacer(1, 6))
            elements.extend(_build_insight_table(rule_insights, styles, is_ai=False))
            elements.append(Spacer(1, 12))

        if has_ai:
            elements.append(Paragraph("AI Design Analysis", styles["subheading"]))
            ai_note = "Generated by Claude Sonnet 4.6 — context-aware design recommendations."
            if ai_cost_usd is not None:
                ai_note += f" (Cost: ${ai_cost_usd:.4f})"
            elements.append(Paragraph(ai_note, styles["body_small"]))
            elements.append(Spacer(1, 6))
            elements.extend(_build_insight_table(ai_insights, styles, is_ai=True))

    # ─── AOI Page (if any) ───────────────────────────────────
    if aois:
        elements.append(PageBreak())
        elements.append(Paragraph("Areas of Interest", styles["heading"]))
        elements.append(
            Paragraph(
                "User-defined regions analyzed for predicted attention capture. "
                "Density > 1.0 means the region captures more attention than its "
                "size would suggest.",
                styles["body"],
            )
        )
        elements.append(Spacer(1, 8))

        aoi_data = [["Region", "Attention %", "Area %", "Density"]]
        for aoi in aois:
            aoi_data.append([
                aoi.label,
                f"{aoi.attention_pct:.1f}%",
                f"{aoi.area_pct:.1f}%",
                f"{aoi.attention_density:.2f}x",
            ])
        t = Table(aoi_data, colWidths=[100, 70, 70, 70])
        t.setStyle(TableStyle([
            ("BACKGROUND", (0, 0), (-1, 0), NAVY),
            ("TEXTCOLOR", (0, 0), (-1, 0), WHITE),
            ("FONTNAME", (0, 0), (-1, 0), "Montserrat-SemiBold"),
            ("FONTNAME", (0, 1), (-1, -1), "Montserrat"),
            ("FONTSIZE", (0, 0), (-1, -1), 9),
            ("GRID", (0, 0), (-1, -1), 0.5, colors.HexColor("#dddddd")),
            ("ROWBACKGROUNDS", (0, 1), (-1, -1), [WHITE, LIGHT_GRAY]),
            ("ALIGN", (1, 0), (-1, -1), "CENTER"),
            ("TOPPADDING", (0, 0), (-1, -1), 5),
            ("BOTTOMPADDING", (0, 0), (-1, -1), 5),
        ]))
        elements.append(t)

    # ─── Footer ──────────────────────────────────────────────
    elements.append(Spacer(1, 30))
    elements.append(
        Paragraph(
            "Generated by OliVAS — Open-Source Visual Attention Software by OLIVER",
            styles["footer"],
        )
    )

    doc.build(elements)
    return buffer.getvalue()