import fitz  # PyMuPDF
import os

# 文件路径
pdf_path = "彩色楼梯灯控制器产品规格书.pdf"
logo_path = "/root/.openclaw/media/inbound/162e2250-6104-4edf-89e4-017349c1dfc1.png"
output_path = "彩色楼梯灯控制器产品规格书_带logo.pdf"

# 打开 PDF
pdf_document = fitz.open(pdf_path)

# 打开 logo 图片
logo_pix = fitz.Pixmap(logo_path)

# 设置 logo 在页眉的显示尺寸（宽度100点，高度按比例缩放）
logo_width = 80
logo_height = logo_width * logo_pix.height / logo_pix.width

# 为每一页添加 logo
for page_num in range(len(pdf_document)):
    page = pdf_document[page_num]
    
    # 获取页面尺寸
    page_rect = page.rect
    
    # 计算 logo 位置：右上角，留一些边距
    margin_right = 30
    margin_top = 20
    
    # logo 放置位置（右上角）
    logo_x = page_rect.width - margin_right - logo_width
    logo_y = margin_top
    
    # 定义 logo 放置的矩形区域
    logo_rect = fitz.Rect(
        logo_x,
        logo_y,
        logo_x + logo_width,
        logo_y + logo_height
    )
    
    # 在页面上插入 logo
    page.insert_image(logo_rect, pixmap=logo_pix)

# 保存新 PDF
pdf_document.save(output_path)
pdf_document.close()

print(f"✅ 已成功将 logo 添加到页眉")
print(f"📄 输出文件: {output_path}")
print(f"📑 总页数: {len(fitz.open(pdf_path))}")
