import fitz
from io import BytesIO


def overlay_pdf(content_pdf_bytes, template_pdf_path):
    """
    Overlay every page of the generated PDF onto the first page
    of the uploaded template PDF.

    Parameters
    ----------
    content_pdf_bytes : bytes
        PDF generated from agreement HTML.

    template_pdf_path : str
        Absolute path of uploaded template PDF.

    Returns
    -------
    bytes
        Final merged PDF.
    """

    # Open generated PDF
    content_pdf = fitz.open(
        stream=content_pdf_bytes,
        filetype="pdf"
    )

    # Open template PDF
    template_pdf = fitz.open(template_pdf_path)

    if len(template_pdf) == 0:
        raise Exception("Template PDF is empty.")

    template_page = template_pdf[0]

    output_pdf = fitz.open()

    for page in content_pdf:

        width = page.rect.width
        height = page.rect.height

        new_page = output_pdf.new_page(
            width=width,
            height=height
        )

        # Draw template FIRST (background)
        new_page.show_pdf_page(
            new_page.rect,
            template_pdf,
            0,
            overlay=False
        )

        # Draw generated agreement ABOVE template
        new_page.show_pdf_page(
            new_page.rect,
            content_pdf,
            page.number,
            overlay=True
        )

    pdf_bytes = output_pdf.tobytes()

    output_pdf.close()
    content_pdf.close()
    template_pdf.close()

    return pdf_bytes