"""Build a print-quality PDF from PAPER-tinyllama-rv32.md.

    python make_paper_pdf.py [PAPER.md] [OUT.pdf]

Renders headings, paragraphs, bullets, tables, fenced code, figures and captions on A4 with
running heads, page numbers and a title block. No LaTeX, no browser.
"""
import os
import re
import sys

from reportlab.lib import colors
from reportlab.lib.enums import TA_CENTER, TA_JUSTIFY, TA_LEFT
from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet
from reportlab.lib.units import mm
from reportlab.platypus import (BaseDocTemplate, Frame, Image, KeepTogether, PageBreak,
                                PageTemplate, Paragraph, Spacer, Table, TableStyle)

SRC = sys.argv[1] if len(sys.argv) > 1 else "PAPER-tinyllama-rv32.md"
OUT = sys.argv[2] if len(sys.argv) > 2 else "PAPER-tinyllama-rv32.pdf"
BASE = os.path.dirname(os.path.abspath(SRC))
INK = colors.HexColor("#111111")
MUTED = colors.HexColor("#555555")
RULE = colors.HexColor("#bbbbbb")
ACCENT = colors.HexColor("#1f3a5f")
CODEBG = colors.HexColor("#f4f5f7")

ss = getSampleStyleSheet()
TITLE = ParagraphStyle("t", parent=ss["Title"], fontName="Times-Bold", fontSize=21, leading=25,
                       textColor=INK, spaceAfter=6, alignment=TA_LEFT)
SUBTITLE = ParagraphStyle("st", parent=ss["Normal"], fontName="Times-Italic", fontSize=11.5,
                          leading=15, textColor=MUTED, spaceAfter=2)
H1 = ParagraphStyle("h1", parent=ss["Heading1"], fontName="Times-Bold", fontSize=14, leading=18,
                    textColor=ACCENT, spaceBefore=14, spaceAfter=6, keepWithNext=1)
H2 = ParagraphStyle("h2", parent=ss["Heading2"], fontName="Times-Bold", fontSize=11.5, leading=15,
                    textColor=INK, spaceBefore=10, spaceAfter=4, keepWithNext=1)
BODY = ParagraphStyle("b", parent=ss["BodyText"], fontName="Times-Roman", fontSize=10, leading=14.2,
                      alignment=TA_JUSTIFY, spaceAfter=5)
BULLET = ParagraphStyle("bl", parent=BODY, leftIndent=12, bulletIndent=2, spaceAfter=3)
QUOTE = ParagraphStyle("q", parent=BODY, fontName="Times-Italic", leftIndent=10, rightIndent=6,
                       textColor=colors.HexColor("#333333"), spaceBefore=2, spaceAfter=6)
CODE = ParagraphStyle("c", parent=ss["Code"], fontName="Courier", fontSize=7.4, leading=9.4,
                      textColor=INK, backColor=CODEBG, borderPadding=(4, 5, 4, 5),
                      leftIndent=6, rightIndent=6, spaceBefore=3, spaceAfter=7)
CAPTION = ParagraphStyle("cap", parent=ss["Normal"], fontName="Times-Italic", fontSize=8.4,
                         leading=11, textColor=MUTED, alignment=TA_CENTER, spaceBefore=3,
                         spaceAfter=9)
CELL = ParagraphStyle("cell", parent=ss["Normal"], fontName="Times-Roman", fontSize=8.2,
                      leading=10.4, textColor=INK)
CELLH = ParagraphStyle("cellh", parent=CELL, fontName="Times-Bold")
CELLR = ParagraphStyle("cellr", parent=CELL, alignment=2)


def inline(text):
    text = text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
    text = re.sub(r"`([^`]+)`", r'<font name="Courier" size="8.6">\1</font>', text)
    text = re.sub(r"\*\*([^*]+)\*\*", r"<b>\1</b>", text)
    text = re.sub(r"(?<!\*)\*([^*]+)\*(?!\*)", r"<i>\1</i>", text)
    text = re.sub(r"\[([^\]]+)\]\(([^)]+)\)", r'<font color="#1f3a5f">\1</font>', text)
    return text


def is_number(cell):
    c = cell.strip().replace(",", "").replace("%", "").replace("**", "")
    try:
        float(c)
        return True
    except ValueError:
        return False


def table_from(rows):
    has_head = len(rows) > 1 and not all(set(c) <= set("-: ") for c in rows[1])
    head = rows[0] if has_head else None
    body = rows[1:] if has_head else rows
    data = []
    if head:
        data.append([Paragraph(inline(c), CELLH) for c in head])
    for r in body:
        data.append([Paragraph(inline(c), CELLR if is_number(c) and i else CELL)
                     for i, c in enumerate(r)])
    ncol = max(len(r) for r in data)
    for r in data:
        while len(r) < ncol:
            r.append(Paragraph("", CELL))
    width = 170 * mm
    first = 58 * mm if ncol > 2 else width / ncol
    cols = [first] + [(width - first) / (ncol - 1)] * (ncol - 1) if ncol > 1 else [width]
    t = Table(data, colWidths=cols, hAlign="LEFT", repeatRows=1 if head else 0)
    style = [("VALIGN", (0, 0), (-1, -1), "TOP"),
             ("TOPPADDING", (0, 0), (-1, -1), 3),
             ("BOTTOMPADDING", (0, 0), (-1, -1), 3),
             ("LEFTPADDING", (0, 0), (-1, -1), 5),
             ("RIGHTPADDING", (0, 0), (-1, -1), 5),
             ("LINEBELOW", (0, 0), (-1, -1), 0.25, RULE)]
    if head:
        style += [("BACKGROUND", (0, 0), (-1, 0), colors.HexColor("#eceff3")),
                  ("LINEABOVE", (0, 0), (-1, 0), 0.9, ACCENT),
                  ("LINEBELOW", (0, 0), (-1, 0), 0.6, ACCENT)]
    t.setStyle(TableStyle(style))
    return t


def figure(path):
    full = os.path.join(BASE, path)
    if not os.path.exists(full):
        return None
    from PIL import Image as PILImage
    w, h = PILImage.open(full).size
    width = min(150 * mm, w * 0.75)
    return Image(full, width=width, height=width * h / w)


def parse(md):
    flow = []
    lines = md.splitlines()
    i = 0
    in_code = False
    code_buf = []
    title_done = False
    para = []
    bullets = []
    table = []

    def flush_para():
        nonlocal para
        if para:
            flow.append(Paragraph(inline(" ".join(para)), BODY))
            para = []

    def flush_bullets():
        nonlocal bullets
        for b in bullets:
            flow.append(Paragraph(inline(b), BULLET, bulletText="\u2022"))
        if bullets:
            flow.append(Spacer(1, 2))
        bullets = []

    def flush_table():
        nonlocal table
        if table:
            flow.append(table_from(table))
            flow.append(Spacer(1, 7))
        table = []

    while i < len(lines):
        ln = lines[i]
        stripped = ln.strip()

        if stripped.startswith("```"):
            flush_para(); flush_bullets(); flush_table()
            if in_code:
                flow.append(Paragraph("<br/>".join(c.replace(" ", "&nbsp;") for c in code_buf)
                                      .replace("&nbsp;&nbsp;&nbsp;&nbsp;", "&nbsp;&nbsp;&nbsp;&nbsp;"),
                                      CODE))
                code_buf = []
                in_code = False
            else:
                in_code = True
            i += 1
            continue
        if in_code:
            code_buf.append(ln if ln else "&nbsp;")
            i += 1
            continue

        if not stripped:
            flush_para(); flush_bullets(); flush_table()
            i += 1
            continue

        if stripped.startswith("|"):
            flush_para(); flush_bullets()
            cells = [c.strip() for c in stripped.strip("|").split("|")]
            if not all(set(c) <= set("-: ") for c in cells):
                table.append(cells)
            i += 1
            continue

        m = re.match(r"^!\[([^\]]*)\]\(([^)]+)\)", stripped)
        if m:
            flush_para(); flush_bullets(); flush_table()
            cap = m.group(1)
            j = i + 1
            while j < len(lines) and not lines[j].strip():
                j += 1
            if j < len(lines):
                nxt = lines[j].strip()
                mm = re.match(r"^\*?(Figure\s+\d+\.[^*]*)\*?$", nxt.lstrip("*_").rstrip("*_"))
                if mm:
                    cap = mm.group(1)
                    i = j
            img = figure(m.group(2))
            if img:
                flow.append(KeepTogether([img, Paragraph(inline(cap), CAPTION)]))
            i += 1
            continue

        m = re.match(r"^#\s+(.*)$", stripped)
        if m and not title_done:
            flush_para(); flush_bullets(); flush_table()
            flow.append(Paragraph(inline(m.group(1)), TITLE))
            title_done = True
            i += 1
            continue

        m = re.match(r"^##\s+(.*)$", stripped)
        if m:
            flush_para(); flush_bullets(); flush_table()
            flow.append(Paragraph(inline(m.group(1)), H1))
            i += 1
            continue

        m = re.match(r"^###\s+(.*)$", stripped)
        if m:
            flush_para(); flush_bullets(); flush_table()
            flow.append(Paragraph(inline(m.group(1)), H2))
            i += 1
            continue

        if stripped.startswith(">"):
            flush_para(); flush_bullets(); flush_table()
            flow.append(Paragraph(inline(stripped.lstrip("> ").strip()), QUOTE))
            i += 1
            continue

        m = re.match(r"^[-*]\s+(.*)$", stripped)
        if m:
            flush_para(); flush_table()
            bullets.append(m.group(1))
            i += 1
            continue

        m = re.match(r"^\d+\.\s+(.*)$", stripped)
        if m:
            flush_para(); flush_table()
            bullets.append(m.group(1))
            i += 1
            continue

        if bullets and (ln.startswith(" ") or ln.startswith("\t")):
            bullets[-1] = bullets[-1] + " " + stripped
            i += 1
            continue

        if stripped in ("---", "***"):
            flush_para(); flush_bullets(); flush_table()
            flow.append(Spacer(1, 6))
            i += 1
            continue

        if title_done and not para and not flow[-1:] == [None] and \
                getattr(flow[-1], "style", None) is TITLE:
            flow.append(Paragraph(inline(stripped), SUBTITLE))
            i += 1
            continue

        para.append(stripped)
        i += 1

    flush_para(); flush_bullets(); flush_table()
    return flow


def decorate(canvas, doc):
    canvas.saveState()
    canvas.setFont("Times-Roman", 8)
    canvas.setFillColor(MUTED)
    if doc.page > 1:
        canvas.drawString(20 * mm, A4[1] - 13 * mm,
                          "TinyLLama-v0 as an int4 workload on a token-level RISC-V machine")
        canvas.setStrokeColor(RULE)
        canvas.setLineWidth(0.3)
        canvas.line(20 * mm, A4[1] - 15 * mm, A4[0] - 20 * mm, A4[1] - 15 * mm)
    canvas.drawRightString(A4[0] - 20 * mm, 12 * mm, str(doc.page))
    canvas.restoreState()


def main():
    md = open(SRC, encoding="utf-8").read()
    flow = parse(md)
    doc = BaseDocTemplate(OUT, pagesize=A4, leftMargin=20 * mm, rightMargin=20 * mm,
                          topMargin=20 * mm, bottomMargin=18 * mm,
                          title="TinyLLama-v0 as an int4 workload on a token-level RISC-V machine",
                          author="OpceanAI", subject="int4 inference on a token-level RV32 machine")
    frame = Frame(doc.leftMargin, doc.bottomMargin, doc.width, doc.height, id="f")
    doc.addPageTemplates([PageTemplate(id="p", frames=[frame], onPage=decorate)])
    doc.build(flow)
    print("escrito", OUT, os.path.getsize(OUT) // 1024, "KB")


if __name__ == "__main__":
    main()