//! Handlers for the kill-tree and kill-zombies request types.

use crate::proto::daemon::{
    DaemonRequest, DaemonResponse, KillTreeResponse, KillZombiesResponse, StatusCode, ZombieReport,
};
use sysinfo::{Pid, System};

use crate::daemon::reaper;

use super::util::error_response;
use super::DaemonState;

/// Handle a `KillTree` request by killing a process and its descendants.
pub fn handle_kill_tree(request: &DaemonRequest, state: &DaemonState) -> DaemonResponse {
    let Some(ref req) = request.kill_tree else {
        return error_response(
            request.id,
            StatusCode::InvalidArgument,
            "missing kill_tree payload".into(),
        );
    };

    let timeout = if req.timeout_seconds > 0.0 {
        req.timeout_seconds
    } else {
        3.0
    };
    let killed = kill_process_tree_impl(req.pid, timeout);

    // Unregister from registry (if tracked).
    state.registry.unregister(req.pid);

    DaemonResponse {
        request_id: request.id,
        code: StatusCode::Ok as i32,
        message: String::new(),
        kill_tree: Some(KillTreeResponse {
            processes_killed: killed,
        }),
        ..Default::default()
    }
}

/// Kill a process tree rooted at `pid`, returning the number of processes killed.
///
/// Collects all descendants via sysinfo, then kills them in reverse order
/// (children before parent) so that parent processes do not respawn children.
fn kill_process_tree_impl(pid: u32, _timeout_seconds: f64) -> u32 {
    use sysinfo::Signal;

    let mut sys = System::new();
    sys.refresh_processes();

    let target = Pid::from_u32(pid);

    // Collect the root and all descendants.
    let mut to_kill = Vec::new();
    collect_descendants(&sys, target, &mut to_kill);
    to_kill.push(target);

    // Kill in reverse order (deepest children first, root last).
    to_kill.reverse();

    let mut killed_count = 0u32;
    for &p in &to_kill {
        if let Some(proc) = sys.process(p) {
            if proc.kill_with(Signal::Kill).unwrap_or(false) {
                killed_count += 1;
            }
        }
    }
    killed_count
}

/// Recursively collect all descendant PIDs of `parent_pid`.
fn collect_descendants(sys: &System, parent_pid: Pid, result: &mut Vec<Pid>) {
    for (child_pid, child_proc) in sys.processes() {
        if child_proc.parent() == Some(parent_pid) {
            result.push(*child_pid);
            collect_descendants(sys, *child_pid, result);
        }
    }
}

/// Handle a `KillZombies` request by scanning for and optionally killing zombie processes.
pub fn handle_kill_zombies(request: &DaemonRequest, state: &DaemonState) -> DaemonResponse {
    let Some(ref req) = request.kill_zombies else {
        return error_response(
            request.id,
            StatusCode::InvalidArgument,
            "missing kill_zombies payload".into(),
        );
    };

    let zombies = reaper::scan_for_zombies(state);
    let orphan_conhosts = reaper::scan_for_orphan_conhosts();

    let mut reports: Vec<ZombieReport> = Vec::new();

    // Registry-based zombies.
    if req.dry_run {
        reports.extend(zombies.iter().map(|z| ZombieReport {
            pid: z.pid,
            command: z.command.clone(),
            reason: z.reason.clone(),
            killed: false,
        }));
        reports.extend(orphan_conhosts.iter().map(|z| ZombieReport {
            pid: z.pid,
            command: z.command.clone(),
            reason: z.reason.clone(),
            killed: false,
        }));
    } else {
        let reg_results = reaper::kill_zombies(state, &zombies);
        reports.extend(
            zombies
                .iter()
                .zip(reg_results.iter())
                .map(|(z, (_pid, killed))| ZombieReport {
                    pid: z.pid,
                    command: z.command.clone(),
                    reason: z.reason.clone(),
                    killed: *killed,
                }),
        );

        let conhost_results = reaper::kill_conhosts(&orphan_conhosts);
        reports.extend(orphan_conhosts.iter().zip(conhost_results.iter()).map(
            |(z, (_pid, killed))| ZombieReport {
                pid: z.pid,
                command: z.command.clone(),
                reason: z.reason.clone(),
                killed: *killed,
            },
        ));
    }

    DaemonResponse {
        request_id: request.id,
        code: StatusCode::Ok as i32,
        message: String::new(),
        kill_zombies: Some(KillZombiesResponse { zombies: reports }),
        ..Default::default()
    }
}
