CamFinder/src/coordinator.rs

use crate::{
    config::Config,
    db::Database,
    models::{ChunkRecord, PendingRange},
    scan::{build_pending_ranges, load_targets},
    terminal::{ScreenState, TerminalUi},
};
use anyhow::{anyhow, Context, Result};
use std::{
    collections::VecDeque,
    net::Ipv4Addr,
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc, Mutex,
    },
    time::Instant,
};
use tokio::{
    io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
    net::{TcpListener, TcpStream},
    task::JoinSet,
};

pub async fn run_coordinator(
    config: Config,
    database: Arc<Database>,
    ui: Arc<TerminalUi>,
    interrupted: Arc<AtomicBool>,
) -> Result<()> {
    let targets = Arc::new(load_targets(&config.cidrs)?);
    let target_count = targets.len();
    let chunk_records = database.load_chunk_records().await?;
    let (tested_before_run, _open_before_run) = database.summarize_chunks().await?;
    let pending_ranges = build_pending_ranges(target_count, config.chunk_size, &chunk_records);
    let pending_count: usize = pending_ranges.iter().map(|range| range.len()).sum();

    if ui.is_enabled() {
        ui.print_banner();
        ui.render(&ScreenState {
            current_ip: if pending_count > 0 {
                targets[pending_ranges[0].start].to_string()
            } else {
                "concluído".to_string()
            },
            tested: tested_before_run,
            total: target_count,
        });
    } else {
        println!(
            "Coordinator ready on {} for {} hosts across {} CIDRs. Pending {}, already tested {}.",
            config.coordinator_bind,
            target_count,
            config.cidrs.len(),
            pending_count,
            tested_before_run
        );
    }

    if pending_ranges.is_empty() {
        if ui.is_enabled() {
            ui.render(&ScreenState {
                current_ip: "concluído".to_string(),
                tested: target_count,
                total: target_count,
            });
            ui.finish();
        }
        return Ok(());
    }

    let listener = TcpListener::bind(config.coordinator_bind)
        .await
        .with_context(|| format!("failed to bind {}", config.coordinator_bind))?;
    let queue = Arc::new(Mutex::new(VecDeque::from(pending_ranges)));
    let tested_this_run = Arc::new(AtomicUsize::new(0));
    let open_this_run = Arc::new(AtomicUsize::new(0));
    let start = Instant::now();
    let mut handlers = JoinSet::new();

    loop {
        if interrupted.load(Ordering::SeqCst) {
            break;
        }

        tokio::select! {
            accept_result = listener.accept() => {
                let (stream, peer) = accept_result?;
                let queue = Arc::clone(&queue);
                let targets = Arc::clone(&targets);
                let database = Arc::clone(&database);
                let ui = Arc::clone(&ui);
                let interrupted = Arc::clone(&interrupted);
                let tested_this_run = Arc::clone(&tested_this_run);
                let open_this_run = Arc::clone(&open_this_run);
                let config = config.clone();

                handlers.spawn(async move {
                    handle_worker(
                        stream,
                        peer,
                        queue,
                        targets,
                        database,
                        ui,
                        interrupted,
                        tested_before_run,
                        tested_this_run,
                        open_this_run,
                        config,
                    )
                    .await
                });
            }
            maybe_finished = handlers.join_next(), if !handlers.is_empty() => {
                if let Some(join_result) = maybe_finished {
                    log_worker_result(join_result);
                }
            }
        }

        if queue.lock().expect("queue mutex poisoned").is_empty() && handlers.is_empty() {
            break;
        }
    }

    while let Some(join_result) = handlers.join_next().await {
        log_worker_result(join_result);
    }

    let duration = start.elapsed().as_secs_f64();
    let total_open_count = database.count_open_endpoints(config.port).await?;

    if interrupted.load(Ordering::SeqCst) {
        if ui.is_enabled() {
            ui.render(&ScreenState {
                current_ip: "interrompido".to_string(),
                tested: tested_before_run + tested_this_run.load(Ordering::SeqCst),
                total: target_count,
            });
            ui.finish();
        } else {
            println!("Interrupted after {:.2}s.", duration);
            println!(
                "Hosts tested in this run: {}",
                tested_this_run.load(Ordering::SeqCst)
            );
            println!(
                "Open endpoints found: {}",
                open_this_run.load(Ordering::SeqCst)
            );
        }
        return Ok(());
    }

    if ui.is_enabled() {
        ui.render(&ScreenState {
            current_ip: "concluído".to_string(),
            tested: target_count,
            total: target_count,
        });
        ui.finish();
    } else {
        println!("Coordinator finished in {:.2}s", duration);
        println!(
            "Hosts tested in this run: {}",
            tested_this_run.load(Ordering::SeqCst)
        );
        println!(
            "Open endpoints found in this run: {}",
            open_this_run.load(Ordering::SeqCst)
        );
        println!("Total open endpoints in DB: {}", total_open_count);
    }

    Ok(())
}

fn log_worker_result(result: std::result::Result<Result<()>, tokio::task::JoinError>) {
    match result {
        Ok(Ok(())) => {}
        Ok(Err(err)) => eprintln!("worker handler ended with error: {err:#}"),
        Err(err) => eprintln!("worker task failed: {err}"),
    }
}

async fn handle_worker(
    stream: TcpStream,
    peer: std::net::SocketAddr,
    queue: Arc<Mutex<VecDeque<PendingRange>>>,
    targets: Arc<Vec<Ipv4Addr>>,
    database: Arc<Database>,
    ui: Arc<TerminalUi>,
    interrupted: Arc<AtomicBool>,
    tested_before_run: usize,
    tested_this_run: Arc<AtomicUsize>,
    open_this_run: Arc<AtomicUsize>,
    config: Config,
) -> Result<()> {
    let (read_half, mut write_half) = stream.into_split();
    let mut lines = BufReader::new(read_half).lines();

    let worker_name = match lines.next_line().await? {
        Some(line) if line.starts_with("HELLO ") => line[6..].trim().to_string(),
        Some(line) => {
            eprintln!("rejected non-worker connection from {peer}: {line}");
            return Ok(());
        }
        None => return Ok(()),
    };

    if !worker_name.is_empty() {
        println!("worker connected from {peer}: {worker_name}");
    } else {
        println!("worker connected from {peer}");
    }

    loop {
        if interrupted.load(Ordering::SeqCst) {
            let _ = write_half.write_all(b"SHUTDOWN\n").await;
            break;
        }

        let Some(range) = next_range(&queue) else {
            let _ = write_half.write_all(b"DONE\n").await;
            break;
        };

        if let Err(err) = process_range(
            &mut write_half,
            &mut lines,
            range,
            &targets,
            &database,
            &ui,
            tested_before_run,
            &tested_this_run,
            &open_this_run,
            &config,
        )
        .await
        {
            requeue_range(&queue, range);
            return Err(err);
        }
    }

    Ok(())
}

async fn process_range(
    write_half: &mut tokio::net::tcp::OwnedWriteHalf,
    lines: &mut tokio::io::Lines<BufReader<tokio::net::tcp::OwnedReadHalf>>,
    range: PendingRange,
    targets: &Arc<Vec<Ipv4Addr>>,
    database: &Arc<Database>,
    ui: &Arc<TerminalUi>,
    tested_before_run: usize,
    tested_this_run: &Arc<AtomicUsize>,
    open_this_run: &Arc<AtomicUsize>,
    config: &Config,
) -> Result<()> {
    let ips = targets[range.start..range.end].to_vec();
    let job_line = format!(
        "JOB {} {} {} {} {}\n",
        range.start,
        range.end,
        config.port,
        config.timeout_ms,
        format_ips(&ips)
    );
    write_half.write_all(job_line.as_bytes()).await?;

    let Some(result_line) = lines.next_line().await? else {
        return Err(anyhow!("worker disconnected before returning job"));
    };

    let result = parse_result_line(&result_line)?;
    if result.chunk_start != range.start || result.chunk_end != range.end {
        return Err(anyhow!(
            "worker returned mismatched range: expected {}..{}, got {}..{}",
            range.start,
            range.end,
            result.chunk_start,
            result.chunk_end
        ));
    }

    for ip in result.open_ips {
        database.upsert_open_endpoint(ip, config.port).await?;
    }

    let previous = database.load_chunk_record(range.chunk_start).await?;
    let tested_before = previous
        .as_ref()
        .map(|record| record.tested_count)
        .unwrap_or(0);
    let open_before = previous
        .as_ref()
        .map(|record| record.open_count)
        .unwrap_or(0);
    let tested_count = tested_before + result.tested;
    let open_count = open_before + result.open_count;
    database
        .save_chunk_record(&ChunkRecord {
            chunk_start: range.chunk_start,
            chunk_end: range.chunk_end,
            tested_count,
            open_count,
            is_complete: result.tested >= range.len(),
        })
        .await?;

    let tested_now = tested_this_run.fetch_add(result.tested, Ordering::SeqCst) + result.tested;
    open_this_run.fetch_add(result.open_count, Ordering::SeqCst);

    if ui.is_enabled() {
        let current_ip = ips
            .last()
            .map(|ip| ip.to_string())
            .unwrap_or_else(|| "concluído".to_string());
        ui.render(&ScreenState {
            current_ip,
            tested: tested_before_run + tested_now,
            total: targets.len(),
        });
    }

    Ok(())
}

fn requeue_range(queue: &Arc<Mutex<VecDeque<PendingRange>>>, range: PendingRange) {
    queue
        .lock()
        .expect("queue mutex poisoned")
        .push_front(range);
}

fn next_range(queue: &Arc<Mutex<VecDeque<PendingRange>>>) -> Option<PendingRange> {
    queue.lock().expect("queue mutex poisoned").pop_front()
}

struct JobResult {
    chunk_start: usize,
    chunk_end: usize,
    tested: usize,
    open_count: usize,
    open_ips: Vec<Ipv4Addr>,
}

fn parse_result_line(line: &str) -> Result<JobResult> {
    let mut parts = line.splitn(6, ' ');
    let kind = parts.next().unwrap_or("");
    if kind != "RESULT" {
        return Err(anyhow!("unexpected worker message: {line}"));
    }

    let chunk_start = parse_usize_part(parts.next(), "chunk_start", line)?;
    let chunk_end = parse_usize_part(parts.next(), "chunk_end", line)?;
    let tested = parse_usize_part(parts.next(), "tested", line)?;
    let open_count = parse_usize_part(parts.next(), "open_count", line)?;
    let ips_part = parts.next().unwrap_or("");
    let open_ips = if ips_part.trim().is_empty() {
        Vec::new()
    } else {
        ips_part
            .split(',')
            .map(str::trim)
            .filter(|entry| !entry.is_empty())
            .map(|entry| {
                entry
                    .parse::<Ipv4Addr>()
                    .with_context(|| format!("invalid IP in worker result: {entry}"))
            })
            .collect::<Result<Vec<_>>>()?
    };

    if open_ips.len() != open_count {
        return Err(anyhow!(
            "open_count mismatch in worker result: declared {}, received {}",
            open_count,
            open_ips.len()
        ));
    }

    Ok(JobResult {
        chunk_start,
        chunk_end,
        tested,
        open_count,
        open_ips,
    })
}

fn parse_usize_part(part: Option<&str>, name: &str, line: &str) -> Result<usize> {
    let value = part.ok_or_else(|| anyhow!("missing {name} in worker message: {line}"))?;
    value
        .trim()
        .parse()
        .map_err(|_| anyhow!("invalid {name} in worker message: {line}"))
}

fn format_ips(ips: &[Ipv4Addr]) -> String {
    ips.iter()
        .map(ToString::to_string)
        .collect::<Vec<_>>()
        .join(",")
}