watcher.py

"""
文件监听模块: 监控文档目录变化, 自动增量索引

使用 watchdog 库监听文件系统事件:
- 新增文件 → 索引新 chunk
- 修改文件 → 删旧 chunk + 重新索引
- 删除文件 → 删对应 chunk

BM25 索引不支持增量更新, 文件变化后整体重建.
ChromaDB 支持增量(upsert/delete).
"""

import os
import threading

from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler

from document_loader import Document, SUPPORTED_EXTENSIONS


class DocChangeHandler(FileSystemEventHandler):
    """
    文件变化事件处理器.

    watchdog 检测到文件变化后调用对应方法:
    - on_created: 新文件
    - on_modified: 文件被修改
    - on_deleted: 文件被删除

    我们不在事件回调里直接做索引(太慢会阻塞监听),
    而是记录到变化队列, 由外部定时批量处理.
    """

    def __init__(self):
        super().__init__()
        # 用 set 记录变化的文件路径, 自动去重
        # 比如一个文件快速保存了 3 次, 只处理一次
        self._changed_files = set()     # 需要重新索引的文件
        self._deleted_files = set()     # 需要删除索引的文件
        self._lock = threading.Lock()   # 线程锁, 防止并发修改

    def _is_supported(self, path: str) -> bool:
        """检查文件是否是支持的类型"""
        ext = os.path.splitext(path)[1].lower()
        return ext in SUPPORTED_EXTENSIONS

    def on_created(self, event):
        if event.is_directory:
            return
        if self._is_supported(event.src_path):
            with self._lock:
                self._changed_files.add(event.src_path)
                # 如果之前标记删除了, 取消删除(文件又回来了)
                self._deleted_files.discard(event.src_path)
            print(f"[watcher] 新文件: {os.path.basename(event.src_path)}")

    def on_modified(self, event):
        if event.is_directory:
            return
        if self._is_supported(event.src_path):
            with self._lock:
                self._changed_files.add(event.src_path)
            print(f"[watcher] 文件修改: {os.path.basename(event.src_path)}")

    def on_deleted(self, event):
        if event.is_directory:
            return
        if self._is_supported(event.src_path):
            with self._lock:
                self._deleted_files.add(event.src_path)
                # 如果之前标记要重新索引, 取消(文件已删除)
                self._changed_files.discard(event.src_path)
            print(f"[watcher] 文件删除: {os.path.basename(event.src_path)}")

    def get_and_clear_changes(self) -> tuple[set, set]:
        """
        获取并清空变化记录.

        返回:
            (changed_files, deleted_files) 两个集合
        """
        with self._lock:
            changed = set(self._changed_files)
            deleted = set(self._deleted_files)
            self._changed_files.clear()
            self._deleted_files.clear()
        return changed, deleted


class FileWatcher:
    """
    文件监听器: 启动 watchdog Observer 监控目录.

    用法:
        watcher = FileWatcher("./test_docs")
        watcher.start()
        ...
        changed, deleted = watcher.get_changes()
        ...
        watcher.stop()
    """

    def __init__(self, watch_dir: str):
        self.watch_dir = watch_dir
        self.handler = DocChangeHandler()
        self.observer = Observer()

    def start(self) -> None:
        """开始监听(后台线程)"""
        self.observer.schedule(
            self.handler,
            self.watch_dir,
            recursive=True,     # 递归监听子目录
        )
        self.observer.daemon = True     # 主程序退出时自动结束
        self.observer.start()
        print(f"[watcher] 开始监听: {self.watch_dir}")

    def stop(self) -> None:
        """停止监听"""
        self.observer.stop()
        self.observer.join()
        print("[watcher] 已停止监听")

    def get_changes(self) -> tuple[set, set]:
        """获取自上次以来的文件变化"""
        return self.handler.get_and_clear_changes()


# ============ 测试代码 ============
if __name__ == "__main__":
    import time

    print("=== 测试文件监听 ===")
    print("监听目录: F:/AI_Program/enterprise-kb/test_docs")
    print("请在 10 秒内修改/新建/删除 test_docs 下的文件...\n")

    watcher = FileWatcher("F:/AI_Program/enterprise-kb/test_docs")
    watcher.start()

    # 等 10 秒, 期间手动操作文件
    for i in range(10):
        time.sleep(1)
        changed, deleted = watcher.get_changes()
        if changed or deleted:
            print(f"  [第{i+1}秒] 变化: {[os.path.basename(f) for f in changed]}, "
                  f"删除: {[os.path.basename(f) for f in deleted]}")

    watcher.stop()
    print("\n测试结束")