import os
import sys
import threading
import time
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
import tkinter as tk
from tkinter import filedialog, messagebox, ttk, simpledialog
from huggingface_hub import HfApi

# 开启 hf_transfer 加速
os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"
# 禁用底层 CLI 的多线程进度条，防止刷屏
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"

class TextHandler(logging.Handler):
    def __init__(self, text_widget):
        super().__init__()
        self.text_widget = text_widget

    def emit(self, record):
        msg = self.format(record)
        def append():
            self.text_widget.configure(state='normal')
            self.text_widget.insert(tk.END, msg + '\n')
            self.text_widget.configure(state='disabled')
            self.text_widget.see(tk.END)
        self.text_widget.after(0, append)

class HFUploaderApp:
    def __init__(self, root):
        self.root = root
        self.root.title("HF 仓库管理 - 大模型批量上传版")
        self.root.geometry("950x700")
        
        self.token_var = tk.StringVar()
        self.repo_id_var = tk.StringVar(value="waveshaker/comfy_models")
        self.selected_tree_path = ""
        self.upload_queue = {}
        
        # 控制停止的全局标志位
        self.is_uploading = False
        
        self.setup_ui()
        self.load_token()
        
    def setup_ui(self):
        # 1. 顶部 Token 区
        top_frame = tk.Frame(self.root, pady=10)
        top_frame.pack(fill=tk.X, padx=10)
        tk.Label(top_frame, text="HF Token:").pack(side=tk.LEFT)
        tk.Entry(top_frame, textvariable=self.token_var, show="*", width=50).pack(side=tk.LEFT, padx=5)
        tk.Button(top_frame, text="保存 Token", command=self.save_token).pack(side=tk.LEFT)
        
        # 2. 仓库配置区
        repo_frame = tk.Frame(self.root)
        repo_frame.pack(fill=tk.X, padx=10, pady=5)
        tk.Label(repo_frame, text="仓库 ID:").pack(side=tk.LEFT)
        tk.Entry(repo_frame, textvariable=self.repo_id_var, width=30).pack(side=tk.LEFT, padx=5)
        
        # 3. 中间分割（左侧：队列，右侧：远端树）
        paned = tk.PanedWindow(self.root, orient=tk.HORIZONTAL, sashwidth=5)
        paned.pack(fill=tk.BOTH, expand=True, padx=10, pady=5)
        
        # --- 左侧：上传队列 ---
        left_frame = tk.Frame(paned)
        paned.add(left_frame, width=550)
        tk.Label(left_frame, text="上传队列 (双击目标路径可修改):").pack(anchor=tk.W)
        
        # 🌟 修正：列名统一为 source
        cols = ("source", "remote", "status")
        self.queue_tree = ttk.Treeview(left_frame, columns=cols, show="headings", height=10)
        self.queue_tree.heading("source", text="文件源 (本地/远端链接)")
        self.queue_tree.heading("remote", text="远端目标目录")
        self.queue_tree.heading("status", text="状态")
        self.queue_tree.column("source", width=250)
        self.queue_tree.column("remote", width=150)
        self.queue_tree.column("status", width=80, anchor="center")
        self.queue_tree.tag_configure("done", foreground="gray")

        self.queue_tree.pack(fill=tk.BOTH, expand=True)
        self.queue_tree.bind("<Double-1>", self.on_edit_remote_path)
        
        queue_btns = tk.Frame(left_frame)
        queue_btns.pack(fill=tk.X, pady=5)
        tk.Button(queue_btns, text="➕ 添加本地文件", command=self.add_to_queue).pack(side=tk.LEFT)
        tk.Button(queue_btns, text="❌ 移除选中", command=self.remove_from_queue).pack(side=tk.LEFT, padx=5)

        # 🌟 极简的链接输入框
        self.remote_url_var = tk.StringVar()
        tk.Entry(queue_btns, textvariable=self.remote_url_var, width=30).pack(side=tk.LEFT, padx=5)
        tk.Button(queue_btns, text="🔗 加入队列", command=self.add_remote_to_queue, ).pack(side=tk.LEFT)

        right_frame = tk.Frame(paned)
        paned.add(right_frame, width=350)
        tk.Label(right_frame, text="远端仓库目录树:").pack(anchor=tk.W)
        
        self.tree_view = ttk.Treeview(right_frame, show="tree", height=10)
        self.tree_view.pack(fill=tk.BOTH, expand=True)
        self.tree_view.bind("<<TreeviewSelect>>", self.on_tree_select)
        
        tree_btns = tk.Frame(right_frame)
        tree_btns.pack(fill=tk.X, pady=5)
        tk.Button(tree_btns, text="🔄 拉取远端目录", command=self.fetch_remote_tree).pack(side=tk.LEFT)
        tk.Button(tree_btns, text="📁 新建文件夹", command=self.create_remote_folder).pack(side=tk.LEFT, padx=5)
        tk.Button(tree_btns, text="⬇️ 设为目标", command=self.set_target_from_tree).pack(side=tk.LEFT)
        
        # 4. 操作区 (包含开始和停止)
        action_frame = tk.Frame(self.root, pady=10)
        action_frame.pack(fill=tk.X, padx=10)
        
        self.upload_btn = tk.Button(action_frame, text="🚀 开始并发上传", command=self.start_upload, font=("Arial", 10, "bold"))
        self.upload_btn.pack(side=tk.LEFT, ipadx=20, ipady=5)
        
        # 停止按钮
        self.stop_btn = tk.Button(action_frame, text="🛑 停止上传", command=self.stop_upload, font=("Arial", 10, "bold"), state=tk.DISABLED)
        self.stop_btn.pack(side=tk.LEFT, padx=10, ipadx=20, ipady=5)

        # 🌟 GUI 内部总进度条
        self.total_progress_var = tk.DoubleVar()
        self.progress_bar = ttk.Progressbar(action_frame, orient="horizontal", length=200, mode="determinate", variable=self.total_progress_var)
        self.progress_bar.pack(side=tk.LEFT, padx=15, fill=tk.X, expand=True)
        
        # 进度文字标签
        self.progress_text = tk.Label(action_frame, text="等待上传...", fg="#2196F3", font=("Arial", 10, "bold"))
        self.progress_text.pack(side=tk.LEFT, padx=5)
        
        # 5. 日志区
        log_frame = tk.Frame(self.root, padx=10, pady=5)
        log_frame.pack(fill=tk.BOTH, expand=True)
        tk.Label(log_frame, text="操作日志:").pack(anchor=tk.W)
        
        self.log_text = tk.Text(log_frame, height=12, state='disabled', bg="#f0f0f0")
        self.log_text.pack(fill=tk.BOTH, expand=True)
        
        self.logger = logging.getLogger("HF_Uploader")
        self.logger.setLevel(logging.INFO)
        handler = TextHandler(self.log_text)
        handler.setFormatter(logging.Formatter('[%(asctime)s] %(message)s', '%H:%M:%S'))
        self.logger.addHandler(handler)

    def load_token(self):
        if os.path.exists("hf_token.txt"):
            with open("hf_token.txt", "r") as f:
                self.token_var.set(f.read().strip())

    def save_token(self):
        with open("hf_token.txt", "w") as f:
            f.write(self.token_var.get().strip())
        messagebox.showinfo("成功", "Token 已保存到 hf_token.txt")

    def add_to_queue(self):
        files = filedialog.askopenfilenames(
            title="选择要上传的模型文件",
            filetypes=[("SafeTensors Model", "*.safetensors"), ("All Files", "*.*")]
        )
        if files:
            for file_path in files:
                if file_path not in self.upload_queue:
                    self.upload_queue[file_path] = ""
            self.refresh_queue()

    def remove_from_queue(self):
        selected = self.queue_tree.selection()
        for item in selected:
            local_path = self.queue_tree.item(item, "values")[0]
            if local_path in self.upload_queue:
                del self.upload_queue[local_path]
        self.refresh_queue()

    def refresh_queue(self):
        # 🌟 在删除之前，先把旧的“已完成”状态存下来
        existing_status = {}
        for item in self.queue_tree.get_children():
            values = self.queue_tree.item(item, "values")
            if len(values) >= 3 and values[2] == "已完成":
                existing_status[values[0]] = "已完成"
                
        # 清空现有表格
        for item in self.queue_tree.get_children():
            self.queue_tree.delete(item)
            
        # 重新插入所有数据
        for local_path, remote_target in self.upload_queue.items():
            # 如果之前传完了，状态保持“已完成”，否则是“待上传”
            status = existing_status.get(local_path, "待上传")
            
            if status == "已完成":
                self.queue_tree.insert("", tk.END, values=(local_path, remote_target, status), tags=("done",))
            else:
                self.queue_tree.insert("", tk.END, values=(local_path, remote_target, status))
            
        # 自动选中最后一个
        children = self.queue_tree.get_children()
        if children:
            last_item = children[-1]
            self.queue_tree.selection_set(last_item)
            self.queue_tree.focus(last_item)
            self.queue_tree.see(last_item)

    def on_edit_remote_path(self, event):
        item = self.queue_tree.selection()
        if not item: return
        item = item[0]
        current_local = self.queue_tree.item(item, "values")[0]
        current_remote = self.queue_tree.item(item, "values")[1]
        
        top = tk.Toplevel(self.root)
        top.title("修改目标路径")
        tk.Label(top, text="输入远端目标文件夹 (例如: moss/ 或留空):").pack(padx=10, pady=5)
        entry = tk.Entry(top, width=40)
        entry.insert(0, current_remote)
        entry.pack(padx=10, pady=5)
        
        def save():
            new_path = entry.get().strip().strip("/") 
            self.upload_queue[current_local] = new_path
            self.refresh_queue()
            top.destroy()
            
        tk.Button(top, text="确定", command=save).pack(pady=10)

    def fetch_remote_tree(self):
        token = self.token_var.get().strip()
        repo_id = self.repo_id_var.get().strip()
        if not token or not repo_id: return

        self.logger.info("正在拉取远端目录结构...")
        self.tree_view.delete(*self.tree_view.get_children())
        
        def fetch():
            try:
                api = HfApi(token=token)
                repo_files = api.list_repo_files(repo_id=repo_id, repo_type="model")
                tree_dict = {}
                for file_path in repo_files:
                    parts = file_path.split('/')
                    current = tree_dict
                    for part in parts:
                        if part not in current: current[part] = {}
                        current = current[part]
                
                def insert_tree(parent, dict_data, current_path=""):
                    for key, children in dict_data.items():
                        path = f"{current_path}/{key}" if current_path else key
                        node_text = f"📄 {key}" if not children else f"📁 {key}"
                        node = self.tree_view.insert(parent, tk.END, text=node_text, values=(path,))
                        if children: insert_tree(node, children, path)
                
                self.root.after(0, lambda: insert_tree("", tree_dict))
                self.logger.info(f"成功加载远端目录结构。")
            except Exception as e:
                self.root.after(0, lambda err=e: self.logger.error(f"拉取失败: {err}"))
        threading.Thread(target=fetch, daemon=True).start()

    def create_remote_folder(self):
        token = self.token_var.get().strip()
        repo_id = self.repo_id_var.get().strip()
        if not token or not repo_id: return

        parent_path = self.selected_tree_path if self.selected_tree_path else ""
        folder_name = simpledialog.askstring("新建文件夹", f"在 '{parent_path or '根目录'}' 下新建文件夹:")
        if not folder_name: return
        folder_name = folder_name.strip().strip("/")
        
        new_folder_path = f"{parent_path}/{folder_name}" if parent_path else folder_name
        gitkeep_path = f"{new_folder_path}/.gitkeep"

        self.logger.info(f"正在新建远端文件夹: {new_folder_path} ...")
        def create():
            try:
                api = HfApi(token=token)
                api.upload_file(path_or_fileobj=b"", path_in_repo=gitkeep_path, repo_id=repo_id, repo_type="model")
                self.root.after(0, lambda: self.logger.info(f"✅ 文件夹新建成功: {new_folder_path}"))
                self.root.after(0, self.fetch_remote_tree)
            except Exception as e:
                self.root.after(0, lambda err=e: self.logger.error(f"❌ 新建失败: {err}"))
        threading.Thread(target=create, daemon=True).start()

    def on_tree_select(self, event):
        selected = self.tree_view.selection()
        if selected:
            self.selected_tree_path = self.tree_view.item(selected[0], "values")[0]
            if "." in os.path.basename(self.selected_tree_path):
                self.selected_tree_path = os.path.dirname(self.selected_tree_path)

    def set_target_from_tree(self):
        if not self.selected_tree_path:
            messagebox.showwarning("提示", "请先在右侧树中选中一个文件夹")
            return
            
        queue_selected = self.queue_tree.selection()
        if not queue_selected:
            messagebox.showwarning("提示", "请先在左侧队列中选中任务")
            return
            
        target_path = self.selected_tree_path.strip("/")
        count = 0
        for item in queue_selected:
            values = self.queue_tree.item(item, "values")
            # 🌟 如果已经是“已完成”，就不改它的路径了
            if len(values) >= 3 and values[2] == "已完成":
                continue
            local_path = values[0]
            self.upload_queue[local_path] = target_path
            count += 1
            
        self.refresh_queue()
        self.logger.info(f"✅ 已将 {count} 个文件的目标路径批量设置为: {target_path}")

    # === 上传与停止逻辑 ===
    def start_upload(self):
        token = self.token_var.get().strip()
        repo_id = self.repo_id_var.get().strip()
        
        # 🌟 核心修复：先检查有哪些文件是“待上传”的
        pending_tasks = []
        for tree_item in self.queue_tree.get_children():
            values = self.queue_tree.item(tree_item, "values")
            status = values[2] if len(values) >= 3 else "待上传"
            if status != "已完成":
                pending_tasks.append(tree_item)
                
        if not token or not repo_id or not pending_tasks:
            messagebox.showwarning("警告", "请检查 Token、仓库 ID，或者当前没有待上传的新文件")
            return
            
        self.is_uploading = True
        self.upload_btn.config(state=tk.DISABLED, text="上传中...")
        self.stop_btn.config(state=tk.NORMAL)
        
        # 计算要上传的真实数量
        total_files = len(pending_tasks)
        self.logger.info(f"🚀 开始并发上传 {total_files} 个待处理文件...")
        
        def upload_worker():
            api = HfApi(token=token)
            
            tasks = []
            for tree_item in pending_tasks:
                values = self.queue_tree.item(tree_item, "values")
                source = values[0]
                remote_dir = values[1]
                
                # 区分是 HTTP 链接还是本地文件
                if source.startswith("http://") or source.startswith("https://"):
                    # 这是一个远端链接，提取文件名
                    filename = source.split("/")[-1]
                    repo_path = f"{remote_dir}/{filename}" if remote_dir else filename
                    tasks.append((tree_item, source, repo_path, True)) # True 代表是远端
                else:
                    # 这是一个本地文件
                    filename = os.path.basename(source)
                    repo_path = f"{remote_dir}/{filename}" if remote_dir else filename
                    tasks.append((tree_item, source, repo_path, False)) # False 代表是本地

            
            completed_files = 0
            progress_lock = threading.Lock()

            def update_ui_progress(completed, total):
                percent = (completed / total) * 100
                self.root.after(0, lambda: self.total_progress_var.set(percent))
                self.root.after(0, lambda: self.progress_text.config(text=f"{completed} / {total} 个文件 ({percent:.1f}%)"))

            def upload_one(tree_item, source, repo_path, is_remote):
                nonlocal completed_files
                if not self.is_uploading:
                    return False, "用户已取消"
                
                try:
                    if is_remote:
                        import subprocess
                        import os
                        
                        clean_path = source.replace("https://huggingface.co/", "")
                        parts = clean_path.split("/")
                        if len(parts) >= 5 and parts[2] in ["blob", "resolve"]:
                            src_repo_id = f"{parts[0]}/{parts[1]}"
                            src_file_path = "/".join(parts[4:]) # 这里已经修复过了
                            
                            # 读取 GUI 里的 token
                            current_token = self.token_var.get().strip()
                            
                            cmd = [
                                "hf", "cp",
                                f"hf://{src_repo_id}/{src_file_path}",
                                f"hf://{repo_id}/{repo_path}"
                            ]
                            
                            # 🌟 强行把 Token 塞给子进程的环境变量，防止它去读错误的系统缓存
                            env = os.environ.copy()
                            env["HF_TOKEN"] = current_token
                            
                            result = subprocess.run(cmd, capture_output=True, text=True, env=env)
                            if result.returncode != 0:
                                raise RuntimeError(result.stderr.strip() or result.stdout.strip())
                    else:
                        # 🌟 本地文件上传逻辑
                        api.upload_file(
                            path_or_fileobj=source, 
                            path_in_repo=repo_path, 
                            repo_id=repo_id, 
                            repo_type="model"
                        )
                    
                    with progress_lock:
                        completed_files += 1
                        update_ui_progress(completed_files, total_files)
                        
                    self.root.after(0, lambda: self.mark_item_done(tree_item))
                    return True, repo_path
                except Exception as e:
                    return False, f"{repo_path} ({e})"

            self.root.after(0, lambda: self.total_progress_var.set(0))
            self.root.after(0, lambda: self.progress_text.config(text=f"0 / {total_files} 个文件", fg="#2196F3"))

            with ThreadPoolExecutor(max_workers=4) as executor:
                future_to_task = {executor.submit(upload_one, t, s, r, is_r): r for t, s, r, is_r in tasks}
                for future in as_completed(future_to_task):
                    if not self.is_uploading: break 
                    success, msg = future.result()
                    if success:
                        self.root.after(0, lambda m=msg: self.logger.info(f"✅ 成功: {m}"))
                    elif msg != "用户已取消":
                        self.root.after(0, lambda m=msg: self.logger.error(f"❌ 失败: {m}"))
            
            self.root.after(0, self.upload_finished)

        threading.Thread(target=upload_worker, daemon=True).start()

    def update_ui_progress(self, percent, completed, total):
        self.total_progress_var.set(percent)
        self.progress_text.config(text=f"{completed} / {total} 个文件 ({percent:.1f}%)")

    def mark_item_done(self, tree_item):
        """将表格中的某一行标记为已完成（灰色）"""
        values = list(self.queue_tree.item(tree_item, "values"))
        if len(values) >= 3:
            values[2] = "已完成"
        else:
            values.append("已完成")
        self.queue_tree.item(tree_item, values=values, tags=("done",))
        
    def stop_upload(self):
        """强行停止上传"""
        if messagebox.askyesno("确认停止", "确定要停止上传吗？\n（已传完的部分已保存在服务器，下次点上传会自动断点续传）"):
            self.is_uploading = False
            self.logger.warning("🛑 已中断后台请求，界面已重置。等待后台线程安全退出...")
            self.upload_btn.config(state=tk.NORMAL, text="🚀 开始并发上传")
            self.stop_btn.config(state=tk.DISABLED)
            self.progress_text.config(text="已完成 / 已停止", fg="#2E7D32")

    def upload_finished(self):
        self.is_uploading = False
        self.upload_btn.config(state=tk.NORMAL, text="🚀 开始并发上传")
        self.stop_btn.config(state=tk.DISABLED)
        self.progress_text.config(text="已完成 / 已停止", fg="#2E7D32")
        self.logger.info("🎉 任务处理完毕！")
        messagebox.showinfo("完成", "上传任务已结束（或已停止），请查看日志。")

    def add_remote_to_queue(self):
        url = self.remote_url_var.get().strip()
        if not url:
            messagebox.showwarning("提示", "请先粘贴 HF 文件链接")
            return
            
        # 自动拼合一个初始的目标路径（比如默认丢到 diffusion_models 下）
        if url not in self.upload_queue:
            self.upload_queue[url] = "" # 默认路径留空
            self.refresh_queue()
            self.remote_url_var.set("") # 加入后清空输入框
            
            # 自动把最后一行高亮
            children = self.queue_tree.get_children()
            if children:
                self.queue_tree.selection_set(children[-1])

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