"""
audit_report.py - generate a bank-letterhead-styled PDF audit report.

Used by app.py to produce the 'Download audit PDF' button.
Returns bytes so it can be streamed directly to the user without disk writes.
"""

import io
from datetime import datetime
from pathlib import Path

from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import cm
from reportlab.lib import colors
from reportlab.platypus import (SimpleDocTemplate, Paragraph, Spacer, Table,
                                 TableStyle, Image as RLImage, PageBreak)
from reportlab.lib.enums import TA_CENTER, TA_LEFT
import cv2
import numpy as np
from PIL import Image as PILImage

import forensics


BAND_COLOURS = {
    "LOW":      colors.HexColor("#16a34a"),
    "MEDIUM":   colors.HexColor("#ca8a04"),
    "HIGH":     colors.HexColor("#ea580c"),
    "CRITICAL": colors.HexColor("#dc2626"),
}


def _styles():
    s = getSampleStyleSheet()
    s.add(ParagraphStyle(name="LetterheadBank", parent=s["Title"],
                          fontSize=22, textColor=colors.HexColor("#1e3a8a"),
                          spaceAfter=4))
    s.add(ParagraphStyle(name="LetterheadSub", parent=s["Normal"],
                          fontSize=10, textColor=colors.grey,
                          spaceAfter=12, alignment=TA_CENTER))
    s.add(ParagraphStyle(name="SectionH", parent=s["Heading2"],
                          fontSize=13, textColor=colors.HexColor("#1e3a8a"),
                          spaceAfter=6, spaceBefore=12))
    s.add(ParagraphStyle(name="Evidence", parent=s["Normal"],
                          fontSize=10, leftIndent=12, bulletIndent=2,
                          spaceAfter=2))
    s.add(ParagraphStyle(name="Mono", parent=s["Normal"],
                          fontName="Courier", fontSize=8,
                          textColor=colors.dimgray))
    return s


def _pil_to_flowable(pil_img, max_width=15 * cm, max_height=8 * cm):
    """Convert a PIL image to a sized RL Image flowable."""
    buf = io.BytesIO()
    pil_img.convert("RGB").save(buf, "PNG")
    buf.seek(0)
    img = RLImage(buf)
    # scale preserving aspect ratio
    iw, ih = pil_img.size
    ratio = min(max_width / iw, max_height / ih)
    img.drawWidth  = iw * ratio
    img.drawHeight = ih * ratio
    return img


def _make_heatmap_image(heat, cmap="hot"):
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots(figsize=(5, 3))
    ax.imshow(heat, cmap=cmap)
    ax.axis("off")
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=110, bbox_inches="tight")
    plt.close(fig)
    buf.seek(0)
    return PILImage.open(buf)


def build_pdf_report(report, source_path):
    """
    report:        the dict returned by forensics.analyse_document(...)
    source_path:   Path to the original document being analysed
    Returns:       bytes of the rendered PDF
    """
    source_path = Path(source_path)
    s = _styles()
    buf = io.BytesIO()
    doc = SimpleDocTemplate(buf, pagesize=A4,
                            leftMargin=2 * cm, rightMargin=2 * cm,
                            topMargin=1.5 * cm, bottomMargin=1.5 * cm)
    story = []

    # ---- Letterhead ----
    story.append(Paragraph("DOCSENTRY - DOCUMENT FORENSICS REPORT",
                           s["LetterheadBank"]))
    story.append(Paragraph("Confidential - For Underwriting Use Only",
                           s["LetterheadSub"]))

    # ---- Document metadata table ----
    meta_data = [
        ["Field",          "Value"],
        ["Document",       source_path.name],
        ["Type",           report.get("type", "-")],
        ["Analysed at",    report.get("analysed_at", "-")[:19].replace("T", " ")],
        ["SHA-256",        report.get("sha256", "-")[:32] + "..."],
    ]
    t = Table(meta_data, colWidths=[4 * cm, 13 * cm])
    t.setStyle(TableStyle([
        ("BACKGROUND", (0, 0), (-1, 0), colors.HexColor("#1e3a8a")),
        ("TEXTCOLOR",  (0, 0), (-1, 0), colors.white),
        ("GRID",       (0, 0), (-1, -1), 0.4, colors.grey),
        ("FONTNAME",   (0, 0), (-1, 0), "Helvetica-Bold"),
        ("FONTSIZE",   (0, 0), (-1, -1), 9),
        ("ROWBACKGROUNDS", (0, 1), (-1, -1),
            [colors.HexColor("#f8fafc"), colors.white]),
    ]))
    story.append(t)
    story.append(Spacer(1, 0.6 * cm))

    # ---- Risk verdict ----
    band_str = report.get("risk_band", "UNKNOWN")
    band_colour = BAND_COLOURS.get(band_str, colors.grey)
    verdict_table = Table([
        [Paragraph(f"<para alignment='center'><font size=22 color='white'>"
                    f"<b>{band_str}</b></font></para>", s["Normal"]),
         Paragraph(f"<b>Risk score:</b> {report.get('risk_score', '-')}<br/>"
                    f"<b>Action:</b> {report.get('recommended_action', '-')}",
                    s["Normal"])],
    ], colWidths=[5 * cm, 12 * cm])
    verdict_table.setStyle(TableStyle([
        ("BACKGROUND", (0, 0), (0, 0), band_colour),
        ("BACKGROUND", (1, 0), (1, 0), colors.HexColor("#f1f5f9")),
        ("VALIGN",     (0, 0), (-1, -1), "MIDDLE"),
        ("BOX",        (0, 0), (-1, -1), 0.4, colors.grey),
        ("LEFTPADDING",  (0, 0), (-1, -1), 12),
        ("RIGHTPADDING", (0, 0), (-1, -1), 12),
        ("TOPPADDING",   (0, 0), (-1, -1), 12),
        ("BOTTOMPADDING",(0, 0), (-1, -1), 12),
    ]))
    story.append(verdict_table)
    story.append(Spacer(1, 0.4 * cm))

    # ---- Sub-score breakdown ----
    story.append(Paragraph("Sub-score breakdown", s["SectionH"]))
    sub_rows = [["Detector", "Score (0=clean, 1=suspicious)"]]
    for k, v in (report.get("sub_scores") or {}).items():
        bar = "#" * int(v * 30)
        sub_rows.append([k, f"{v:.2f}  {bar}"])
    if len(sub_rows) > 1:
        sub_t = Table(sub_rows, colWidths=[5 * cm, 12 * cm])
        sub_t.setStyle(TableStyle([
            ("BACKGROUND", (0, 0), (-1, 0), colors.HexColor("#1e3a8a")),
            ("TEXTCOLOR",  (0, 0), (-1, 0), colors.white),
            ("GRID",       (0, 0), (-1, -1), 0.4, colors.grey),
            ("FONTNAME",   (1, 1), (1, -1), "Courier"),
            ("FONTSIZE",   (0, 0), (-1, -1), 9),
        ]))
        story.append(sub_t)
    story.append(Spacer(1, 0.4 * cm))

    # ---- Evidence ----
    story.append(Paragraph("Forensic evidence", s["SectionH"]))
    for ev in report.get("evidence", []):
        story.append(Paragraph(f"&bull; {ev}", s["Evidence"]))
    story.append(Spacer(1, 0.3 * cm))

    # ---- Image-specific: heatmaps ----
    if report.get("type") == "image":
        try:
            story.append(PageBreak())
            story.append(Paragraph("Forensic visualizations", s["SectionH"]))

            # ELA
            ela_img, ela_s = forensics.error_level_analysis(source_path)
            story.append(Paragraph(f"<b>Error Level Analysis</b> (score: {ela_s:.2f})",
                                    s["Normal"]))
            story.append(_pil_to_flowable(ela_img))
            story.append(Spacer(1, 0.3 * cm))

            # Copy-move
            viz, n_cm, _ = forensics.copy_move_detect(source_path)
            if viz is not None:
                story.append(Paragraph(f"<b>Copy-move matches:</b> {n_cm}",
                                        s["Normal"]))
                viz_rgb = cv2.cvtColor(viz, cv2.COLOR_BGR2RGB)
                story.append(_pil_to_flowable(PILImage.fromarray(viz_rgb)))
                story.append(Spacer(1, 0.3 * cm))

            # Noise heatmap
            heat, ratio = forensics.noise_inconsistency(source_path)
            story.append(Paragraph(f"<b>Noise outlier ratio:</b> {ratio:.2%}",
                                    s["Normal"]))
            story.append(_pil_to_flowable(_make_heatmap_image(heat)))
        except Exception as e:
            story.append(Paragraph(f"Could not render heatmaps: {e}", s["Normal"]))

    # ---- PDF-specific audit details ----
    if report.get("type") == "pdf":
        audit = report.get("pdf_audit", {})
        fonts = report.get("font_audit", {})
        story.append(Paragraph("PDF structural audit", s["SectionH"]))
        meta = audit.get("metadata", {}) or {}
        pdf_rows = [
            ["Pages",       str(audit.get("pages", "-"))],
            ["EOF markers", str(audit.get("eof_markers", "-"))],
            ["Producer",    str(meta.get("producer", "-"))],
            ["Creator",     str(meta.get("creator", "-"))],
            ["Fonts used",  ", ".join(fonts.get("fonts", []) or ["-"])],
        ]
        pdf_t = Table(pdf_rows, colWidths=[5 * cm, 12 * cm])
        pdf_t.setStyle(TableStyle([
            ("GRID",     (0, 0), (-1, -1), 0.4, colors.grey),
            ("FONTSIZE", (0, 0), (-1, -1), 9),
            ("BACKGROUND", (0, 0), (0, -1), colors.HexColor("#f1f5f9")),
        ]))
        story.append(pdf_t)
        story.append(Spacer(1, 0.3 * cm))
        for f in audit.get("flags", []) + fonts.get("flags", []):
            if f not in ("clean", "ok"):
                story.append(Paragraph(f"&bull; {f}", s["Evidence"]))

    # ---- ML model verdict (if present) ----
    if "ml_prediction" in report:
        ml = report["ml_prediction"]
        story.append(Paragraph("Trained ML model verdict", s["SectionH"]))
        ml_t = Table([
            ["Tamper probability", f"{ml['tamper_probability']:.1%}"],
            ["Verdict",            ml["verdict"]],
        ], colWidths=[5 * cm, 12 * cm])
        ml_t.setStyle(TableStyle([
            ("GRID", (0, 0), (-1, -1), 0.4, colors.grey),
            ("FONTSIZE", (0, 0), (-1, -1), 9),
            ("BACKGROUND", (0, 0), (0, -1), colors.HexColor("#f1f5f9")),
        ]))
        story.append(ml_t)

    # ---- Footer ----
    story.append(Spacer(1, 0.6 * cm))
    story.append(Paragraph(
        "<i>This report was generated automatically by DocSentry. "
        "Findings are based on forensic-signal heuristics and an explainable "
        "Random Forest classifier. Manual verification is required for any "
        "file in HIGH or CRITICAL bands.</i>", s["Mono"]))

    doc.build(story)
    buf.seek(0)
    return buf.read()


if __name__ == "__main__":
    # CLI smoke test
    import sys
    if len(sys.argv) < 2:
        print("Usage: python audit_report.py <file>")
        sys.exit(1)
    src = Path(sys.argv[1])
    report = forensics.analyse_document(src)
    pdf_bytes = build_pdf_report(report, src)
    out = src.with_suffix(".audit.pdf")
    out.write_bytes(pdf_bytes)
    print(f"Wrote {out}  ({len(pdf_bytes)} bytes)")