import os
import json
import threading
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
from queue import Queue
from PIL import Image, ImageTk
import numpy as np


class ImageProcessorApp:
    def __init__(self, root):
        self.root = root
        self.root.title("智能图片合成器 (带断点续传)")
        self.root.geometry("800x600")

        # 初始化状态
        self.is_running = False
        self.pause_flag = False
        self.status_file = "progress_status.json"
        self.task_queue = Queue()
        self.current_task = {"bg_index": 0, "input_index": 0}

        # 边距设置（可配置）
        self.margins = self.load_margins() or {"top": 76, "bottom": 76, "left": 57, "right": 57}

        # 创建界面
        self.create_widgets()
        self.load_progress()

    def create_widgets(self):
        # 控制面板
        control_frame = ttk.Frame(self.root)
        control_frame.pack(fill=tk.X, padx=10, pady=5)

        self.start_btn = ttk.Button(control_frame, text="开始", command=self.start_processing)
        self.pause_btn = ttk.Button(control_frame, text="暂停", state=tk.DISABLED, command=self.toggle_pause)
        self.reset_btn = ttk.Button(control_frame, text="重置", command=self.reset_progress)
        self.start_btn.pack(side=tk.LEFT, padx=5)
        self.pause_btn.pack(side=tk.LEFT, padx=5)
        self.reset_btn.pack(side=tk.LEFT, padx=5)

        # 文件夹选择
        self.setup_folder_controls()

        # 边距配置
        self.setup_margin_controls()

        # 进度显示
        self.progress = ttk.Progressbar(self.root, orient=tk.HORIZONTAL)
        self.progress.pack(fill=tk.X, padx=10, pady=5)
        self.status_label = ttk.Label(self.root, text="就绪")
        self.status_label.pack()

        # 背景图片尺寸显示
        self.bg_dimensions_label = ttk.Label(self.root, text="背景尺寸: 未加载")
        self.bg_dimensions_label.pack()

    def setup_folder_controls(self):
        frame = ttk.LabelFrame(self.root, text="文件夹设置")
        frame.pack(fill=tk.X, padx=10, pady=5)

        # 输入文件夹
        ttk.Label(frame, text="输入:").grid(row=0, column=0, padx=5)
        self.input_dir = ttk.Entry(frame, width=40)
        self.input_dir.grid(row=0, column=1, padx=5)
        ttk.Button(frame, text="浏览", command=lambda: self.select_dir(self.input_dir)).grid(row=0, column=2)

        # 背景文件夹
        ttk.Label(frame, text="背景:").grid(row=1, column=0, padx=5)
        self.bg_dir = ttk.Entry(frame, width=40)
        self.bg_dir.grid(row=1, column=1, padx=5)
        ttk.Button(frame, text="浏览", command=lambda: self.select_dir(self.bg_dir)).grid(row=1, column=2)

        # 输出文件夹
        ttk.Label(frame, text="输出:").grid(row=2, column=0, padx=5)
        self.output_dir = ttk.Entry(frame, width=40)
        self.output_dir.grid(row=2, column=1, padx=5)
        ttk.Button(frame, text="浏览", command=lambda: self.select_dir(self.output_dir)).grid(row=2, column=2)

    def setup_margin_controls(self):
        frame = ttk.LabelFrame(self.root, text="边距设置")
        frame.pack(fill=tk.X, padx=10, pady=5)

        ttk.Label(frame, text="上:").grid(row=0, column=0, padx=5)
        self.top_margin = ttk.Entry(frame, width=10)
        self.top_margin.insert(0, str(self.margins["top"]))
        self.top_margin.grid(row=0, column=1, padx=5)

        ttk.Label(frame, text="下:").grid(row=0, column=2, padx=5)
        self.bottom_margin = ttk.Entry(frame, width=10)
        self.bottom_margin.insert(0, str(self.margins["bottom"]))
        self.bottom_margin.grid(row=0, column=3, padx=5)

        ttk.Label(frame, text="左:").grid(row=1, column=0, padx=5)
        self.left_margin = ttk.Entry(frame, width=10)
        self.left_margin.insert(0, str(self.margins["left"]))
        self.left_margin.grid(row=1, column=1, padx=5)

        ttk.Label(frame, text="右:").grid(row=1, column=2, padx=5)
        self.right_margin = ttk.Entry(frame, width=10)
        self.right_margin.insert(0, str(self.margins["right"]))
        self.right_margin.grid(row=1, column=3, padx=5)

        ttk.Button(frame, text="应用", command=self.update_margins).grid(row=1, column=4, padx=5)

    def update_margins(self):
        try:
            self.margins = {
                "top": int(self.top_margin.get()),
                "bottom": int(self.bottom_margin.get()),
                "left": int(self.left_margin.get()),
                "right": int(self.right_margin.get()),
            }
            messagebox.showinfo("成功", "边距已更新！")
        except ValueError:
            messagebox.showerror("错误", "请输入有效的数字！")

    def load_margins(self):
        if os.path.exists("margins.json"):
            try:
                with open("margins.json", "r") as f:
                    return json.load(f)
            except:
                pass
        return None

    def save_margins(self):
        with open("margins.json", "w") as f:
            json.dump(self.margins, f)

    def select_dir(self, entry):
        path = filedialog.askdirectory()
        if path:
            entry.delete(0, tk.END)
            entry.insert(0, path)

    def load_progress(self):
        if os.path.exists(self.status_file):
            try:
                with open(self.status_file, 'r') as f:
                    data = json.load(f)
                    self.current_task = data.get('current_task', self.current_task)
                    self.input_dir.insert(0, data.get('input_folder', ''))
                    self.bg_dir.insert(0, data.get('bg_folder', ''))
                    self.output_dir.insert(0, data.get('output_folder', ''))
                    self.progress["value"] = data.get('progress', 0)
            except:
                pass

    def save_progress(self):
        data = {
            "input_folder": self.input_dir.get(),
            "bg_folder": self.bg_dir.get(),
            "output_folder": self.output_dir.get(),
            "current_task": self.current_task,
            "progress": self.progress["value"]
        }
        with open(self.status_file, 'w') as f:
            json.dump(data, f)

    def start_processing(self):
        if not self.validate_paths():
            return

        self.is_running = True
        self.pause_flag = False
        self.start_btn.config(state=tk.DISABLED)
        self.pause_btn.config(state=tk.NORMAL)

        # 初始化任务队列
        bg_images = self.get_image_files(self.bg_dir.get())
        input_images = self.get_image_files(self.input_dir.get())
        total = len(bg_images) * len(input_images)
        self.progress["maximum"] = total

        # 创建处理线程
        self.worker_thread = threading.Thread(
            target=self.process_images,
            args=(bg_images, input_images),
            daemon=True
        )
        self.worker_thread.start()
        self.monitor_thread()

    def toggle_pause(self):
        self.pause_flag = not self.pause_flag
        self.pause_btn.config(text="继续" if self.pause_flag else "暂停")
        self.save_progress()

    def reset_progress(self):
        self.progress["value"] = 0
        self.current_task = {"bg_index": 0, "input_index": 0}
        if os.path.exists(self.status_file):
            os.remove(self.status_file)

    def validate_paths(self):
        required = [
            (self.input_dir.get(), "请输入输入文件夹"),
            (self.bg_dir.get(), "请输入背景文件夹"),
            (self.output_dir.get(), "请指定输出文件夹")
        ]
        for path, msg in required:
            if not path:
                messagebox.showerror("错误", msg)
                return False
            if not os.path.exists(path):
                messagebox.showerror("错误", f"路径不存在: {path}")
                return False
        return True

    def monitor_thread(self):
        if self.worker_thread.is_alive():
            self.root.after(100, self.monitor_thread)
        else:
            self.start_btn.config(state=tk.NORMAL)
            self.pause_btn.config(state=tk.DISABLED)
            self.is_running = False
            self.save_progress()

    def process_images(self, bg_images, input_images):
        bg_path = self.bg_dir.get()
        input_path = self.input_dir.get()
        output_path = self.output_dir.get()

        for bg_idx in range(self.current_task["bg_index"], len(bg_images)):
            if not self.is_running or self.pause_flag:
                break

            bg_name = bg_images[bg_idx]
            try:
                with Image.open(os.path.join(bg_path, bg_name)) as img:
                    bg_img = self.process_background(img)
                    # 更新背景图片尺寸显示
                    self.root.after(0, lambda: self.bg_dimensions_label.config(text=f"背景尺寸: {bg_img.size[0]}x{bg_img.size[1]}"))
            except Exception as e:
                continue

            target_size = self.calculate_target_size(bg_img.size)
            subfolder = self.create_subfolder(output_path, bg_name)

            for input_idx in range(self.current_task["input_index"], len(input_images)):
                if not self.is_running or self.pause_flag:
                    self.current_task.update({"bg_index": bg_idx, "input_index": input_idx})
                    break

                input_name = input_images[input_idx]
                try:
                    self.process_single_image(
                        os.path.join(input_path, input_name),
                        bg_img,
                        target_size,
                        os.path.join(subfolder, f"composite_{input_name}")
                    )
                    self.progress["value"] += 1
                    self.status_label.config(text=f"处理中: {self.progress['value']}/{self.progress['maximum']}")
                except Exception as e:
                    continue

                self.current_task["input_index"] = input_idx + 1
                self.save_progress()

            self.current_task.update({"bg_index": bg_idx + 1, "input_index": 0})
            self.save_progress()

    def process_background(self, img):
        processed = img.copy().convert("RGBA")
        data = np.array(processed)
        r, g, b, a = data.T
        white_mask = (r > 240) & (g > 240) & (b > 240)
        data[..., :][white_mask.T] = (255, 255, 255, 0)
        return Image.fromarray(data)

    def calculate_target_size(self, bg_size):
        return (
            bg_size[0] - self.margins["left"] - self.margins["right"],
            bg_size[1] - self.margins["top"] - self.margins["bottom"]
        )

    def create_subfolder(self, output_path, bg_name):
        subfolder = os.path.join(output_path, os.path.splitext(bg_name)[0])
        os.makedirs(subfolder, exist_ok=True)
        return subfolder

    def process_single_image(self, input_path, bg_img, target_size, output_path):
        with Image.open(input_path) as img:
            input_img = img.convert('RGBA')

        resized = self.adaptive_resize(input_img, target_size)
        composite = Image.new('RGBA', bg_img.size)
        composite.paste(bg_img, (0, 0))
        composite.paste(resized, (self.margins["left"], self.margins["top"]), mask=resized)
        composite.save(output_path)

    @staticmethod
    def adaptive_resize(img, target_size):
        width, height = img.size
        target_w, target_h = target_size

        ratio = min(target_w / width, target_h / height)
        new_size = (int(width * ratio), int(height * ratio))

        resized = img.resize(new_size, Image.LANCZOS)
        canvas = Image.new("RGBA", target_size, (0, 0, 0, 0))
        canvas.paste(resized, (
            (target_w - new_size[0]) // 2,
            (target_h - new_size[1]) // 2
        ))
        return canvas

    @staticmethod
    def get_image_files(folder):
        return [f for f in os.listdir(folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.tif'))]


if __name__ == "__main__":
    root = tk.Tk()
    app = ImageProcessorApp(root)
    root.mainloop()