"""PDF export for the Praxa Behavioral Alignment explorer.

Uses ReportLab (pure Python, no Chromium). Charts stay on the Space;
the PDF is a readable briefing: executive summary, then one section per domain.
"""

from __future__ import annotations

from datetime import datetime, timezone
from io import BytesIO
from typing import Any

import pandas as pd
from reportlab.lib import colors
from reportlab.lib.enums import TA_CENTER, TA_JUSTIFY
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet
from reportlab.lib.units import inch
from reportlab.platypus import (
    PageBreak,
    Paragraph,
    SimpleDocTemplate,
    Spacer,
    Table,
    TableStyle,
)

from data import (
    ALL_CONDITIONS,
    DOMAIN_LABELS,
    DOMAIN_PREFIX,
    MATCHED_8B_KEYS,
    best_by_family,
    disagreement_summary,
    encoder_comparator_frame,
    generative_ladder_frame,
    get_metrics,
    refresh_cache,
    rfdt_encoder_disagreement,
)
from narrative import (
    EXPERIMENTS_DATASET_REPO,
    PUBLISHED_MODEL_REPOS,
    narrative_bundle,
)

NAVY = colors.HexColor("#0c1220")
# Match Space card / section chrome (dark panel blues, not bright accent).
PANEL = colors.HexColor("#12203a")
PANEL_EDGE = colors.HexColor("#1e3a5f")
INK = colors.HexColor("#0f172a")
MUTED = colors.HexColor("#334155")
RULE = colors.HexColor("#94a3b8")
ROW_ALT = colors.HexColor("#f1f5f9")
RESEARCHER = "Eidan J. Rosado, PhD"


def _fmt(value: float | None) -> str:
    if value is None or (isinstance(value, float) and pd.isna(value)):
        return "—"
    return f"{float(value):.3f}"


def _styles() -> dict[str, ParagraphStyle]:
    base = getSampleStyleSheet()
    return {
        "cover_title": ParagraphStyle(
            "cover_title",
            parent=base["Title"],
            fontName="Helvetica-Bold",
            fontSize=22,
            leading=26,
            textColor=NAVY,
            alignment=TA_CENTER,
            spaceAfter=4,
        ),
        "cover_sub": ParagraphStyle(
            "cover_sub",
            parent=base["Normal"],
            fontName="Helvetica",
            fontSize=11,
            leading=15,
            textColor=PANEL_EDGE,
            alignment=TA_CENTER,
            spaceAfter=4,
        ),
        "byline": ParagraphStyle(
            "byline",
            parent=base["Normal"],
            fontName="Helvetica",
            fontSize=10,
            leading=13,
            textColor=MUTED,
            alignment=TA_CENTER,
            spaceBefore=0,
            spaceAfter=8,
        ),
        "h1": ParagraphStyle(
            "h1",
            parent=base["Heading1"],
            fontName="Helvetica-Bold",
            fontSize=14,
            leading=18,
            textColor=NAVY,
            spaceBefore=14,
            spaceAfter=8,
        ),
        "h2": ParagraphStyle(
            "h2",
            parent=base["Heading2"],
            fontName="Helvetica-Bold",
            fontSize=12,
            leading=15,
            textColor=PANEL,
            spaceBefore=10,
            spaceAfter=6,
        ),
        "body": ParagraphStyle(
            "body",
            parent=base["Normal"],
            fontName="Helvetica",
            fontSize=9.5,
            leading=13,
            textColor=INK,
            alignment=TA_JUSTIFY,
            spaceAfter=8,
        ),
        "bullet": ParagraphStyle(
            "bullet",
            parent=base["Normal"],
            fontName="Helvetica",
            fontSize=9.5,
            leading=13,
            textColor=INK,
            leftIndent=12,
            spaceAfter=4,
        ),
        "meta": ParagraphStyle(
            "meta",
            parent=base["Normal"],
            fontName="Helvetica",
            fontSize=8,
            leading=10,
            textColor=MUTED,
            alignment=TA_CENTER,
        ),
        "footer": ParagraphStyle(
            "footer",
            parent=base["Normal"],
            fontName="Helvetica",
            fontSize=8,
            textColor=MUTED,
            alignment=TA_CENTER,
        ),
        "cell": ParagraphStyle(
            "cell",
            parent=base["Normal"],
            fontName="Helvetica",
            fontSize=8,
            leading=10,
            textColor=INK,
        ),
        "cell_head": ParagraphStyle(
            "cell_head",
            parent=base["Normal"],
            fontName="Helvetica-Bold",
            fontSize=8,
            leading=10,
            textColor=colors.white,
        ),
    }


def _escape(text: str) -> str:
    return (
        text.replace("&", "&")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
    )


def _domain_frame(domain: str) -> pd.DataFrame:
    df_raw = get_metrics(domain, matched_8b_only=False)
    if df_raw.empty:
        return df_raw
    keep = df_raw["model_key"].isin(MATCHED_8B_KEYS) | (df_raw["family"] == "encoder")
    df_raw = df_raw[keep].reset_index(drop=True)
    ladder = generative_ladder_frame(df_raw)
    narrative = pd.concat(
        [ladder, df_raw[df_raw["family"] == "encoder"]],
        ignore_index=True,
    )
    if not narrative.empty:
        narrative = narrative.drop_duplicates(
            subset=["model_key", "condition"], keep="last"
        )
    return narrative


def _bests_table(stats: dict[str, Any], styles: dict[str, ParagraphStyle]) -> Table:
    rows = [
        [
            Paragraph("Method", styles["cell_head"]),
            Paragraph("Macro-F1", styles["cell_head"]),
            Paragraph("Model", styles["cell_head"]),
            Paragraph("Condition", styles["cell_head"]),
        ]
    ]
    mapping = [
        ("Best frozen", "best_frozen"),
        ("Best LoRA", "best_lora"),
        ("Best Behavioral KV", "best_kv"),
        ("Best RFDT-LoRA", "best_rfdt"),
        ("Encoder-FT", "best_encoder"),
    ]
    for label, key in mapping:
        hit = stats.get(key)
        if not hit:
            rows.append(
                [
                    Paragraph(label, styles["cell"]),
                    Paragraph("—", styles["cell"]),
                    Paragraph("—", styles["cell"]),
                    Paragraph("—", styles["cell"]),
                ]
            )
            continue
        rows.append(
            [
                Paragraph(label, styles["cell"]),
                Paragraph(_fmt(hit.get("value")), styles["cell"]),
                Paragraph(_escape(str(hit.get("model") or "—")), styles["cell"]),
                Paragraph(
                    _escape(str(hit.get("condition_label") or "—")), styles["cell"]
                ),
            ]
        )
    table = Table(rows, colWidths=[1.55 * inch, 0.85 * inch, 1.5 * inch, 1.6 * inch])
    table.setStyle(
        TableStyle(
            [
                ("BACKGROUND", (0, 0), (-1, 0), PANEL),
                ("BACKGROUND", (0, 1), (-1, -1), colors.white),
                ("ROWBACKGROUNDS", (0, 1), (-1, -1), [colors.white, ROW_ALT]),
                ("GRID", (0, 0), (-1, -1), 0.4, PANEL_EDGE),
                ("VALIGN", (0, 0), (-1, -1), "MIDDLE"),
                ("LEFTPADDING", (0, 0), (-1, -1), 5),
                ("RIGHTPADDING", (0, 0), (-1, -1), 5),
                ("TOPPADDING", (0, 0), (-1, -1), 4),
                ("BOTTOMPADDING", (0, 0), (-1, -1), 4),
            ]
        )
    )
    return table


def _metrics_snapshot(
    df: pd.DataFrame, styles: dict[str, ParagraphStyle]
) -> Table | None:
    if df.empty:
        return None
    focus_conds = (
        "base",
        "4shot",
        "lora",
        "cartridge",
        "rfdt_lora",
        "encoder_ft",
    )
    show = df[df["condition"].isin(focus_conds)].copy()
    if show.empty:
        return None
    order = {c: i for i, c in enumerate(ALL_CONDITIONS)}
    show["_ord"] = show["condition"].map(lambda c: order.get(c, 99))
    show = show.sort_values(["model", "_ord"])

    rows = [
        [
            Paragraph("Model", styles["cell_head"]),
            Paragraph("Condition", styles["cell_head"]),
            Paragraph("Macro-F1", styles["cell_head"]),
            Paragraph("Acc", styles["cell_head"]),
            Paragraph("Allowed-edge", styles["cell_head"]),
            Paragraph("Denied refuse", styles["cell_head"]),
        ]
    ]
    for row in show.itertuples(index=False):
        rows.append(
            [
                Paragraph(_escape(str(row.model)), styles["cell"]),
                Paragraph(_escape(str(row.condition_label)), styles["cell"]),
                Paragraph(_fmt(getattr(row, "macro_f1", None)), styles["cell"]),
                Paragraph(_fmt(getattr(row, "accuracy", None)), styles["cell"]),
                Paragraph(_fmt(getattr(row, "allowed_edge", None)), styles["cell"]),
                Paragraph(
                    _fmt(getattr(row, "denied_edge_refuse_recall", None)),
                    styles["cell"],
                ),
            ]
        )
    table = Table(
        rows,
        colWidths=[
            1.15 * inch,
            1.15 * inch,
            0.75 * inch,
            0.55 * inch,
            0.95 * inch,
            0.95 * inch,
        ],
    )
    table.setStyle(
        TableStyle(
            [
                ("BACKGROUND", (0, 0), (-1, 0), PANEL),
                ("ROWBACKGROUNDS", (0, 1), (-1, -1), [colors.white, ROW_ALT]),
                ("GRID", (0, 0), (-1, -1), 0.35, PANEL_EDGE),
                ("VALIGN", (0, 0), (-1, -1), "MIDDLE"),
                ("LEFTPADDING", (0, 0), (-1, -1), 4),
                ("RIGHTPADDING", (0, 0), (-1, -1), 4),
                ("TOPPADDING", (0, 0), (-1, -1), 3),
                ("BOTTOMPADDING", (0, 0), (-1, -1), 3),
            ]
        )
    )
    return table


def _disagreement_table(
    domain: str, styles: dict[str, ParagraphStyle]
) -> Table | None:
    disagree = rfdt_encoder_disagreement(domain, matched_8b_only=True)
    summary = disagreement_summary(disagree)
    if summary.empty:
        return None
    order = ["both_correct", "rfdt_only", "encoder_only", "both_wrong"]
    rows = [
        [
            Paragraph("RFDT model", styles["cell_head"]),
            Paragraph("Both correct", styles["cell_head"]),
            Paragraph("RFDT only", styles["cell_head"]),
            Paragraph("Encoder only", styles["cell_head"]),
            Paragraph("Both wrong", styles["cell_head"]),
            Paragraph("N", styles["cell_head"]),
        ]
    ]
    for model, group in summary.groupby("model"):
        rates = {r.outcome: float(r.rate) for r in group.itertuples(index=False)}
        total = int(group["total"].iloc[0])
        rows.append(
            [
                Paragraph(_escape(str(model)), styles["cell"]),
                *[
                    Paragraph(_fmt(rates.get(key)), styles["cell"])
                    for key in order
                ],
                Paragraph(str(total), styles["cell"]),
            ]
        )
    table = Table(
        rows,
        colWidths=[
            1.25 * inch,
            0.9 * inch,
            0.85 * inch,
            0.95 * inch,
            0.85 * inch,
            0.45 * inch,
        ],
    )
    table.setStyle(
        TableStyle(
            [
                ("BACKGROUND", (0, 0), (-1, 0), PANEL),
                ("ROWBACKGROUNDS", (0, 1), (-1, -1), [colors.white, ROW_ALT]),
                ("GRID", (0, 0), (-1, -1), 0.35, PANEL_EDGE),
                ("VALIGN", (0, 0), (-1, -1), "MIDDLE"),
                ("LEFTPADDING", (0, 0), (-1, -1), 4),
                ("RIGHTPADDING", (0, 0), (-1, -1), 4),
                ("TOPPADDING", (0, 0), (-1, -1), 3),
                ("BOTTOMPADDING", (0, 0), (-1, -1), 3),
            ]
        )
    )
    return table


def _exec_summary_bits(
    domain_packets: list[tuple[str, pd.DataFrame, dict[str, Any], Any]],
) -> list[str]:
    bits: list[str] = []
    for domain, _df, stats, bundle in domain_packets:
        label = DOMAIN_LABELS.get(domain, domain)
        industry = "healthcare" if domain == "medicareplus" else "government"
        best_rfdt = stats.get("best_rfdt")
        best_enc = stats.get("best_encoder")
        best_lora = stats.get("best_lora")
        best_frozen = stats.get("best_frozen")
        if best_lora and best_frozen:
            bits.append(
                f"{label} ({industry}): best LoRA {_fmt(best_lora['value'])} "
                f"({best_lora['model']}) vs best frozen {_fmt(best_frozen['value'])} "
                f"({best_frozen['condition_label']})."
            )
        if best_rfdt and best_enc:
            bits.append(
                f"{label}: RFDT-LoRA {_fmt(best_rfdt['value'])} "
                f"({best_rfdt['model']}) vs ModernBERT Encoder-FT "
                f"{_fmt(best_enc['value'])} on the shared held-out contract."
            )
        elif best_rfdt:
            bits.append(
                f"{label}: best RFDT-LoRA {_fmt(best_rfdt['value'])} "
                f"({best_rfdt['model']}); Encoder-FT not in this export packet."
            )
        if bundle.findings and bundle.findings.body:
            # One compressed sentence from findings for the brief.
            first = bundle.findings.body.split("\n\n")[0].split(". ")[0].strip()
            if first and not first.endswith("."):
                first += "."
            if first:
                bits.append(first)
    return bits[:8]


def _make_footer(generated: str):
    def _footer(canvas, doc) -> None:
        canvas.saveState()
        canvas.setStrokeColor(PANEL_EDGE)
        canvas.setLineWidth(0.4)
        canvas.line(0.75 * inch, 0.55 * inch, letter[0] - 0.75 * inch, 0.55 * inch)
        canvas.setFont("Helvetica", 7.5)
        canvas.setFillColor(MUTED)
        canvas.drawString(
            0.75 * inch,
            0.38 * inch,
            f"COMPASS-PRAXA Policy Alignment · Generated {generated}",
        )
        canvas.drawRightString(
            letter[0] - 0.75 * inch,
            0.38 * inch,
            f"Page {doc.page}",
        )
        canvas.restoreState()

    return _footer


def build_alignment_report_pdf(*, force_refresh: bool = False) -> bytes:
    """Build a multi-domain PDF briefing from the current Hub cache."""

    refresh_cache(force=force_refresh)
    styles = _styles()
    domains = tuple(DOMAIN_PREFIX.keys())

    packets: list[tuple[str, pd.DataFrame, dict[str, Any], Any]] = []
    for domain in domains:
        df = _domain_frame(domain)
        stats = best_by_family(df, metric="macro_f1")
        bundle = narrative_bundle(df, domain)
        packets.append((domain, df, stats, bundle))

    buffer = BytesIO()
    doc = SimpleDocTemplate(
        buffer,
        pagesize=letter,
        leftMargin=0.75 * inch,
        rightMargin=0.75 * inch,
        topMargin=0.7 * inch,
        bottomMargin=0.75 * inch,
        title="Praxa Behavioral Alignment Report",
        author=RESEARCHER,
    )
    story: list[Any] = []
    generated = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")

    story.append(
        Paragraph(f"Researcher · {RESEARCHER}", styles["byline"])
    )
    story.append(Paragraph("Praxa Behavioral Alignment", styles["cover_title"]))
    story.append(
        Paragraph("COMPASS-PRAXA Policy Alignment", styles["cover_sub"])
    )
    story.append(Spacer(1, 0.14 * inch))

    story.append(Paragraph("Executive summary", styles["h1"]))
    story.append(
        Paragraph(
            "Matched-primary 8B instruct models (frozen prompting, LoRA, "
            "Behavioral KV, RFDT) plus ModernBERT-base Encoder-FT comparator. "
            "Numbers below are taken from published Hub metrics "
            f"({EXPERIMENTS_DATASET_REPO}). Domains are split by industry "
            "scenario. The Space remains the interactive source; this PDF is a "
            "static briefing for sharing.",
            styles["body"],
        )
    )
    for bit in _exec_summary_bits(packets):
        story.append(Paragraph(f"• {_escape(bit)}", styles["bullet"]))
    if not any(not p[1].empty for p in packets):
        story.append(
            Paragraph(
                "• No published metrics were available when this report was built.",
                styles["bullet"],
            )
        )

    story.append(Paragraph("Cross-domain headline scores", styles["h2"]))
    cross_rows = [
        [
            Paragraph("Domain", styles["cell_head"]),
            Paragraph("Industry", styles["cell_head"]),
            Paragraph("Best LoRA", styles["cell_head"]),
            Paragraph("Best RFDT", styles["cell_head"]),
            Paragraph("Encoder-FT", styles["cell_head"]),
        ]
    ]
    for domain, _df, stats, _bundle in packets:
        industry = "Healthcare" if domain == "medicareplus" else "Government"
        cross_rows.append(
            [
                Paragraph(
                    _escape(DOMAIN_LABELS.get(domain, domain).split(" (")[0]),
                    styles["cell"],
                ),
                Paragraph(industry, styles["cell"]),
                Paragraph(
                    _fmt((stats.get("best_lora") or {}).get("value")),
                    styles["cell"],
                ),
                Paragraph(
                    _fmt((stats.get("best_rfdt") or {}).get("value")),
                    styles["cell"],
                ),
                Paragraph(
                    _fmt((stats.get("best_encoder") or {}).get("value")),
                    styles["cell"],
                ),
            ]
        )
    cross = Table(
        cross_rows,
        colWidths=[1.4 * inch, 1.0 * inch, 1.1 * inch, 1.1 * inch, 1.1 * inch],
    )
    cross.setStyle(
        TableStyle(
            [
                ("BACKGROUND", (0, 0), (-1, 0), PANEL),
                ("ROWBACKGROUNDS", (0, 1), (-1, -1), [colors.white, ROW_ALT]),
                ("GRID", (0, 0), (-1, -1), 0.35, PANEL_EDGE),
                ("VALIGN", (0, 0), (-1, -1), "MIDDLE"),
                ("LEFTPADDING", (0, 0), (-1, -1), 5),
                ("RIGHTPADDING", (0, 0), (-1, -1), 5),
                ("TOPPADDING", (0, 0), (-1, -1), 4),
                ("BOTTOMPADDING", (0, 0), (-1, -1), 4),
            ]
        )
    )
    story.append(cross)

    for domain, df, stats, bundle in packets:
        story.append(PageBreak())
        label = DOMAIN_LABELS.get(domain, domain)
        industry = "Healthcare" if domain == "medicareplus" else "Government"
        story.append(Paragraph(label, styles["h1"]))
        story.append(
            Paragraph(
                f"Industry scenario: {industry}. Held-out COMPASS evaluation on "
                "published metrics.json / predictions.jsonl.",
                styles["body"],
            )
        )

        if df.empty:
            story.append(
                Paragraph(
                    "No published runs for this domain in the current Hub cache.",
                    styles["body"],
                )
            )
            continue

        story.append(Paragraph("Headline methods", styles["h2"]))
        story.append(_bests_table(stats, styles))
        story.append(Spacer(1, 0.12 * inch))

        if bundle.findings and bundle.findings.body:
            story.append(Paragraph("What the data shows", styles["h2"]))
            for part in bundle.findings.body.split("\n\n"):
                part = part.strip()
                if part:
                    story.append(Paragraph(_escape(part), styles["body"]))

        if bundle.implications and bundle.implications.body:
            story.append(
                Paragraph(
                    bundle.implications.title or "What this means",
                    styles["h2"],
                )
            )
            for part in str(bundle.implications.body).split("\n\n"):
                text = part.strip()
                if text:
                    story.append(Paragraph(_escape(text), styles["body"]))

        story.append(Paragraph("Metrics snapshot", styles["h2"]))
        story.append(
            Paragraph(
                "Selected ladder conditions plus Encoder-FT when published.",
                styles["body"],
            )
        )
        snap = _metrics_snapshot(df, styles)
        if snap is not None:
            story.append(snap)
        else:
            story.append(Paragraph("No snapshot rows available.", styles["body"]))

        story.append(
            Paragraph("RFDT-LoRA vs Encoder-FT disagreement", styles["h2"])
        )
        story.append(
            Paragraph(
                "Paired on sample_id. Rates are shares of the held-out test set "
                "where both methods agree or only one is correct.",
                styles["body"],
            )
        )
        disagree_table = _disagreement_table(domain, styles)
        if disagree_table is not None:
            story.append(disagree_table)
        else:
            story.append(
                Paragraph(
                    "Predictions not available for a paired join in this export.",
                    styles["body"],
                )
            )

        # Light comparator callout from encoder frame.
        comp = encoder_comparator_frame(df)
        if not comp.empty and "macro_f1" in comp.columns:
            story.append(Paragraph("Comparator macro-F1", styles["h2"]))
            for row in comp.sort_values(["condition", "model"]).itertuples(
                index=False
            ):
                story.append(
                    Paragraph(
                        f"• {_escape(str(row.model))} · "
                        f"{_escape(str(row.condition_label))}: "
                        f"{_fmt(getattr(row, 'macro_f1', None))}",
                        styles["bullet"],
                    )
                )

    story.append(PageBreak())
    story.append(Paragraph("Published artifacts", styles["h1"]))
    story.append(
        Paragraph(
            "Weight repos (adapters, Behavioral KV cartridges, Encoder-FT "
            "classifiers) and the experiments dataset used by this Space.",
            styles["body"],
        )
    )
    story.append(Paragraph("Weight repos", styles["h2"]))
    for entry in PUBLISHED_MODEL_REPOS:
        story.append(
            Paragraph(
                f"• <b>{_escape(entry['label'])}</b> — "
                f"https://huggingface.co/{_escape(entry['repo'])}<br/>"
                f"&nbsp;&nbsp;{_escape(entry['blurb'])}",
                styles["bullet"],
            )
        )
    story.append(Paragraph("Experiments dataset", styles["h2"]))
    story.append(
        Paragraph(
            f"• https://huggingface.co/datasets/{EXPERIMENTS_DATASET_REPO}",
            styles["bullet"],
        )
    )
    story.append(Paragraph("Notes", styles["h2"]))
    story.append(
        Paragraph(
            "• COMPASS query and policy plaintext are not redistributed here. "
            "This briefing reports metrics and derived summaries only. Obtain "
            "the source text from the upstream COMPASS dataset under its terms: "
            "https://huggingface.co/datasets/AIM-Intelligence/COMPASS-Policy-Alignment-Testbed-Dataset. "
            "Cite: Choi, D., Lee, D., Kartono, B. J., Berndt, H., Kwon, T., "
            "Jang, J., Park, H., Yu, H., &amp; Kahng, M. (2026). "
            "<i>COMPASS: A framework for evaluating organization-specific "
            "policy alignment in LLMs</i> (arXiv:2601.01836). arXiv. "
            "https://arxiv.org/abs/2601.01836",
            styles["bullet"],
        )
    )
    story.append(
        Paragraph(
            "• Interactive charts (ladders, slice heatmaps, edge panels) live on "
            "the Hugging Face Space. Refresh the Space before exporting if you "
            "need the newest Hub slots.",
            styles["bullet"],
        )
    )

    doc.build(
        story,
        onFirstPage=_make_footer(generated),
        onLaterPages=_make_footer(generated),
    )
    return buffer.getvalue()


def report_filename(*, when: datetime | None = None) -> str:
    stamp = (when or datetime.now(timezone.utc)).strftime("%Y%m%d")
    return f"praxa-behavioral-alignment-report-{stamp}.pdf"