PictureEdit/main_window.py

"""
作者：太一
微信：taiyi1224
邮箱：shuobo1224@qq.com
"""

import tkinter as tk
from tkinter import ttk, messagebox, filedialog, simpledialog
from pathlib import Path
from models import Project, Template
from canvas_widget import CanvasWidget
from template_dialog import TemplateDialog
from export_dialog import ExportDialog
from exporter import Exporter

# 确保中文显示正常
import matplotlib

matplotlib.use('Agg')
import matplotlib.pyplot as plt

plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]


class MainWindow(tk.Tk):
    """主窗口"""

    def __init__(self):
        super().__init__()
        self.title("批量图片布局工具")
        self.project = Project()
        self.init_ui()

    def init_ui(self):
        """初始化UI"""
        # 设置样式
        style = ttk.Style()
        style.configure("TButton", padding=5)
        style.configure("TLabel", padding=2)

        # 菜单栏
        self.create_menu()

        # 工具栏
        toolbar = ttk.Frame(self, padding=5)
        toolbar.pack(fill=tk.X)

        ttk.Button(toolbar, text="打开背景", command=self.open_background).pack(side=tk.LEFT, padx=2)
        ttk.Button(toolbar, text="添加图片", command=self.add_foregrounds).pack(side=tk.LEFT, padx=2)
        ttk.Button(toolbar, text="清除图片", command=self.clear_foregrounds).pack(side=tk.LEFT, padx=2)
        ttk.Button(toolbar, text="模板设置", command=self.set_template).pack(side=tk.LEFT, padx=2)
        ttk.Button(toolbar, text="批量导出", command=self.batch_export).pack(side=tk.LEFT, padx=2)

        # 主布局
        main_frame = ttk.Frame(self, padding=5)
        main_frame.pack(fill=tk.BOTH, expand=True)

        # 左侧：图片列表
        left_frame = ttk.Frame(main_frame, width=200)
        left_frame.pack(side=tk.LEFT, fill=tk.Y, padx=(0, 5))

        ttk.Label(left_frame, text="前景图片列表", font=("Arial", 10, "bold")).pack(anchor=tk.W, pady=5)

        self.image_list_frame = ttk.Frame(left_frame)
        self.image_list_frame.pack(fill=tk.BOTH, expand=True)

        # 滚动条
        scrollbar = ttk.Scrollbar(self.image_list_frame)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)

        self.image_listbox = tk.Listbox(
            self.image_list_frame,
            yscrollcommand=scrollbar.set,
            selectmode=tk.SINGLE,
            height=20
        )
        self.image_listbox.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        scrollbar.config(command=self.image_listbox.yview)

        # 绑定列表事件
        self.image_listbox.bind('<<ListboxSelect>>', self.on_image_select)

        # 添加按钮
        ttk.Button(left_frame, text="添加图片", command=self.add_foregrounds).pack(fill=tk.X, pady=5)
        ttk.Button(left_frame, text="移除选中", command=self.remove_selected).pack(fill=tk.X)

        # 中间：画布区域
        center_frame = ttk.Frame(main_frame)
        center_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        # 画布滚动条
        hscroll = ttk.Scrollbar(center_frame, orient=tk.HORIZONTAL)
        hscroll.pack(side=tk.BOTTOM, fill=tk.X)

        vscroll = ttk.Scrollbar(center_frame, orient=tk.VERTICAL)
        vscroll.pack(side=tk.RIGHT, fill=tk.Y)

        # 创建画布
        self.canvas = CanvasWidget(center_frame, self.project, main_window=self)
        # self.canvas = CanvasWidget(center_frame, self.project)
        self.canvas.pack(fill=tk.BOTH, expand=True)

        # 关联滚动条
        self.canvas.config(xscrollcommand=hscroll.set, yscrollcommand=vscroll.set)
        hscroll.config(command=self.canvas.xview)
        vscroll.config(command=self.canvas.yview)

        # 右侧：属性面板
        right_frame = ttk.Frame(main_frame, width=200)
        right_frame.pack(side=tk.RIGHT, fill=tk.Y, padx=(5, 0))

        ttk.Label(right_frame, text="属性设置", font=("Arial", 10, "bold")).pack(anchor=tk.W, pady=5)

        self.property_frame = ttk.Frame(right_frame)
        self.property_frame.pack(fill=tk.BOTH, expand=True, padx=5)

        # 初始化属性面板内容
        self.clear_property_panel()

        # 设置窗口大小
        self.geometry("1200x700")

    def create_menu(self):
        """创建菜单栏"""
        menubar = tk.Menu(self)

        # 文件菜单
        file_menu = tk.Menu(menubar, tearoff=0)
        file_menu.add_command(label="新建项目", command=self.new_project)
        file_menu.add_command(label="打开项目", command=self.open_project)
        file_menu.add_command(label="保存项目", command=self.save_project)
        file_menu.add_command(label="另存为", command=self.save_project_as)
        file_menu.add_separator()
        file_menu.add_command(label="退出", command=self.quit)
        menubar.add_cascade(label="文件", menu=file_menu)

        # 编辑菜单
        edit_menu = tk.Menu(menubar, tearoff=0)
        edit_menu.add_command(label="清除所有图片", command=self.clear_foregrounds)
        edit_menu.add_separator()
        edit_menu.add_command(label="上移图层", command=lambda: self.move_layer(1))
        edit_menu.add_command(label="下移图层", command=lambda: self.move_layer(-1))
        menubar.add_cascade(label="编辑", menu=edit_menu)

        # 视图菜单
        view_menu = tk.Menu(menubar, tearoff=0)
        view_menu.add_command(label="缩放重置", command=self.reset_zoom)
        menubar.add_cascade(label="视图", menu=view_menu)

        # 帮助菜单
        help_menu = tk.Menu(menubar, tearoff=0)
        help_menu.add_command(label="关于", command=self.show_about)
        menubar.add_cascade(label="帮助", menu=help_menu)

        # 设置菜单栏
        self.config(menu=menubar)

    def open_background(self):
        """打开背景图"""
        file_path = filedialog.askopenfilename(
            filetypes=[
                ("图像文件", "*.jpg;*.jpeg;*.png;*.webp;*.bmp"),
                ("所有文件", "*.*")
            ]
        )
        if file_path:
            if self.project.template.load_background(file_path):
                self.canvas.draw_preview()

    def add_foregrounds(self):
        """添加前景图"""
        file_paths = filedialog.askopenfilenames(
            filetypes=[
                ("图像文件", "*.jpg;*.jpeg;*.png;*.webp;*.bmp"),
                ("所有文件", "*.*")
            ]
        )

        if file_paths:
            # 检查是否是第一次添加图片
            is_first_batch = len(self.project.foregrounds) == 0

            # 添加所有选择的图片
            added_foregrounds = []
            for file_path in file_paths:
                # 检查是否已添加
                if not any(fg.file_path == file_path for fg in self.project.foregrounds):
                    fg = self.project.add_foreground(file_path)
                    added_foregrounds.append(fg)

            # 如果是第一次批量添加，则设置批量处理模式
            if is_first_batch and len(added_foregrounds) > 1:
                # 第一张图片保持可见和可编辑
                first_fg = added_foregrounds[0]
                first_fg.hidden = False
                first_fg.visible = True

                # 其他图片设置为隐藏
                for fg in added_foregrounds[1:]:
                    fg.hidden = True
                    fg.visible = False

                # 选中第一张图片
                self.canvas.selected_item = first_fg
                self.update_property_panel(first_fg)
            elif not is_first_batch and self.project.foregrounds:
                # 如果不是第一次添加，且已存在图片，则复制第一张图片的配置
                reference_fg = self.project.foregrounds[0]
                for fg in added_foregrounds:
                    fg.vertices = [tuple(v) for v in reference_fg.vertices]
                    fg.angle = reference_fg.angle
                    fg.mask_shape = reference_fg.mask_shape
                    fg.locked = reference_fg.locked
                    fg.visible = reference_fg.visible
                    fg.hidden = reference_fg.hidden

            # 更新列表
            self.update_image_list()
            # 重绘画布
            self.canvas.draw_preview()

    def update_image_list(self):
        """更新图片列表"""
        self.image_listbox.delete(0, tk.END)
        for fg in self.project.foregrounds:
            name = Path(fg.file_path).name
            self.image_listbox.insert(tk.END, name)

    def on_image_select(self, event):
        """处理列表选择事件"""
        selection = self.image_listbox.curselection()
        if selection:
            index = selection[0]
            if 0 <= index < len(self.project.foregrounds):
                fg = self.project.foregrounds[index]
                self.canvas.selected_item = fg
                self.update_property_panel(fg)
                self.canvas.draw_preview()
        else:
            # 如果没有选择任何项目，显示默认面板
            self.clear_property_panel()

    def remove_selected(self):
        """移除选中的图片"""
        selection = self.image_listbox.curselection()
        if selection:
            index = selection[0]
            if 0 <= index < len(self.project.foregrounds):
                fg = self.project.foregrounds[index]
                self.project.remove_foreground(fg)
                self.update_image_list()
                self.clear_property_panel()
                self.canvas.selected_item = None
                self.canvas.draw_preview()

    def clear_foregrounds(self):
        """清除所有前景图"""
        if messagebox.askyesno("确认", "确定要清除所有前景图片吗？"):
            self.project.foregrounds.clear()
            self.update_image_list()
            self.clear_property_panel()
            self.canvas.selected_item = None
            self.canvas.draw_preview()

    def get_canvas(self):
        """获取画布组件"""
        return self.canvas

    def set_template(self):
        """设置模板"""
        dialog = TemplateDialog(self, self.project.template)
        if dialog.result:
            # 处理所有类型的模板设置结果
            if dialog.result["type"] == "preset_loaded":
                # 直接更新现有模板的属性
                self.project.template.ratio = dialog.result["ratio"]
                self.project.template.width_px = dialog.result["width"]
                self.project.template.height_px = dialog.result["height"]
                self.project.template.bg_color = dialog.result["bg_color"]
                if "bg_image_path" in dialog.result and dialog.result["bg_image_path"]:
                    self.project.template.load_background(dialog.result["bg_image_path"])
            elif dialog.result["type"] == "preset":
                # 更新现有模板的属性
                self.project.template.ratio = tuple(dialog.result["ratio"])
                self.project.template.width_px = dialog.result["width"]
                self.project.template.height_px = dialog.result["height"]
                self.project.template.bg_color = dialog.result["bg_color"]
                # 设置背景图
                if "bg_image_path" in dialog.result and dialog.result["bg_image_path"]:
                    self.project.template.load_background(dialog.result["bg_image_path"])
                elif "bg_image" in dialog.result:
                    self.project.template.bg_image = dialog.result["bg_image"]
            elif dialog.result["type"] == "custom":
                # 更新现有模板的属性
                self.project.template.set_custom_size(
                    dialog.result["width"],
                    dialog.result["height"]
                )
                # 设置ratio属性
                if "ratio" in dialog.result:
                    self.project.template.ratio = dialog.result["ratio"]
                self.project.template.bg_color = dialog.result["bg_color"]
                # 设置背景图
                if "bg_image_path" in dialog.result and dialog.result["bg_image_path"]:
                    self.project.template.load_background(dialog.result["bg_image_path"])
                elif "bg_image" in dialog.result:
                    self.project.template.bg_image = dialog.result["bg_image"]

            # 确保更新画布显示
            self.canvas.draw_preview()

    def batch_export(self):
        """批量导出"""
        if not self.project.foregrounds:
            messagebox.showinfo("提示", "请先添加前景图片")
            return

        # 在导出前，将第一张图片的配置应用到所有隐藏的图片
        if self.project.foregrounds:
            first_fg = self.project.foregrounds[0]
            # 将第一张图片的配置应用到所有隐藏的图片
            for fg in self.project.foregrounds[1:]:
                if fg.hidden:
                    fg.vertices = [tuple(v) for v in first_fg.vertices]
                    fg.angle = first_fg.angle
                    fg.mask_shape = first_fg.mask_shape

        # 设置默认导出路径为第一个图片所在目录下的output文件夹
        default_output_dir = None
        if self.project.foregrounds:
            first_image_dir = Path(self.project.foregrounds[0].file_path).parent
            default_output_dir = first_image_dir / "output"

        dialog = ExportDialog(self, len(self.project.foregrounds))
        if dialog.result:
            # 禁用主窗口防止重复操作
            self.grab_set()
            try:
                # 执行导出
                success = Exporter.export_batch(
                    self.project,
                    dialog.output_dir_var.get(),
                    dialog.filename_var.get(),
                    dialog.format_var.get(),
                    dialog.quality_var.get(),
                    progress_callback=self.update_export_progress
                )

                # 导出完成
                if success:
                    messagebox.showinfo("成功", f"已成功导出 {len(self.project.foregrounds)} 张图片")
                else:
                    messagebox.showerror("错误", "导出过程中发生错误")
            finally:
                # 重新启用主窗口
                self.grab_release()

    def update_export_progress(self, current, total):
        """更新导出进度（这个方法需要在ExportDialog中实现）"""
        # 注意：这个方法需要在ExportDialog中调用，而不是直接在这里实现
        pass

    def update_property_panel(self, fg):
        """更新属性面板"""
        # 清除现有内容
        for widget in self.property_frame.winfo_children():
            widget.destroy()

        # 位置设置
        ttk.Label(self.property_frame, text="位置", font=("Arial", 9, "bold")).pack(anchor=tk.W, pady=5)

        pos_frame = ttk.Frame(self.property_frame)
        pos_frame.pack(fill=tk.X, pady=2)

        ttk.Label(pos_frame, text="X:").grid(row=0, column=0, sticky=tk.W)
        x_var = tk.DoubleVar(value=fg.x)
        x_entry = ttk.Entry(pos_frame, textvariable=x_var, width=10)
        x_entry.grid(row=0, column=1, padx=2)
        x_entry.bind("<FocusOut>", lambda e: self.update_property(fg, "x", x_var.get()))

        ttk.Label(pos_frame, text="Y:").grid(row=0, column=2, sticky=tk.W, padx=5)
        y_var = tk.DoubleVar(value=fg.y)
        y_entry = ttk.Entry(pos_frame, textvariable=y_var, width=10)
        y_entry.grid(row=0, column=3, padx=2)
        y_entry.bind("<FocusOut>", lambda e: self.update_property(fg, "y", y_var.get()))

        # 尺寸设置
        ttk.Label(self.property_frame, text="尺寸", font=("Arial", 9, "bold")).pack(anchor=tk.W, pady=5)

        size_frame = ttk.Frame(self.property_frame)
        size_frame.pack(fill=tk.X, pady=2)

        ttk.Label(size_frame, text="宽:").grid(row=0, column=0, sticky=tk.W)
        w_var = tk.DoubleVar(value=fg.w)
        w_entry = ttk.Entry(size_frame, textvariable=w_var, width=10)
        w_entry.grid(row=0, column=1, padx=2)
        w_entry.bind("<FocusOut>", lambda e: self.update_property(fg, "w", w_var.get()))

        ttk.Label(size_frame, text="高:").grid(row=0, column=2, sticky=tk.W, padx=5)
        h_var = tk.DoubleVar(value=fg.h)
        h_entry = ttk.Entry(size_frame, textvariable=h_var, width=10)
        h_entry.grid(row=0, column=3, padx=2)
        h_entry.bind("<FocusOut>", lambda e: self.update_property(fg, "h", h_var.get()))

        # 旋转设置
        ttk.Label(self.property_frame, text="旋转", font=("Arial", 9, "bold")).pack(anchor=tk.W, pady=5)

        rot_frame = ttk.Frame(self.property_frame)
        rot_frame.pack(fill=tk.X, pady=2)

        angle_var = tk.DoubleVar(value=fg.angle)
        angle_entry = ttk.Entry(rot_frame, textvariable=angle_var, width=10)
        angle_entry.grid(row=0, column=0, padx=2)
        ttk.Label(rot_frame, text="度").grid(row=0, column=1, sticky=tk.W)
        angle_entry.bind("<FocusOut>", lambda e: self.update_property(fg, "angle", angle_var.get()))

        # 遮罩设置
        ttk.Label(self.property_frame, text="遮罩", font=("Arial", 9, "bold")).pack(anchor=tk.W, pady=5)

        mask_frame = ttk.Frame(self.property_frame)
        mask_frame.pack(fill=tk.X, pady=2)

        mask_var = tk.StringVar(value=fg.mask_shape)
        masks = [("矩形", "rect"), ("圆形", "round"), ("心形", "heart"), ("星形", "star")]
        for i, (text, value) in enumerate(masks):
            ttk.Radiobutton(
                mask_frame,
                text=text,
                variable=mask_var,
                value=value,
                command=lambda: self.update_property(fg, "mask_shape", mask_var.get())
            ).grid(row=i // 2, column=i % 2, sticky=tk.W, padx=2, pady=2)

        # 其他设置
        ttk.Label(self.property_frame, text="其他", font=("Arial", 9, "bold")).pack(anchor=tk.W, pady=5)

        lock_var = tk.BooleanVar(value=fg.locked)
        ttk.Checkbutton(
            self.property_frame,
            text="锁定位置",
            variable=lock_var,
            command=lambda: self.update_property(fg, "locked", lock_var.get())
        ).pack(anchor=tk.W, pady=2)

        visible_var = tk.BooleanVar(value=fg.visible)
        ttk.Checkbutton(
            self.property_frame,
            text="显示图片",
            variable=visible_var,
            command=lambda: self.update_property(fg, "visible", visible_var.get())
        ).pack(anchor=tk.W, pady=2)

    def update_property(self, fg, prop, value):
        """更新属性值"""
        if hasattr(fg, prop):
            setattr(fg, prop, value)
            self.canvas.draw_preview()

    def clear_property_panel(self):
        """清除属性面板"""
        for widget in self.property_frame.winfo_children():
            widget.destroy()

        ttk.Label(self.property_frame, text="未选择图片", foreground="#999").pack(pady=20)

    def move_layer(self, direction):
        """移动图层顺序"""
        if not self.canvas.selected_item:
            return

        fg = self.canvas.selected_item
        current_idx = fg.z_index
        new_idx = current_idx + direction

        # 检查边界
        if new_idx < 0 or new_idx >= len(self.project.foregrounds):
            return

        # 交换z_index
        for other in self.project.foregrounds:
            if other.z_index == new_idx:
                other.z_index = current_idx
                break

        fg.z_index = new_idx
        self.canvas.draw_preview()

    def reset_zoom(self):
        """重置缩放"""
        self.project.canvas_zoom = 1.0
        self.canvas.zoom = 1.0
        self.canvas.draw_preview()

    def new_project(self):
        """新建项目"""
        if messagebox.askyesno("确认", "确定要创建新项目吗？当前项目将被清空。"):
            self.project = Project()
            self.canvas.project = self.project
            self.update_image_list()
            self.clear_property_panel()
            self.canvas.selected_item = None
            self.canvas.draw_preview()

    def open_project(self):
        """打开项目"""
        file_path = filedialog.askopenfilename(
            filetypes=[("项目文件", "*.json"), ("所有文件", "*.*")]
        )

        if file_path:
            new_project = Project()
            if new_project.load(file_path):
                self.project = new_project
                self.canvas.project = self.project
                self.canvas.zoom = self.project.canvas_zoom
                self.update_image_list()
                self.clear_property_panel()
                self.canvas.selected_item = None
                self.canvas.draw_preview()
            else:
                messagebox.showerror("错误", "打开项目失败")

    def save_project(self):
        """保存项目"""
        if self.project.file_path:
            if self.project.save():
                messagebox.showinfo("成功", "项目保存成功")
            else:
                messagebox.showerror("错误", "项目保存失败")
        else:
            self.save_project_as()

    def save_project_as(self):
        """另存为项目"""
        file_path = filedialog.asksaveasfilename(
            defaultextension=".json",
            filetypes=[("项目文件", "*.json"), ("所有文件", "*.*")]
        )

        if file_path:
            if self.project.save(file_path):
                messagebox.showinfo("成功", "项目保存成功")
            else:
                messagebox.showerror("错误", "项目保存失败")

    def show_about(self):
        """显示关于对话框"""
        messagebox.showinfo(
            "关于",
            "批量图片布局工具\n"
            "版本: 1.0\n"
            "用于批量处理图片布局的工具，可以将多张图片按照相同的布局应用到指定模板中。"
        )