import os, cv2, zipfile, gc, gradio as gr
import numpy as np
import fitz  # PyMuPDF
from PIL import Image

# 設定
OUTPUT_DIR = "output"
ZIP_PATH = os.path.join(OUTPUT_DIR, "results.zip")
BG_PATH = "background.jpg"
TARGET_WIDTH, TARGET_HEIGHT = 2133, 1200
POS1 = (1032, 0)
POS2 = (4666, 0)

def ensure_dir(path):
    os.makedirs(path, exist_ok=True)
    if not os.access(path, os.W_OK):
        raise RuntimeError(f"❌ 無法寫入：{path}")
ensure_dir(OUTPUT_DIR)

# 解析頁碼範圍字串（如 "1-3,5,7"）
def parse_page_selection(selection_text):
    if not selection_text.strip():
        return None
    pages = set()
    for part in selection_text.split(','):
        part = part.strip()
        if '-' in part:
            start, end = map(int, part.split('-'))
            pages.update(range(start, end + 1))
        elif part.isdigit():
            pages.add(int(part))
    return sorted(pages)

# PDF轉圖片（np.array 格式，BGR）
def pdf_to_images(pdf_file, page_list=None):
    images = []
    with fitz.open(pdf_file.name) as doc:
        total_pages = len(doc)
        selected = page_list or list(range(1, total_pages + 1))
        for page_num in selected:
            if 1 <= page_num <= total_pages:
                pix = doc[page_num - 1].get_pixmap(matrix=fitz.Matrix(2, 2), colorspace=fitz.csRGB)
                img = np.frombuffer(pix.samples, dtype=np.uint8).reshape((pix.height, pix.width, 3))
                img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
                images.append(img)
    return images

# 合成圖片
def overlay_image(bg, fg_img, x, y):
    fg = cv2.resize(fg_img, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_AREA)
    bg[y:y+TARGET_HEIGHT, x:x+TARGET_WIDTH] = fg
    return bg

# 處理與合成流程
def process_and_combine(images, prefix):
    bg = cv2.imread(BG_PATH)
    if bg is None:
        raise gr.Error("❌ 無法讀取 background.jpg")

    for f in os.listdir(OUTPUT_DIR):
        if f.endswith(".jpg"):
            os.remove(os.path.join(OUTPUT_DIR, f))

    results = []
    for i, img in enumerate(images, 1):
        combined = overlay_image(bg.copy(), img, *POS1)
        combined = overlay_image(combined, img, *POS2)
        output_path = os.path.join(OUTPUT_DIR, f"{prefix}{i:03d}.jpg")
        cv2.imwrite(output_path, combined)
        results.append((output_path, f"{prefix}{i:03d}"))
        if i % 10 == 0:
            gc.collect()

    if not results:
        raise gr.Error("⚠️ 沒有產生圖片")

    with zipfile.ZipFile(ZIP_PATH, "w", zipfile.ZIP_DEFLATED) as zipf:
        for f in sorted(os.listdir(OUTPUT_DIR)):
            if f.endswith(".jpg"):
                zipf.write(os.path.join(OUTPUT_DIR, f), arcname=f)

    return results, ZIP_PATH, f"✅ 成功合成 {len(results)} 張圖，已打包 ZIP"

# 主邏輯：上傳處理
def process_inputs(files, prefix, page_selection):
    prefix = prefix.strip() or "result_"
    selected_pages = parse_page_selection(page_selection)

    images = []
    for file in files:
        name = file.name.lower()
        if name.endswith(".pdf"):
            images.extend(pdf_to_images(file, selected_pages))
        elif name.endswith((".png", ".jpg", ".jpeg")):
            img = Image.open(file).convert("RGB")
            img = np.array(img)
            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
            images.append(img)
        else:
            raise gr.Error(f"不支援的檔案格式：{name}")

    if not images:
        raise gr.Error("⚠️ 沒有有效圖片可處理，請確認頁碼或檔案內容")

    return process_and_combine(images, prefix)

# Gradio UI 元件定義
page_selection = gr.Textbox(label="選擇頁碼（例如：1,3,5-7）", placeholder="留空代表全部頁面")

# 介面建構
demo = gr.Interface(
    fn=process_inputs,
    inputs=[
        gr.Files(label="📎 上傳 PDF 或圖片", file_types=[".pdf", ".png", ".jpg", ".jpeg"]),
        gr.Textbox(label="檔名前綴（可選）", placeholder="例如：poster_"),
        page_selection
    ],
    outputs=[
        gr.Gallery(label="🖼️ 預覽", columns=3, height="auto"),
        gr.File(label="⬇️ 下載 ZIP"),
        gr.Textbox(label="📄 處理日誌")
    ],
    title="PDF 圖片合成工具",
    description="📄🖼️ 將 PDF 或圖片套用固定背景合成為寬圖，背景為 7872x1200，支援 PDF 選擇頁數與多檔處理。"
)

if __name__ == "__main__":
    demo.launch()