import os
import re
import tkinter as tk
from tkinter import ttk, filedialog, scrolledtext, messagebox, simpledialog
import multiprocessing as mp
from multiprocessing import Pool, Manager
import queue
import threading
from functools import partial
import mimetypes

# 定義文件類型分類（可擴展）
FILE_TYPE_CATEGORIES = {
    "音頻文件": ['.mp3', '.wav', '.flac', '.aac', '.ogg', '.wma', '.m4a', '.ape', '.alac'],
    "視頻文件": ['.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv', '.mpg', '.mpeg', '.rmvb', '.3gp'],
    "圖像文件": ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp', '.svg', '.psd', '.ai'],
    "文檔文件": ['.doc', '.docx', '.pdf', '.txt', '.xls', '.xlsx', '.ppt', '.pptx', '.md', '.rtf'],
    "壓縮文件": ['.zip', '.rar', '.7z', '.tar', '.gz', '.bz2'],
    "程序文件": ['.exe', '.dll', '.py', '.java', '.c', '.cpp', '.js', '.html', '.css']
}

class FileSearchApp:
    def __init__(self, root):
        self.root = root
        self.root.title("文件智能搜索工具")
        self.root.geometry("1200x700")
        self.root.minsize(1000, 600)
        
        # 設(shè)置中文字體
        self.style = ttk.Style()
        self.style.configure(".", font=("SimHei", 10))
        
        # 搜索參數(shù)
        self.search_path = tk.StringVar()
        self.keyword = tk.StringVar()
        self.search_type = tk.StringVar(value="name")  # name 或 content
        self.is_searching = False
        self.manager = Manager()  # 用于創(chuàng)建跨進程共享對象
        self.stop_event = self.manager.Event()  # 跨進程停止事件
        
        # 文件類型篩選狀態(tài)（只搜索選中的類型）
        self.category_vars = {}  # 存儲類別勾選狀態(tài)
        self.extension_vars = {}  # 存儲擴展名勾選狀態(tài)
        self.included_extensions = set()  # 當(dāng)前選中的要包含的擴展名
        self.other_files_frame = None  # "其他文件"類別的擴展架子框架
        
        # 創(chuàng)建UI
        self.create_widgets()
        # 初始化文件類型篩選狀態(tài)（默認(rèn)不選中任何類型）
        self.init_file_type_filters()
        
        # 進程和隊列
        self.progress_queue = None
        self.result_queue = None
        self.search_process = None
        
        # 搜索統(tǒng)計
        self.total_files = 0
        self.processed_files = 0
        self.matched_files = 0

    def create_widgets(self):
        # 主框架分割為左右兩部分
        main_paned = ttk.PanedWindow(self.root, orient=tk.HORIZONTAL)
        main_paned.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 左側(cè)：文件類型篩選面板
        filter_frame = ttk.LabelFrame(main_paned, text="文件類型篩選（只搜索選中類型）", padding="10")
        main_paned.add(filter_frame, weight=1)
        
        # 篩選面板滾動區(qū)域
        self.filter_canvas = tk.Canvas(filter_frame)
        self.filter_scrollbar = ttk.Scrollbar(filter_frame, orient="vertical", command=self.filter_canvas.yview)
        self.filter_scrollable_frame = ttk.Frame(self.filter_canvas)
        
        self.filter_scrollable_frame.bind(
            "<Configure>",
            lambda e: self.filter_canvas.configure(scrollregion=self.filter_canvas.bbox("all"))
        )
        
        self.filter_canvas.create_window((0, 0), window=self.filter_scrollable_frame, anchor="nw")
        self.filter_canvas.configure(yscrollcommand=self.filter_scrollbar.set)
        
        self.filter_canvas.pack(side="left", fill="both", expand=True)
        self.filter_scrollbar.pack(side="right", fill="y")
        
        # 添加文件類型分類和擴展名復(fù)選框
        for category, extensions in FILE_TYPE_CATEGORIES.items():
            # 類別總復(fù)選框
            cat_var = tk.BooleanVar()
            self.category_vars[category] = cat_var
            cat_check = ttk.Checkbutton(
                self.filter_scrollable_frame, 
                text=category, 
                variable=cat_var,
                command=partial(self.toggle_category, category)
            )
            cat_check.pack(anchor=tk.W, pady=5)
            
            # 該類別下的擴展名復(fù)選框
            ext_frame = ttk.Frame(self.filter_scrollable_frame)
            ext_frame.pack(anchor=tk.W, padx=20)
            
            for ext in extensions:
                ext_var = tk.BooleanVar()
                self.extension_vars[ext] = ext_var
                ext_check = ttk.Checkbutton(
                    ext_frame, 
                    text=ext, 
                    variable=ext_var,
                    command=partial(self.update_category_state, category)
                )
                ext_check.pack(side=tk.LEFT, padx=5, pady=2)
        
        # 添加自定義擴展名按鈕
        ttk.Button(
            filter_frame, 
            text="添加自定義擴展名", 
            command=self.add_custom_extension
        ).pack(pady=10, fill=tk.X)
        
        # 右側(cè)：主功能區(qū)
        right_frame = ttk.Frame(main_paned)
        main_paned.add(right_frame, weight=3)
        
        # 頂部框架 - 路徑選擇
        path_frame = ttk.Frame(right_frame, padding="10")
        path_frame.pack(fill=tk.X)
        
        ttk.Label(path_frame, text="搜索路徑:").pack(side=tk.LEFT, padx=5)
        ttk.Entry(path_frame, textvariable=self.search_path, width=50).pack(side=tk.LEFT, padx=5, fill=tk.X, expand=True)
        ttk.Button(path_frame, text="瀏覽...", command=self.browse_path).pack(side=tk.LEFT, padx=5)
        
        # 中間框架 - 搜索設(shè)置
        search_frame = ttk.Frame(right_frame, padding="10")
        search_frame.pack(fill=tk.X)
        
        ttk.Label(search_frame, text="搜索關(guān)鍵詞:").pack(side=tk.LEFT, padx=5)
        ttk.Entry(search_frame, textvariable=self.keyword, width=30).pack(side=tk.LEFT, padx=5)
        
        ttk.Radiobutton(search_frame, text="按文件名", variable=self.search_type, value="name").pack(side=tk.LEFT, padx=5)
        ttk.Radiobutton(search_frame, text="按文件內(nèi)容", variable=self.search_type, value="content").pack(side=tk.LEFT, padx=5)
        
        ttk.Button(search_frame, text="開始搜索", command=self.start_search).pack(side=tk.LEFT, padx=5)
        ttk.Button(search_frame, text="停止搜索", command=self.stop_search).pack(side=tk.LEFT, padx=5)
        
        # 進度框架
        progress_frame = ttk.Frame(right_frame, padding="10")
        progress_frame.pack(fill=tk.X)
        
        ttk.Label(progress_frame, text="總體進度:").pack(side=tk.LEFT, padx=5)
        self.overall_progress = ttk.Progressbar(progress_frame, orient="horizontal", length=100, mode="determinate")
        self.overall_progress.pack(side=tk.LEFT, padx=5, fill=tk.X, expand=True)
        
        # 階段進度框架
        stage_frame = ttk.Frame(right_frame, padding="10")
        stage_frame.pack(fill=tk.X)
        
        ttk.Label(stage_frame, text="當(dāng)前階段:").pack(side=tk.LEFT, padx=5)
        self.stage_label = ttk.Label(stage_frame, text="準(zhǔn)備就緒")
        self.stage_label.pack(side=tk.LEFT, padx=5)
        
        self.stage_progress = ttk.Progressbar(stage_frame, orient="horizontal", length=100, mode="determinate")
        self.stage_progress.pack(side=tk.LEFT, padx=5, fill=tk.X, expand=True)
        
        # 狀態(tài)框架
        status_frame = ttk.Frame(right_frame, padding="10")
        status_frame.pack(fill=tk.X)
        
        self.status_label = ttk.Label(status_frame, text="等待開始搜索...")
        self.status_label.pack(anchor=tk.W)
        
        # 結(jié)果框架
        result_frame = ttk.LabelFrame(right_frame, text="搜索結(jié)果", padding="10")
        result_frame.pack(fill=tk.BOTH, expand=True)
        
        self.result_display = scrolledtext.ScrolledText(result_frame, wrap=tk.WORD)
        self.result_display.pack(fill=tk.BOTH, expand=True, pady=5)
        # 配置不同內(nèi)容的樣式標(biāo)簽
        self.result_display.tag_configure("filename", foreground="blue", font=("SimHei", 10, "bold"))
        self.result_display.tag_configure("linenum", foreground="green", font=("SimHei", 10))
        self.result_display.tag_configure("match", background="yellow")
        self.result_display.tag_configure("separator", foreground="gray")
        
        # 統(tǒng)計信息框架
        stats_frame = ttk.Frame(right_frame, padding="10")
        stats_frame.pack(fill=tk.X)
        
        self.stats_label = ttk.Label(stats_frame, text="文件總數(shù): 0 | 已處理: 0 | 匹配: 0")
        self.stats_label.pack(anchor=tk.W)

    def init_file_type_filters(self):
        """初始化文件類型篩選狀態(tài)（默認(rèn)不選中任何類型）"""
        pass  # 保持默認(rèn)不選中

    def toggle_category(self, category, init=False):
        """切換整個類別的選中狀態(tài)，同步子項"""
        target_state = self.category_vars[category].get()
        for ext in FILE_TYPE_CATEGORIES[category]:
            self.extension_vars[ext].set(target_state)
        if not init:
            self.update_included_extensions()

    def update_category_state(self, category):
        """根據(jù)子項狀態(tài)更新類別復(fù)選框狀態(tài)"""
        extensions = FILE_TYPE_CATEGORIES[category]
        checked_count = sum(1 for ext in extensions if self.extension_vars[ext].get())
        
        # 全選則勾選類別，否則不勾選
        self.category_vars[category].set(checked_count == len(extensions))
        self.update_included_extensions()

    def update_included_extensions(self):
        """更新當(dāng)前選中的要包含的擴展名集合"""
        self.included_extensions = {
            ext for ext, var in self.extension_vars.items() if var.get()
        }

    def add_custom_extension(self):
        """添加自定義擴展名"""
        ext = simpledialog.askstring("添加自定義擴展名", "請輸入擴展名（帶點，如 .log）:")
        if not ext:
            return
            
        # 格式化擴展名
        if not ext.startswith('.'):
            ext = '.' + ext
        ext = ext.lower()
        
        # 檢查是否已存在
        if ext in self.extension_vars:
            messagebox.showinfo("提示", f"擴展名 {ext} 已存在")
            return
            
        # 添加到"其他文件"類別（如果不存在則創(chuàng)建）
        if "其他文件" not in FILE_TYPE_CATEGORIES:
            FILE_TYPE_CATEGORIES["其他文件"] = []
            # 創(chuàng)建類別復(fù)選框
            cat_var = tk.BooleanVar()
            self.category_vars["其他文件"] = cat_var
            cat_check = ttk.Checkbutton(
                self.filter_scrollable_frame,
                text="其他文件", 
                variable=cat_var,
                command=partial(self.toggle_category, "其他文件")
            )
            cat_check.pack(anchor=tk.W, pady=5)
            
            # 創(chuàng)建擴展架子框架
            self.other_files_frame = ttk.Frame(self.filter_scrollable_frame)
            self.other_files_frame.pack(anchor=tk.W, padx=20)
        
        # 添加到類別列表
        FILE_TYPE_CATEGORIES["其他文件"].append(ext)
        
        # 創(chuàng)建復(fù)選框
        ext_var = tk.BooleanVar()
        self.extension_vars[ext] = ext_var
        ext_check = ttk.Checkbutton(
            self.other_files_frame,
            text=ext, 
            variable=ext_var,
            command=partial(self.update_category_state, "其他文件")
        )
        ext_check.pack(side=tk.LEFT, padx=5, pady=2)
        
        messagebox.showinfo("成功", f"已添加自定義擴展名 {ext}")

    def browse_path(self):
        path = filedialog.askdirectory()
        if path:
            self.search_path.set(path)

    def start_search(self):
        if self.is_searching:
            messagebox.showinfo("提示", "正在搜索中，請先停止當(dāng)前搜索")
            return
            
        path = self.search_path.get()
        keyword = self.keyword.get()
        
        if not path:
            messagebox.showerror("錯誤", "請選擇搜索路徑")
            return
            
        if not keyword:
            messagebox.showerror("錯誤", "請輸入搜索關(guān)鍵詞")
            return
            
        if not os.path.exists(path):
            messagebox.showerror("錯誤", "所選路徑不存在")
            return
        
        # 檢查是否選擇了文件類型
        self.update_included_extensions()
        if not self.included_extensions:
            messagebox.showerror("錯誤", "請至少選擇一種文件類型")
            return
            
        # 初始化搜索狀態(tài)
        self.is_searching = True
        self.stop_event.clear()  # 重置停止事件
        self.result_display.delete(1.0, tk.END)
        self.total_files = 0
        self.processed_files = 0
        self.matched_files = 0
        self.update_stats()
        
        # 創(chuàng)建隊列
        self.progress_queue = self.manager.Queue()
        self.result_queue = self.manager.Queue()
        
        # 啟動搜索進程
        self.search_process = threading.Thread(
            target=self.perform_search,
            args=(path, keyword, self.search_type.get())
        )
        self.search_process.daemon = True
        self.search_process.start()
        
        # 啟動進度和結(jié)果處理線程
        self.root.after(100, self.process_queue_updates)

    def stop_search(self):
        if self.is_searching and self.search_process:
            self.stop_event.set()
            self.status_label.config(text="正在停止搜索...")

    def perform_search(self, root_path, keyword, search_type):
        try:
            # 階段1: 掃描所有文件并過濾（只保留選中類型）
            self.progress_queue.put(("stage", "掃描文件并過濾（只保留選中類型）", 0))
            
            # 獲取所有文件列表
            all_files = []
            for dirpath, _, filenames in os.walk(root_path):
                if self.stop_event.is_set():
                    self.progress_queue.put(("done", "搜索已停止"))
                    return
                    
                for filename in filenames:
                    file_path = os.path.join(dirpath, filename)
                    # 檢查文件擴展名是否在選中的類型中
                    ext = os.path.splitext(filename)[1].lower()
                    if ext in self.included_extensions:  # 只搜索選中類型
                        all_files.append(file_path)
                
                # 更新進度
                total_scanned = len(all_files)
                progress = min(100, int((total_scanned / (total_scanned + 1)) * 100))  # 避免除以0
                self.progress_queue.put(("stage", "掃描文件并過濾（只保留選中類型）", progress))
            
            self.total_files = len(all_files)
            self.progress_queue.put(("total_files", self.total_files))
            self.progress_queue.put(("stage", "掃描文件并過濾完成", 100))
            
            if self.total_files == 0:
                self.progress_queue.put(("done", "未找到符合條件的文件類型"))
                return
                
            # 階段2: 搜索文件
            self.progress_queue.put(("stage", f"正在{('搜索文件名' if search_type == 'name' else '搜索文件內(nèi)容')}", 0))
            
            # 使用多進程搜索
            num_processes = min(mp.cpu_count(), self.total_files)
            chunk_size = max(1, self.total_files // num_processes)
            
            with Pool(processes=num_processes) as pool:
                # 部分應(yīng)用函數(shù)參數(shù)
                if search_type == 'name':
                    search_func = partial(
                        search_filename, 
                        keyword=keyword, 
                        stop_event=self.stop_event,
                        progress_queue=self.progress_queue
                    )
                else:
                    search_func = partial(
                        search_file_content, 
                        keyword=keyword, 
                        stop_event=self.stop_event,
                        progress_queue=self.progress_queue
                    )
                
                # 異步處理文件列表
                results = []
                for i in range(0, self.total_files, chunk_size):
                    chunk = all_files[i:i+chunk_size]
                    results.append(pool.apply_async(search_func, args=(chunk,)))
                
                # 收集結(jié)果
                for result in results:
                    matched_in_chunk = result.get()
                    for file_path, matches in matched_in_chunk:
                        self.result_queue.put((file_path, matches))
                    
                    if self.stop_event.is_set():
                        pool.terminate()
                        self.progress_queue.put(("done", "搜索已停止"))
                        return
            
            self.progress_queue.put(("done", "搜索完成"))
            
        except Exception as e:
            self.progress_queue.put(("error", str(e)))

    def process_queue_updates(self):
        if not self.is_searching:
            return
            
        # 處理進度更新
        try:
            while not self.progress_queue.empty():
                item = self.progress_queue.get_nowait()
                if item[0] == "stage":
                    stage_name, progress = item[1], item[2]
                    self.stage_label.config(text=stage_name)
                    self.stage_progress["value"] = progress
                elif item[0] == "progress":
                    self.processed_files = item[1]
                    overall_progress = (self.processed_files / self.total_files) * 100 if self.total_files > 0 else 0
                    self.overall_progress["value"] = overall_progress
                    self.status_label.config(text=f"正在處理: {item[2]}")
                    self.update_stats()
                elif item[0] == "total_files":
                    self.total_files = item[1]
                    self.update_stats()
                elif item[0] == "done":
                    self.status_label.config(text=item[1])
                    self.is_searching = False
                elif item[0] == "error":
                    messagebox.showerror("錯誤", f"搜索過程中發(fā)生錯誤: {item[1]}")
                    self.is_searching = False
        
        except queue.Empty:
            pass
        
        # 處理結(jié)果更新
        try:
            while not self.result_queue.empty():
                item = self.result_queue.get_nowait()
                file_path, matches = item
                self.matched_files += 1
                self.display_result(file_path, matches)
                self.update_stats()
        except queue.Empty:
            pass
        
        # 繼續(xù)檢查隊列或結(jié)束
        if self.is_searching:
            self.root.after(100, self.process_queue_updates)
        else:
            self.overall_progress["value"] = 100
            self.stage_progress["value"] = 100

    def display_result(self, file_path, matches):
        # 插入分割線
        self.result_display.insert(tk.END, "----------------------------\n", "separator")
        
        # 顯示文件名（藍(lán)色加粗）
        self.result_display.insert(tk.END, f"文件路徑: {file_path}\n", "filename")
        
        # 如果有匹配內(nèi)容，顯示行數(shù)和匹配片段
        if matches and self.search_type.get() == "content":
            self.result_display.insert(tk.END, "  匹配內(nèi)容:\n")
            for line_num, line_content in matches[:5]:  # 只顯示前5個匹配
                # 顯示行數(shù)（綠色）
                self.result_display.insert(tk.END, f"    第{line_num}行: ", "linenum")
                start_pos = self.result_display.index(tk.END)
                self.result_display.insert(tk.END, line_content + "\n")
                end_pos = self.result_display.index(tk.END)
                
                # 標(biāo)記匹配的關(guān)鍵詞（黃色高亮）
                start = start_pos
                keyword = self.keyword.get()
                while True:
                    start = self.result_display.search(keyword, start, end_pos, nocase=True)
                    if not start:
                        break
                    line, col = map(int, start.split('.'))
                    end = f"{line}.{col + len(keyword)}"
                    self.result_display.tag_add("match", start, end)
                    start = end
        
        self.result_display.see(tk.END)

    def update_stats(self):
        self.stats_label.config(
            text=f"文件總數(shù): {self.total_files} | 已處理: {self.processed_files} | 匹配: {self.matched_files}"
        )

# 多進程輔助函數(shù) - 搜索文件名
def search_filename(files, keyword, stop_event, progress_queue):
    matched = []
    pattern = re.compile(re.escape(keyword), re.IGNORECASE)
    
    for i, file_path in enumerate(files):
        if stop_event.is_set():
            return []
            
        filename = os.path.basename(file_path)
        if pattern.search(filename):
            matched.append((file_path, []))
            
        # 更新進度
        if i % 10 == 0 or i == len(files) - 1:
            progress_queue.put(("progress", i + 1, file_path))
    
    return matched

# 多進程輔助函數(shù) - 搜索文件內(nèi)容
def search_file_content(files, keyword, stop_event, progress_queue):
    matched = []
    pattern = re.compile(re.escape(keyword), re.IGNORECASE)
    
    for i, file_path in enumerate(files):
        if stop_event.is_set():
            return []
            
        try:
            # 嘗試確定文件類型，跳過二進制文件
            mime_type, _ = mimetypes.guess_type(file_path)
            if mime_type and mime_type.startswith(('image/', 'audio/', 'video/')):
                continue
                
            # 嘗試以文本方式打開文件
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                matches_in_file = []
                for line_num, line in enumerate(f, 1):
                    if pattern.search(line):
                        # 截取匹配行的上下文
                        start = max(0, line.find(keyword) - 30)
                        end = min(len(line), line.find(keyword) + len(keyword) + 30)
                        snippet = line[start:end].replace('\n', ' ').replace('\r', '')
                        matches_in_file.append((line_num, snippet))
                        
                        # 限制每個文件最多記錄10個匹配
                        if len(matches_in_file) >= 10:
                            break
                
                if matches_in_file:
                    matched.append((file_path, matches_in_file))
                    
        except (IOError, UnicodeDecodeError):
            # 無法讀取的文件跳過
            pass
            
        # 更新進度
        if i % 10 == 0 or i == len(files) - 1:
            progress_queue.put(("progress", i + 1, file_path))
    
    return matched

if __name__ == "__main__":
    # 在Windows上運行多進程需要保護主模塊
    if os.name == 'nt':
        mp.set_start_method('spawn')
    
    root = tk.Tk()
    app = FileSearchApp(root)
    root.mainloop()