import tkinter as tk
from tkinter import filedialog, messagebox
from predict import separate as spt


class Application(tk.Frame):
    def __init__(self, root):
        super().__init__(root)

        self.master = root  # 主窗口程序
        self.pack()

        self.wav_root = tk.StringVar()
        self.save_root = tk.StringVar()

        self.master.minsize(500, 300)
        self.master.title('爆破小队')

        # 在窗口中加入控件
        self.add_wind()
    
    def choose_wav(self):
        """ 选择需要分离的音频文件 """
        filetypes = [("音频文件", "*.wav")]
        file_name = filedialog.askopenfilename(
            title="选择混合音频文件",
            filetypes=filetypes,
            initialdir="./"
        )
        self.wav_root.set(file_name)
    
    def choose_save_root(self):
        """ 选择分离音频保存路径 """
        save_root = filedialog.askdirectory(title="选择分离音频存储路径", initialdir="./")
        self.save_root.set(save_root)
    
    def separate(self):
        """ 执行分离程序 """
        wav_root = self.wav_root.get()
        save_root = self.save_root.get()

        if len(wav_root) == 0:
            messagebox.showinfo('提示: ', 'wav 文件路径未指定')
        elif len(save_root) == 0:
            messagebox.showinfo('提示: ', '文件输出路径未指定')
        else:
            spt(wav_root, save_root)
            messagebox.showinfo('提示: ', '音频分离完成')

    def add_wind(self):
        """ 在窗口中加入控件
        """

        frame1 = tk.Frame(self)
        tk.Label(frame1, text="爆破小队——音频多人声分离", bg='#61afef', width=70, height=3).grid()

        frame2 = tk.Frame(self)
        tk.Entry(frame2, textvariable=self.wav_root, width=50, state='readonly').grid(
            row=0, column=0, ipadx=2, ipady=2,
            padx=10, pady=30
        )
        tk.Button(frame2, text="选择混合音频文件", command=self.choose_wav, bg='#95da72').grid(row=0, column=1)

        frame3 = tk.Frame(self)
        tk.Entry(frame3, textvariable=self.save_root, width=50, state='readonly').grid(
            row=0, column=0, ipadx=2, ipady=2, padx=10
        )
        tk.Button(frame3, text="选择输出路径", command=self.choose_save_root, bg='#95da72').grid(
            row=0, column=1, ipadx=12
        )

        frame4 = tk.Frame(self)
        tk.Button(frame4, text="分离音频", command=self.separate, bg='#ffcf49').grid(
            ipadx=40, ipady=10
        )

        frame1.grid()
        frame2.grid()
        frame3.grid()
        frame4.grid(
            pady=40
        )


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