emod-cli/src/debug/hotreload.rs

use std::{
    collections::{HashMap, HashSet},
    path::{Path, PathBuf},
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, Ordering},
    },
    thread::{self, JoinHandle},
    time::Duration,
};

#[cfg(not(windows))]
use std::{fs, time::SystemTime};

use serde_json::json;
#[cfg(not(windows))]
use walkdir::WalkDir;

use super::{
    config::ResolvedModDir,
    ipc::DebugIpcServer,
    log::{self, ConsoleColor},
};

const DEBOUNCE_MS: u64 = 100;

pub struct HotReloadTask {
    stop: Arc<AtomicBool>,

    file_thread: Option<JoinHandle<()>>,
    foreground_thread: Option<JoinHandle<()>>,
}

impl HotReloadTask {
    pub fn start(
        process_id: u32,
        mod_dirs: &[ResolvedModDir],
        ipc: DebugIpcServer,
    ) -> Option<Self> {
        let roots: Vec<PathBuf> = mod_dirs
            .iter()
            .filter(|dir| dir.hot_reload)
            .map(|dir| dir.path.clone())
            .collect();
        if roots.is_empty() {
            return None;
        }

        println!("[HotReload] 追踪目录列表：");
        for root in &roots {
            println!("  └── {}", root.display());
        }

        let stop = Arc::new(AtomicBool::new(false));
        let need_update = Arc::new(AtomicBool::new(false));
        let is_foreground = Arc::new(AtomicBool::new(false));
        let cached_paths = Arc::new(Mutex::new(HashSet::new()));

        let file_thread = {
            let stop = Arc::clone(&stop);
            let need_update = Arc::clone(&need_update);
            let is_foreground = Arc::clone(&is_foreground);
            let cached_paths = Arc::clone(&cached_paths);
            let ipc = ipc.clone();
            let roots = roots.clone();
            thread::spawn(move || {
                watch_py_files(roots, stop, need_update, is_foreground, cached_paths, ipc)
            })
        };

        let foreground_thread = {
            let stop = Arc::clone(&stop);
            let need_update = Arc::clone(&need_update);
            let is_foreground = Arc::clone(&is_foreground);
            let cached_paths = Arc::clone(&cached_paths);
            let ipc = ipc.clone();
            let roots = roots.clone();
            thread::spawn(move || {
                watch_foreground(
                    process_id,
                    roots,
                    stop,
                    need_update,
                    is_foreground,
                    cached_paths,
                    ipc,
                )
            })
        };

        Some(Self {
            stop,

            file_thread: Some(file_thread),
            foreground_thread: Some(foreground_thread),
        })
    }

    pub fn safe_exit(&mut self) {
        self.stop.store(true, Ordering::Relaxed);
        if let Some(handle) = self.file_thread.take() {
            let _ = handle.join();
        }
        if let Some(handle) = self.foreground_thread.take() {
            let _ = handle.join();
        }
    }
}

impl Drop for HotReloadTask {
    fn drop(&mut self) {
        self.safe_exit();
    }
}

fn watch_py_files(
    roots: Vec<PathBuf>,
    stop: Arc<AtomicBool>,
    need_update: Arc<AtomicBool>,
    is_foreground: Arc<AtomicBool>,
    cached_paths: Arc<Mutex<HashSet<PathBuf>>>,
    ipc: DebugIpcServer,
) {
    watch_py_files_native(roots, stop, need_update, is_foreground, cached_paths, ipc);
}

#[cfg(windows)]
fn watch_py_files_native(
    roots: Vec<PathBuf>,
    stop: Arc<AtomicBool>,
    need_update: Arc<AtomicBool>,
    is_foreground: Arc<AtomicBool>,
    cached_paths: Arc<Mutex<HashSet<PathBuf>>>,
    ipc: DebugIpcServer,
) {
    use std::{mem, os::windows::ffi::OsStrExt, ptr, time::Instant};
    use windows_sys::Win32::{
        Foundation::{HANDLE, INVALID_HANDLE_VALUE, WAIT_FAILED, WAIT_OBJECT_0, WAIT_TIMEOUT},
        Storage::FileSystem::{
            CreateFileW, FILE_ACTION_MODIFIED, FILE_FLAG_BACKUP_SEMANTICS, FILE_FLAG_OVERLAPPED,
            FILE_LIST_DIRECTORY, FILE_NOTIFY_CHANGE_LAST_WRITE, FILE_NOTIFY_INFORMATION,
            FILE_SHARE_DELETE, FILE_SHARE_READ, FILE_SHARE_WRITE, OPEN_EXISTING,
            ReadDirectoryChangesW,
        },
        System::{
            IO::{GetOverlappedResult, OVERLAPPED},
            Threading::{CreateEventW, WaitForMultipleObjects},
        },
    };

    struct WatchItem {
        dir: PathBuf,
        h_dir: HANDLE,
        event: HANDLE,
        overlapped: OVERLAPPED,
        buffer: Vec<u8>,
    }

    impl Drop for WatchItem {
        fn drop(&mut self) {
            unsafe {
                if self.h_dir != INVALID_HANDLE_VALUE && !self.h_dir.is_null() {
                    windows_sys::Win32::Foundation::CloseHandle(self.h_dir);
                }
                if !self.event.is_null() {
                    windows_sys::Win32::Foundation::CloseHandle(self.event);
                }
            }
        }
    }

    fn start_watch(item: &mut WatchItem) -> bool {
        item.overlapped = unsafe { mem::zeroed() };
        item.overlapped.hEvent = item.event;
        let mut bytes_returned = 0u32;
        unsafe {
            ReadDirectoryChangesW(
                item.h_dir,
                item.buffer.as_mut_ptr().cast(),
                item.buffer.len() as u32,
                1,
                FILE_NOTIFY_CHANGE_LAST_WRITE,
                &mut bytes_returned,
                &mut item.overlapped,
                None,
            ) != 0
        }
    }

    let mut items = Vec::new();
    for root in &roots {
        if !root.is_dir() {
            continue;
        }
        let mut wide: Vec<u16> = root.as_os_str().encode_wide().collect();
        wide.push(0);
        let event = unsafe { CreateEventW(ptr::null(), 0, 0, ptr::null()) };
        if event.is_null() {
            continue;
        }
        let h_dir = unsafe {
            CreateFileW(
                wide.as_ptr(),
                FILE_LIST_DIRECTORY,
                FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
                ptr::null(),
                OPEN_EXISTING,
                FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OVERLAPPED,
                ptr::null_mut(),
            )
        };
        if h_dir == INVALID_HANDLE_VALUE {
            unsafe {
                windows_sys::Win32::Foundation::CloseHandle(event);
            }
            continue;
        }
        items.push(WatchItem {
            dir: root.clone(),
            h_dir,
            event,
            overlapped: unsafe { mem::zeroed() },
            buffer: vec![0u8; 64 * 1024],
        });
    }

    if items.is_empty() {
        return;
    }
    if items.len() > 64 {
        log::print_colored(
            "[HotReload] 追踪目录超过 Windows WaitForMultipleObjects 限制，已忽略超出部分。",
            ConsoleColor::Yellow,
        );
        items.truncate(64);
    }

    for item in &mut items {
        start_watch(item);
    }

    let mut debounce = HashMap::<PathBuf, Instant>::new();
    while !stop.load(Ordering::Relaxed) {
        let handles: Vec<HANDLE> = items.iter().map(|item| item.event).collect();
        let result =
            unsafe { WaitForMultipleObjects(handles.len() as u32, handles.as_ptr(), 0, 100) };
        if result == WAIT_TIMEOUT {
            continue;
        }
        if result == WAIT_FAILED {
            break;
        }
        let index = (result - WAIT_OBJECT_0) as usize;
        if index >= items.len() {
            continue;
        }
        let item = &mut items[index];
        let mut bytes = 0u32;
        let ok = unsafe { GetOverlappedResult(item.h_dir, &mut item.overlapped, &mut bytes, 0) };
        if ok == 0 || bytes == 0 {
            start_watch(item);
            continue;
        }

        let mut offset = 0usize;
        while offset < bytes as usize {
            let info =
                unsafe { &*(item.buffer.as_ptr().add(offset) as *const FILE_NOTIFY_INFORMATION) };
            let name_len = (info.FileNameLength / 2) as usize;
            let name = unsafe { std::slice::from_raw_parts(info.FileName.as_ptr(), name_len) };
            let path = item.dir.join(String::from_utf16_lossy(name));
            if info.Action == FILE_ACTION_MODIFIED
                && path.extension().and_then(|ext| ext.to_str()) == Some("py")
            {
                let now = Instant::now();
                let should_trigger = debounce
                    .get(&path)
                    .map(|last| now.duration_since(*last) >= Duration::from_millis(DEBOUNCE_MS))
                    .unwrap_or(true);
                if should_trigger {
                    debounce.insert(path.clone(), now);
                    record_changed_path(
                        &roots,
                        &path,
                        &cached_paths,
                        &need_update,
                        &is_foreground,
                        &ipc,
                    );
                }
            }
            if info.NextEntryOffset == 0 {
                break;
            }
            offset += info.NextEntryOffset as usize;
        }
        start_watch(item);
    }
}

#[cfg(not(windows))]
fn watch_py_files_native(
    roots: Vec<PathBuf>,
    stop: Arc<AtomicBool>,
    need_update: Arc<AtomicBool>,
    is_foreground: Arc<AtomicBool>,
    cached_paths: Arc<Mutex<HashSet<PathBuf>>>,
    ipc: DebugIpcServer,
) {
    let mut last_seen = snapshot_py_files(&roots);
    let mut debounce = HashMap::<PathBuf, SystemTime>::new();

    while !stop.load(Ordering::Relaxed) {
        thread::sleep(Duration::from_millis(DEBOUNCE_MS));
        let now = SystemTime::now();
        let current = snapshot_py_files(&roots);
        for (path, modified) in &current {
            let changed = match last_seen.get(path) {
                Some(previous) => modified > previous,
                None => false,
            };
            if !changed {
                continue;
            }
            if let Some(last) = debounce.get(path) {
                if now.duration_since(*last).unwrap_or_default()
                    < Duration::from_millis(DEBOUNCE_MS)
                {
                    continue;
                }
            }
            debounce.insert(path.clone(), now);
            record_changed_path(
                &roots,
                path,
                &cached_paths,
                &need_update,
                &is_foreground,
                &ipc,
            );
        }
        last_seen = current;
    }
}

fn record_changed_path(
    roots: &[PathBuf],
    path: &Path,
    cached_paths: &Arc<Mutex<HashSet<PathBuf>>>,
    need_update: &Arc<AtomicBool>,
    is_foreground: &Arc<AtomicBool>,
    ipc: &DebugIpcServer,
) {
    log::print_colored(
        &format!("[HotReload] Detected change in: {}", path.display()),
        ConsoleColor::Yellow,
    );
    cached_paths.lock().unwrap().insert(path.to_path_buf());
    need_update.store(true, Ordering::Relaxed);
    if is_foreground.load(Ordering::Relaxed) {
        trigger_reload(roots, cached_paths, need_update, ipc);
    }
}

fn watch_foreground(
    process_id: u32,
    roots: Vec<PathBuf>,
    stop: Arc<AtomicBool>,
    need_update: Arc<AtomicBool>,
    is_foreground: Arc<AtomicBool>,
    cached_paths: Arc<Mutex<HashSet<PathBuf>>>,
    ipc: DebugIpcServer,
) {
    let mut last_state = false;
    while !stop.load(Ordering::Relaxed) {
        let current = foreground_process_id() == Some(process_id);
        is_foreground.store(current, Ordering::Relaxed);
        if current && !last_state && need_update.load(Ordering::Relaxed) {
            trigger_reload(&roots, &cached_paths, &need_update, &ipc);
        }
        last_state = current;
        thread::sleep(Duration::from_millis(80));
    }
}

fn trigger_reload(
    roots: &[PathBuf],
    cached_paths: &Arc<Mutex<HashSet<PathBuf>>>,
    need_update: &Arc<AtomicBool>,
    ipc: &DebugIpcServer,
) {
    let paths: Vec<PathBuf> = {
        let mut guard = cached_paths.lock().unwrap();
        guard.drain().collect()
    };
    let modules: Vec<String> = paths
        .iter()
        .filter_map(|path| py_path_to_module_name(path, roots))
        .collect();
    need_update.store(false, Ordering::Relaxed);
    if modules.is_empty() {
        return;
    }
    log::print_colored(
        "[HotReload] 检测到修改，已触发热更新。",
        ConsoleColor::Yellow,
    );
    let payload = json!(modules).to_string();
    let _ = ipc.send_message(2, payload.as_bytes());
}

#[cfg(not(windows))]
fn snapshot_py_files(roots: &[PathBuf]) -> HashMap<PathBuf, SystemTime> {
    let mut files = HashMap::new();
    for root in roots {
        if !root.is_dir() {
            continue;
        }
        for entry in WalkDir::new(root).into_iter().filter_map(Result::ok) {
            let path = entry.path();
            if !path.is_file() || path.extension().and_then(|ext| ext.to_str()) != Some("py") {
                continue;
            }
            if let Ok(metadata) = fs::metadata(path) {
                if let Ok(modified) = metadata.modified() {
                    files.insert(path.to_path_buf(), modified);
                }
            }
        }
    }
    files
}

pub fn py_path_to_module_name(file_path: &Path, mod_roots: &[PathBuf]) -> Option<String> {
    if file_path.extension().and_then(|ext| ext.to_str()) != Some("py") {
        return None;
    }
    let mut cur = file_path.parent()?;
    loop {
        if cur.join("manifest.json").is_file() || cur.join("pack_manifest.json").is_file() {
            let rel = file_path.strip_prefix(cur).ok()?;
            let mut parts: Vec<String> = rel
                .iter()
                .map(|part| part.to_string_lossy().into_owned())
                .collect();
            let last = parts.last_mut()?;
            if !last.ends_with(".py") {
                return None;
            }
            last.truncate(last.len() - 3);
            return Some(parts.join("."));
        }
        if mod_roots.iter().any(|root| same_path(root, cur)) {
            return None;
        }
        cur = cur.parent()?;
    }
}

fn same_path(left: &Path, right: &Path) -> bool {
    left == right || (left.canonicalize().ok().as_deref() == right.canonicalize().ok().as_deref())
}

#[cfg(windows)]
fn foreground_process_id() -> Option<u32> {
    use windows_sys::Win32::UI::WindowsAndMessaging::{
        GetForegroundWindow, GetWindowThreadProcessId,
    };
    unsafe {
        let hwnd = GetForegroundWindow();
        if hwnd.is_null() {
            return None;
        }
        let mut pid = 0u32;
        GetWindowThreadProcessId(hwnd, &mut pid);
        if pid == 0 { None } else { Some(pid) }
    }
}

#[cfg(not(windows))]
fn foreground_process_id() -> Option<u32> {
    None
}

#[cfg(test)]
mod tests {
    use std::fs;

    use super::py_path_to_module_name;

    #[test]
    fn maps_py_file_to_module_name_under_manifest() {
        let root = std::env::temp_dir().join(format!("emod-hot-{}", std::process::id()));
        let pack = root.join("behavior_pack");
        let module = pack.join("foo/bar.py");
        fs::create_dir_all(module.parent().unwrap()).unwrap();
        fs::write(pack.join("manifest.json"), "{}").unwrap();
        fs::write(&module, "").unwrap();

        assert_eq!(
            py_path_to_module_name(&module, std::slice::from_ref(&root)),
            Some("foo.bar".to_string())
        );
        let outside = root.join("x.py");
        fs::write(&outside, "").unwrap();
        assert_eq!(
            py_path_to_module_name(&outside, std::slice::from_ref(&root)),
            None
        );

        let _ = fs::remove_dir_all(root);
    }
}