import os
import sys
import shutil
import subprocess
import datetime
import zipfile
import json
import tkinter as tk
from tkinter import filedialog, messagebox, scrolledtext
import threading

class ConfigManager:
    def __init__(self):
        self.config_file = "backup_config.json"
        self.local_config = {
            "project_root": "",
            "pg_bin_path": "",
            "last_backup_path": ""
        }
        self.env_config = {}
        self.load_local_config()
        
        # 如果有项目根目录，尝试加载 .env
        if self.local_config["project_root"]:
            self.load_env(os.path.join(self.local_config["project_root"], "server", ".env"))

    def load_local_config(self):
        """加载本地工具配置"""
        if os.path.exists(self.config_file):
            try:
                with open(self.config_file, 'r', encoding='utf-8') as f:
                    saved = json.load(f)
                    self.local_config.update(saved)
            except Exception as e:
                print(f"加载配置文件失败: {e}")

    def save_local_config(self):
        """保存本地工具配置"""
        try:
            with open(self.config_file, 'w', encoding='utf-8') as f:
                json.dump(self.local_config, f, indent=2, ensure_ascii=False)
        except Exception as e:
            print(f"保存配置文件失败: {e}")

    def load_env(self, env_path):
        """解析 .env 文件"""
        self.env_config = {}
        if not os.path.exists(env_path):
            return False
        
        try:
            with open(env_path, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if not line or line.startswith('#'):
                        continue
                    if '=' in line:
                        key, value = line.split('=', 1)
                        key = key.strip()
                        value = value.strip()
                        if (value.startswith('"') and value.endswith('"')) or \
                           (value.startswith("'") and value.endswith("'")):
                            value = value[1:-1]
                        self.env_config[key] = value
            return True
        except Exception:
            return False

    def get_env(self, key, default=None):
        return self.env_config.get(key, default)

    def get_upload_paths(self):
        """获取所有需要备份的文件路径"""
        paths = []
        server_root = os.path.join(self.local_config["project_root"], "server")
        
        # 1. 主目录
        main_dir = self.get_env('UPLOAD_DIR', 'uploads')
        if os.path.isabs(main_dir):
            paths.append(main_dir)
        else:
            paths.append(os.path.abspath(os.path.join(server_root, main_dir)))
            
        # 2. 备用目录
        fallback_str = self.get_env('UPLOAD_FALLBACK_DIRS')
        if fallback_str:
            fallbacks = [p.strip() for p in fallback_str.split(',') if p.strip()]
            for p in fallbacks:
                abs_path = p if os.path.isabs(p) else os.path.abspath(os.path.join(server_root, p))
                if abs_path not in paths:
                    paths.append(abs_path)
                    
        # 3. 默认 uploads 目录兜底
        default_uploads = os.path.abspath(os.path.join(server_root, 'uploads'))
        if default_uploads not in paths:
            paths.append(default_uploads)
            
        return paths

class BackupApp:
    def __init__(self, root):
        self.root = root
        self.root.title("ST系统 备份与恢复工具 (可配置版)")
        self.root.geometry("700x600")
        
        self.config_mgr = ConfigManager()
        
        # 如果没有配置项目路径，尝试自动检测当前目录
        if not self.config_mgr.local_config["project_root"]:
            current_dir = os.path.dirname(os.path.abspath(__file__))
            if os.path.exists(os.path.join(current_dir, "server", ".env")):
                self.config_mgr.local_config["project_root"] = current_dir
                self.config_mgr.load_env(os.path.join(current_dir, "server", ".env"))
                self.config_mgr.save_local_config()

        # 如果没有配置 PG 路径，尝试自动检测
        if not self.config_mgr.local_config["pg_bin_path"]:
            found_pg = self.find_pg_bin()
            if found_pg:
                self.config_mgr.local_config["pg_bin_path"] = found_pg
                self.config_mgr.save_local_config()

        self.setup_ui()
        self.check_status()

    def find_pg_bin(self):
        """尝试自动查找 PostgreSQL bin 目录"""
        try:
            subprocess.run(['pg_dump', '--version'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            return "SYSTEM_PATH" # 表示在系统环境变量中
        except FileNotFoundError:
            pass

        common_paths = [
            r"C:\Program Files\PostgreSQL",
            r"C:\Program Files (x86)\PostgreSQL"
        ]
        for base in common_paths:
            if os.path.exists(base):
                for version in os.listdir(base):
                    bin_path = os.path.join(base, version, "bin")
                    if os.path.exists(os.path.join(bin_path, "pg_dump.exe")):
                        return bin_path
        return ""

    def setup_ui(self):
        # 1. 设置区域
        frame_settings = tk.LabelFrame(self.root, text="环境配置", padx=10, pady=10)
        frame_settings.pack(fill=tk.X, padx=10, pady=5)

        # 项目路径
        tk.Label(frame_settings, text="项目根目录:").grid(row=0, column=0, sticky="w")
        self.entry_project = tk.Entry(frame_settings, width=50)
        self.entry_project.grid(row=0, column=1, padx=5)
        self.entry_project.insert(0, self.config_mgr.local_config["project_root"])
        tk.Button(frame_settings, text="浏览...", command=self.browse_project).grid(row=0, column=2)

        # PG 路径
        tk.Label(frame_settings, text="PostgreSQL bin:").grid(row=1, column=0, sticky="w")
        self.entry_pg = tk.Entry(frame_settings, width=50)
        self.entry_pg.grid(row=1, column=1, padx=5)
        self.entry_pg.insert(0, self.config_mgr.local_config["pg_bin_path"])
        tk.Button(frame_settings, text="浏览...", command=self.browse_pg).grid(row=1, column=2)
        
        tk.Button(frame_settings, text="保存并重新加载配置", command=self.save_settings, bg="#FF9800", fg="white").grid(row=2, column=1, pady=10)

        # 2. 备份路径区域
        frame_path = tk.LabelFrame(self.root, text="备份保存位置", padx=10, pady=5)
        frame_path.pack(fill=tk.X, padx=10, pady=5)
        
        default_backup = self.config_mgr.local_config["last_backup_path"] or os.getcwd()
        self.lbl_backup_path = tk.Label(frame_path, text=default_backup, relief=tk.SUNKEN, anchor="w", bg="white")
        self.lbl_backup_path.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 10))
        tk.Button(frame_path, text="更改路径...", command=self.choose_backup_dir).pack(side=tk.RIGHT)

        # 3. 操作按钮
        frame_actions = tk.Frame(self.root, pady=10)
        frame_actions.pack(fill=tk.X)
        
        self.btn_backup = tk.Button(frame_actions, text="一键备份", command=self.start_backup, 
                             bg="#4CAF50", fg="white", font=("Microsoft YaHei", 12, "bold"), height=2, width=20)
        self.btn_backup.pack(side=tk.LEFT, padx=40)
        
        self.btn_restore = tk.Button(frame_actions, text="一键恢复", command=self.start_restore, 
                              bg="#2196F3", fg="white", font=("Microsoft YaHei", 12, "bold"), height=2, width=20)
        self.btn_restore.pack(side=tk.RIGHT, padx=40)

        # 3.5 提示信息
        frame_tips = tk.Frame(self.root)
        frame_tips.pack(fill=tk.X, pady=5)
        tk.Label(frame_tips, text="⚠️ 恢复注意事项：请确保数据库为空（不要运行 init-db），否则会因表已存在而失败。", 
                 fg="red", font=("Microsoft YaHei", 9)).pack()

        # 4. 信息显示
        self.lbl_status = tk.Label(self.root, text="正在检查配置...", fg="blue")
        self.lbl_status.pack(pady=5)

        # 5. 日志
        tk.Label(self.root, text="运行日志:").pack(anchor="w", padx=10)
        self.log_area = scrolledtext.ScrolledText(self.root, height=12)
        self.log_area.pack(fill=tk.BOTH, expand=True, padx=10, pady=(0, 10))

    def log(self, message):
        timestamp = datetime.datetime.now().strftime("%H:%M:%S")
        self.log_area.insert(tk.END, f"[{timestamp}] {message}\n")
        self.log_area.see(tk.END)

    def browse_project(self):
        path = filedialog.askdirectory(title="选择项目根目录 (包含 server 文件夹)")
        if path:
            self.entry_project.delete(0, tk.END)
            self.entry_project.insert(0, path)

    def browse_pg(self):
        path = filedialog.askdirectory(title="选择 PostgreSQL bin 目录")
        if path:
            self.entry_pg.delete(0, tk.END)
            self.entry_pg.insert(0, path)

    def save_settings(self):
        project_root = self.entry_project.get().strip()
        pg_bin = self.entry_pg.get().strip()
        
        self.config_mgr.local_config["project_root"] = project_root
        self.config_mgr.local_config["pg_bin_path"] = pg_bin
        self.config_mgr.save_local_config()
        
        # 重新加载 .env
        env_path = os.path.join(project_root, "server", ".env")
        if self.config_mgr.load_env(env_path):
            self.log("配置保存成功，已加载 .env 文件")
        else:
            self.log("配置保存成功，但无法找到或解析 .env 文件！")
            
        self.check_status()

    def check_status(self):
        ready = True
        msg = []
        
        if not self.config_mgr.env_config:
            ready = False
            msg.append("未加载 .env 配置")
        else:
            db = self.config_mgr.get_env('DB_DATABASE')
            msg.append(f"数据库: {db}")

        pg_path = self.config_mgr.local_config["pg_bin_path"]
        if not pg_path and self.find_pg_bin() == "":
             msg.append("PG工具: 未配置")
             ready = False
        elif pg_path == "SYSTEM_PATH":
             msg.append("PG工具: 系统环境变量")
        else:
             msg.append(f"PG工具: {pg_path}")

        if ready:
            self.lbl_status.config(text=" | ".join(msg), fg="green")
            self.btn_backup.config(state=tk.NORMAL)
            self.btn_restore.config(state=tk.NORMAL)
        else:
            self.lbl_status.config(text="配置不完整: " + " | ".join(msg), fg="red")
            self.btn_backup.config(state=tk.DISABLED)
            self.btn_restore.config(state=tk.DISABLED)

    def choose_backup_dir(self):
        path = filedialog.askdirectory(initialdir=self.lbl_backup_path.cget("text"), title="选择备份保存位置")
        if path:
            self.lbl_backup_path.config(text=path)
            self.config_mgr.local_config["last_backup_path"] = path
            self.config_mgr.save_local_config()
            self.log(f"备份保存路径已更改为: {path}")

    def get_pg_cmd(self, cmd):
        pg_path = self.config_mgr.local_config["pg_bin_path"]
        if pg_path and pg_path != "SYSTEM_PATH":
            return os.path.join(pg_path, cmd)
        return cmd

    def start_backup(self):
        if not messagebox.askyesno("确认备份", "确定要开始备份吗？"):
            return
        threading.Thread(target=self.run_backup, daemon=True).start()

    def run_backup(self):
        try:
            self.log("=== 开始备份 ===")
            save_path = self.lbl_backup_path.cget("text")
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            
            # 临时目录
            temp_dir = os.path.join(save_path, f"temp_backup_{timestamp}")
            os.makedirs(temp_dir, exist_ok=True)
            
            # 1. 数据库
            self.log("正在导出数据库...")
            db_file = os.path.join(temp_dir, "database.sql")
            env = os.environ.copy()
            env['PGPASSWORD'] = self.config_mgr.get_env('DB_PASSWORD', '')
            
            cmd = [
                self.get_pg_cmd('pg_dump'),
                '-h', self.config_mgr.get_env('DB_HOST', 'localhost'),
                '-p', self.config_mgr.get_env('DB_PORT', '5432'),
                '-U', self.config_mgr.get_env('DB_USER', 'postgres'),
                '-F', 'p', '-f', db_file,
                self.config_mgr.get_env('DB_DATABASE', 'stsystem')
            ]
            
            startupinfo = subprocess.STARTUPINFO()
            startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            
            proc = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, startupinfo=startupinfo)
            out, err = proc.communicate()
            if proc.returncode != 0:
                raise Exception(f"数据库备份失败: {err.decode('gbk', errors='ignore')}")
            
            # 2. 文件
            self.log("正在备份文件...")
            files_dir = os.path.join(temp_dir, "files")
            os.makedirs(files_dir)
            
            paths = self.config_mgr.get_upload_paths()
            for src in paths:
                if os.path.exists(src):
                    dirname = os.path.basename(src)
                    dest = os.path.join(files_dir, dirname)
                    if os.path.exists(dest): dest += f"_{paths.index(src)}"
                    shutil.copytree(src, dest)
            
            # 3. 打包
            self.log("正在打包...")
            zip_name = f"STSystem_Backup_{timestamp}.zip"
            zip_full_path = os.path.join(save_path, zip_name)
            
            with zipfile.ZipFile(zip_full_path, 'w', zipfile.ZIP_DEFLATED) as zf:
                for root, dirs, files in os.walk(temp_dir):
                    for file in files:
                        fp = os.path.join(root, file)
                        arcname = os.path.relpath(fp, temp_dir)
                        zf.write(fp, arcname)
            
            shutil.rmtree(temp_dir)
            self.log(f"备份成功: {zip_full_path}")
            messagebox.showinfo("成功", f"备份已保存:\n{zip_full_path}")
            
        except Exception as e:
            self.log(f"错误: {e}")
            messagebox.showerror("错误", str(e))

    def start_restore(self):
        zip_path = filedialog.askopenfilename(filetypes=[("Zip", "*.zip")])
        if not zip_path: return
        
        msg = (
            "⚠️ 高风险操作警告 ⚠️\n\n"
            "1. 此操作将覆盖数据库和文件。\n"
            "2. 请确保目标数据库是【空数据库】（不要运行 init-db）。\n"
            "   如果数据库中已有表，恢复将失败！\n"
            "3. 请确保已停止 Node 后端服务。\n\n"
            "确定要继续吗？"
        )
        
        if not messagebox.askyesno("确认恢复", msg, icon='warning'):
            return
        threading.Thread(target=self.run_restore, args=(zip_path,), daemon=True).start()

    def run_restore(self, zip_path):
        try:
            self.log("=== 开始恢复 ===")
            temp_dir = os.path.join(os.path.dirname(zip_path), "temp_restore_extract")
            if os.path.exists(temp_dir): shutil.rmtree(temp_dir)
            os.makedirs(temp_dir)
            
            with zipfile.ZipFile(zip_path, 'r') as zf:
                zf.extractall(temp_dir)
                
            # 1. 数据库
            self.log("恢复数据库...")
            db_file = os.path.join(temp_dir, "database.sql")
            if not os.path.exists(db_file): raise Exception("未找到 database.sql")
            
            env = os.environ.copy()
            env['PGPASSWORD'] = self.config_mgr.get_env('DB_PASSWORD', '')
            
            cmd = [
                self.get_pg_cmd('psql'),
                '-h', self.config_mgr.get_env('DB_HOST', 'localhost'),
                '-p', self.config_mgr.get_env('DB_PORT', '5432'),
                '-U', self.config_mgr.get_env('DB_USER', 'postgres'),
                '-d', self.config_mgr.get_env('DB_DATABASE', 'stsystem'),
                '-v', 'ON_ERROR_STOP=1', # 遇到错误立即停止
                '-f', db_file
            ]
            
            startupinfo = subprocess.STARTUPINFO()
            startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            proc = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, startupinfo=startupinfo)
            out, err = proc.communicate()
            
            # 记录输出以便调试
            if out: self.log(f"PSQL OUT: {out.decode('gbk', errors='ignore')}")
            if err: self.log(f"PSQL ERR: {err.decode('gbk', errors='ignore')}")

            if proc.returncode != 0:
                err_msg = err.decode('gbk', errors='ignore')
                if "already exists" in err_msg or "已经存在" in err_msg:
                    raise Exception(f"恢复失败: 数据库中已存在表或数据。\n请先清空数据库，或使用 scripts 中的脚本重置数据库。\n详细错误: {err_msg}")
                raise Exception(f"数据库恢复失败: {err_msg}")

            # 1.5 修复序列 (防止主键冲突)
            self.log("正在修复数据库序列...")
            fix_seq_sql = """
            DO $$
            DECLARE
                r RECORD;
            BEGIN
                FOR r IN 
                    SELECT table_name, column_name
                    FROM information_schema.columns
                    WHERE table_schema = 'public' AND column_default LIKE 'nextval%'
                LOOP
                    EXECUTE 'SELECT setval(' || quote_literal(pg_get_serial_sequence(r.table_name, r.column_name)) || ', COALESCE((SELECT MAX(' || quote_ident(r.column_name) || ') FROM ' || quote_ident(r.table_name) || '), 1), false)';
                END LOOP;
            END $$;
            """
            
            cmd_fix = [
                self.get_pg_cmd('psql'),
                '-h', self.config_mgr.get_env('DB_HOST', 'localhost'),
                '-p', self.config_mgr.get_env('DB_PORT', '5432'),
                '-U', self.config_mgr.get_env('DB_USER', 'postgres'),
                '-d', self.config_mgr.get_env('DB_DATABASE', 'stsystem'),
                '-c', fix_seq_sql
            ]
            
            proc_fix = subprocess.Popen(cmd_fix, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, startupinfo=startupinfo)
            out_fix, err_fix = proc_fix.communicate()
            
            if proc_fix.returncode != 0:
                self.log(f"警告: 序列修复失败 (但这可能不影响使用): {err_fix.decode('gbk', errors='ignore')}")
            else:
                self.log("序列修复完成。")
                
            # 2. 文件
            self.log("恢复文件...")
            # 恢复到主目录
            server_root = os.path.join(self.config_mgr.local_config["project_root"], "server")
            target_dir = self.config_mgr.get_env('UPLOAD_DIR', 'uploads')
            if not os.path.isabs(target_dir):
                target_dir = os.path.abspath(os.path.join(server_root, target_dir))
            
            if not os.path.exists(target_dir): os.makedirs(target_dir)
            
            files_src = os.path.join(temp_dir, "files")
            if os.path.exists(files_src):
                for root, dirs, files in os.walk(files_src):
                    for file in files:
                        shutil.copy2(os.path.join(root, file), os.path.join(target_dir, file))
            
            shutil.rmtree(temp_dir)
            self.log("恢复成功！")
            messagebox.showinfo("成功", "恢复完成！请重启服务。")
            
        except Exception as e:
            self.log(f"错误: {e}")
            messagebox.showerror("错误", str(e))

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