import cv2
import numpy as np
import os
import gradio as gr
import gc
import zipfile

# 常量設定
MAX_FILES = 250
BG_PATH = "background.jpg"  # 預設背景圖，請自行放置
FG_DIR = "foregrounds"
OUTPUT_DIR = "output"
ZIP_PATH = "output/results.zip"
TARGET_WIDTH, TARGET_HEIGHT = 2133, 1200
POS1 = (1032, 0)
POS2 = (4666, 0)

# 初始化資料夾與權限檢查
def check_write_permission(path):
    if not os.path.exists(path):
        try:
            os.makedirs(path)
        except Exception as e:
            raise RuntimeError(f"❌ 無法建立目錄 '{path}'：{e}")
    elif not os.access(path, os.W_OK):
        raise RuntimeError(f"❌ 目錄 '{path}' 存在，但沒有寫入權限！")

check_write_permission(FG_DIR)
check_write_permission(OUTPUT_DIR)

# 合成圖片
def overlay_image(bg, fg, x, y):
    h, w = fg.shape[:2]
    bh, bw = bg.shape[:2]
    if x + w > bw or y + h > bh:
        raise ValueError(f"前景圖超出背景邊界：({x},{y},{w},{h}) > 背景({bw},{bh})")

    if fg.shape[2] == 4:
        alpha = fg[:, :, 3:] / 255.0
        bg[y:y+h, x:x+w] = alpha * fg[:, :, :3] + (1 - alpha) * bg[y:y+h, x:x+w]
    else:
        bg[y:y+h, x:x+w] = fg[:, :, :3]
    return bg

# 壓縮成 zip
def zip_output_files():
    with zipfile.ZipFile(ZIP_PATH, "w", zipfile.ZIP_DEFLATED) as zipf:
        for file in sorted(os.listdir(OUTPUT_DIR)):
            if file.endswith(".jpg"):
                zipf.write(os.path.join(OUTPUT_DIR, file), arcname=file)
    return ZIP_PATH

# 處理合成流程
def process_images(fg_files):
    try:
        progress = gr.Progress()
        progress(0, desc="準備中...")

        if len(fg_files) == 0:
            return [], None, "請至少上傳一張前景圖"
        if len(fg_files) > MAX_FILES:
            raise gr.Error(f"最多支援 {MAX_FILES} 張圖片")

        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, fg_file in enumerate(fg_files, 1):
            progress(i / len(fg_files), f"合成第 {i} 張圖...")

            try:
                with open(fg_file.name, 'rb') as f:
                    fg_data = np.frombuffer(f.read(), np.uint8)
                fg = cv2.imdecode(fg_data, cv2.IMREAD_UNCHANGED)
                if fg is None:
                    continue

                fg = cv2.resize(fg, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_AREA)

                result = overlay_image(bg.copy(), fg, *POS1)
                result = overlay_image(result, fg, *POS2)

                output_path = f"{OUTPUT_DIR}/result_{i:03d}.jpg"
                cv2.imwrite(output_path, result)
                results.append((output_path, f"Result {i:03d}"))

                if i % 10 == 0:
                    gc.collect()

            except Exception as e:
                print(f"處理 {fg_file.name} 失敗: {str(e)}")
                continue

        progress(1, "完成！")
        zip_file = zip_output_files()
        return results, zip_file, f"成功處理 {len(results)}/{len(fg_files)} 張圖片，已打包為 ZIP 可下載"

    except Exception as e:
        raise gr.Error(f"處理出錯：{str(e)}")

# Gradio 介面
with gr.Blocks() as demo:
    gr.Markdown("## 🖼️ 寬幅圖片合成工具（背景固定為 7872x1200）")

    fg_upload = gr.Files(label="上傳前景圖（JPG 或 PNG，自動套用背景）", file_types=[".jpg", ".png"])
    btn_run = gr.Button("開始合成", variant="primary")

    with gr.Row():
        gallery = gr.Gallery(label="合成結果預覽", columns=3, preview=True)
        zip_download = gr.File(label="下載結果 ZIP 檔")
        log_output = gr.Textbox(label="日誌", interactive=False)

    btn_run.click(
        fn=process_images,
        inputs=fg_upload,
        outputs=[gallery, zip_download, log_output],
        concurrency_limit=4
    )

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