import os
import shutil
import pandas as pd
from typing import Dict
import pyzipper

class DataSaver:
    def __init__(self, directory: str):
        """初始化DataSaver类。

        Args:
            directory: 数据保存的目录。
        """
        self.directory = directory
    
    def save_dataframes(self, dataframes: Dict[str, pd.DataFrame]):
        """将多份数据保存到指定目录下，得到多份文件。

        Args:
            dataframes: 待保存的DataFrame字典，键为文件名，值为DataFrame数据。

        Returns:
            保存的文件列表。
        """
        if not os.path.exists(self.directory):
            os.makedirs(self.directory)
        self.dataframes = dataframes

        file_list = []
        for file_name, df in dataframes.items():
            file_path = os.path.join(self.directory, file_name)
            if file_path.endswith(".csv"):
                df.to_csv(file_path, index=False)
            elif file_path.endswith(".xlsx"):
                df.to_excel(file_path, index=False)
            else:
                raise ValueError("不支持的文件类型")
            file_list.append(file_path)

        return file_list
    
    def zip_directory(self, zip_filename='data.zip', password=None):
        """
        将目录中的所有文件压缩为一个zip文件，可选择是否进行加密。

        Args:
            zip_filename (str): zip文件的文件名
            password (str or bytes): 加密密码，默认不加密
        """
        # 压缩目录下所有文件
        zip_path = os.path.join(self.directory, zip_filename)
        
        flist = []
        for root, dirs, files in os.walk(self.directory):
            for file in files:
                if file in self.dataframes.keys():
                    file_path = os.path.join(root, file)
                    flist.append((file_path, os.path.relpath(file_path, self.directory)))
            
        with pyzipper.AESZipFile(zip_path, 'w', 
                                 compression=pyzipper.ZIP_LZMA,
                                 encryption=pyzipper.WZ_AES if password else None) as zf:
            if password:
                if isinstance(password, str):
                    # 密码需要转bytes类型
                    password = password.encode()
                zf.setpassword(password)
            
            for ob in flist:
                file_path = ob[0]
                rel_file_path = ob[1]
                zf.write(file_path, arcname=rel_file_path)

    def delete_files_except_zip(self):
        """删除目录下除了zip文件以外的本次被打包文件。"""
        for file_name in self.dataframes.keys():
            file_path = os.path.join(self.directory, file_name)
            os.remove(file_path)


# 初始化DataSaver
saver = DataSaver("D:/新建文件夹/data")

# 保存数据
df1 = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})
df2 = pd.DataFrame({"A": [7, 8, 9], "B": [10, 11, 12]})
dataframes = {"file1.csv": df1}
file_list = saver.save_dataframes(dataframes)

# 打包文件
saver.zip_directory('train.zip', '123456')

# 删除非zip文件
saver.delete_files_except_zip()