import argparse
import fitz  # PyMuPDF


def side_by_side_preserve_links(input_pdf: str, output_pdf: str, start_page: int = 1):
    src = fitz.open(input_pdf)
    out = fitz.open()

    # Map each source page -> (output_page_index, x_offset, y_offset)
    # In this simple layout:
    # - even source pages go on the LEFT at x_offset=0
    # - odd source pages go on the RIGHT at x_offset=width_of_left_page (of that pair)
    placement = {}

    # 1) Create output pages and place page content, while recording placement
    out_page_index = 0
    i = start_page - 1
    while i < src.page_count:
        left = src.load_page(i)
        left_rect = left.rect

        right = src.load_page(i + 1) if (i + 1) < src.page_count else None
        right_rect = right.rect if right else fitz.Rect(0, 0, 0, 0)

        new_width = left_rect.width + (right_rect.width if right else 0) - 1
        new_height = max(left_rect.height, right_rect.height if right else 0)

        new_page = out.new_page(width=new_width, height=new_height)

        # Show left page
        left_target = fitz.Rect(0, 0, left_rect.width, left_rect.height)
        new_page.show_pdf_page(left_target, src, i)
        placement[i] = (out_page_index, 0.0, 0.0)

        # Show right page
        if right:
            right_x = left_rect.width - 1
            right_target = fitz.Rect(right_x, 0, right_x + right_rect.width, right_rect.height)
            new_page.show_pdf_page(right_target, src, i + 1)
            placement[i + 1] = (out_page_index, right_x, 0.0)

        out_page_index += 1
        i += 2

    # Helper: map a source destination (page + point) into the 2-up output coordinate space
    def map_dest(src_page_num: int, pt: fitz.Point | None):
        # Destination output page index
        dest_out_page, dest_xoff, dest_yoff = placement[src_page_num]
        if pt is None:
            return dest_out_page, None
        return dest_out_page, fitz.Point(pt.x + dest_xoff, pt.y + dest_yoff)

    # Helper: map a source link rectangle ("from") into the output coordinate space
    def map_from_rect(src_page_num: int, r: fitz.Rect):
        out_page_num, xoff, yoff = placement[src_page_num]
        return out_page_num, fitz.Rect(r.x0 + xoff, r.y0 + yoff, r.x1 + xoff, r.y1 + yoff)

    # 2) Copy links (annotations) from each source page to the correct output page
    for spn in range(src.page_count):
        links = src.load_page(spn).get_links()
        if not links:
            continue

        for lnk in links:
            # Map the clickable rectangle
            out_pn, new_from = map_from_rect(spn, lnk["from"])
            out_page = out.load_page(out_pn)

            kind = lnk.get("kind")

            # Common fields
            new_link = {
                "kind": kind,
                "from": new_from,
            }

            # External URL links
            if kind == fitz.LINK_URI:
                new_link["uri"] = lnk.get("uri", "")

            # Internal "go to page" links
            elif kind == fitz.LINK_GOTO:
                dest_src_page = lnk.get("page")
                if dest_src_page is None or dest_src_page < 0 or dest_src_page >= src.page_count:
                    continue  # skip broken destinations

                dest_out_page, new_to = map_dest(dest_src_page, lnk.get("to"))
                new_link["page"] = dest_out_page
                if new_to is not None:
                    new_link["to"] = new_to
                if "zoom" in lnk and lnk["zoom"] is not None:
                    new_link["zoom"] = lnk["zoom"]

            # Named destinations / file launches etc.
            # You can add more cases if you use them, but most PDFs are URI + GOTO.
            else:
                # Try a best-effort copy for other kinds (may or may not be supported)
                for k in ("page", "to", "zoom", "file", "uri", "name"):
                    if k in lnk and lnk[k] is not None:
                        new_link[k] = lnk[k]

            out_page.insert_link(new_link)

    out.save(output_pdf)
    out.close()
    src.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Places pages side by side")
    parser.add_argument("input_pdf", help="Path to the source PDF file.")
    parser.add_argument("output_pdf", help="Path to the destination PDF file.")
    parser.add_argument("start_page", help="Allows to skip n number of pages", type=int)
    args = parser.parse_args()
    side_by_side_preserve_links(args.input_pdf, args.output_pdf, args.start_page)