scatter-proxy 0.9.0

use std::sync::Arc;
use std::time::{Duration, Instant};

use tokio::sync::{broadcast, Semaphore};
use tokio::task::JoinSet;
use tracing::{debug, info, warn};

use crate::challenge::ChallengeResolver;
use crate::classifier::{BodyClassifier, BodyVerdict};
use crate::config::ScatterProxyConfig;
use crate::health::HealthTracker;
use crate::metrics::ThroughputTracker;
use crate::mutator::RequestMutator;
use crate::proxy::{ProxyManager, ProxyState};
use crate::rate_limit::RateLimiter;
use crate::score::{adaptive_k, compute_score};
use crate::task::{ScatterResponse, TaskEntry, TaskPool};

enum RaceOutcome {
    Response(reqwest::Response),
    RequestError(String),
    Timeout,
}

pub(crate) struct Scheduler {
    config: Arc<ScatterProxyConfig>,
    task_pool: Arc<TaskPool>,
    health: Arc<HealthTracker>,
    rate_limiter: Arc<RateLimiter>,
    proxy_manager: Arc<ProxyManager>,
    classifier: Arc<dyn BodyClassifier>,
    semaphore: Arc<Semaphore>,
    throughput: Arc<ThroughputTracker>,
    challenge_resolver: Option<Arc<dyn ChallengeResolver>>,
    request_mutator: Option<Arc<dyn RequestMutator>>,
}

impl Scheduler {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        config: Arc<ScatterProxyConfig>,
        task_pool: Arc<TaskPool>,
        health: Arc<HealthTracker>,
        rate_limiter: Arc<RateLimiter>,
        proxy_manager: Arc<ProxyManager>,
        classifier: Arc<dyn BodyClassifier>,
        semaphore: Arc<Semaphore>,
        throughput: Arc<ThroughputTracker>,
        challenge_resolver: Option<Arc<dyn ChallengeResolver>>,
        request_mutator: Option<Arc<dyn RequestMutator>>,
    ) -> Self {
        Self {
            config,
            task_pool,
            health,
            rate_limiter,
            proxy_manager,
            classifier,
            semaphore,
            throughput,
            challenge_resolver,
            request_mutator,
        }
    }

    pub async fn run(self, mut shutdown: broadcast::Receiver<()>) {
        debug!("scheduler started");
        let scheduler = Arc::new(self);
        let worker_count = scheduler.config.max_inflight.clamp(1, 32);
        let mut workers = JoinSet::new();

        for worker_id in 0..worker_count {
            let scheduler = Arc::clone(&scheduler);
            let mut worker_shutdown = shutdown.resubscribe();
            workers.spawn(async move {
                scheduler.worker_loop(worker_id, &mut worker_shutdown).await;
            });
        }

        let _ = shutdown.recv().await;
        workers.abort_all();
        while workers.join_next().await.is_some() {}
        debug!(workers = worker_count, "scheduler stopped");
    }

    async fn worker_loop(
        self: Arc<Self>,
        worker_id: usize,
        shutdown: &mut broadcast::Receiver<()>,
    ) {
        debug!(worker_id, "scheduler worker started");
        loop {
            let delay = self
                .task_pool
                .next_delayed_ready_in()
                .unwrap_or(Duration::from_millis(10));
            tokio::select! {
                _ = shutdown.recv() => break,
                _ = self.task_pool.notified() => {},
                _ = tokio::time::sleep(delay.min(Duration::from_millis(50))) => {},
            }
            if let Some(task) = self.pick_task() {
                self.task_pool.mark_dispatch();
                self.process_task(task).await;
            }
        }
        debug!(worker_id, "scheduler worker stopped");
    }

    fn is_in_cooldown(&self, proxy: &str, host: &str) -> bool {
        let consecutive_fails = self.health.get_consecutive_fails(proxy, host);
        let threshold = self.config.cooldown_consecutive_fails as u32;
        if consecutive_fails < threshold {
            return false;
        }
        let exponent = (consecutive_fails - threshold).min(32);
        let factor = 2f64.powi(exponent as i32);
        let cooldown_secs = self.config.cooldown_base.as_secs_f64() * factor;
        let cooldown_secs = cooldown_secs.min(self.config.cooldown_max.as_secs_f64());
        self.health.seconds_since_last_access(proxy, host) < cooldown_secs
    }

    fn pick_task(&self) -> Option<TaskEntry> {
        self.task_pool.pick_next(&std::collections::HashSet::new())
    }

    /// Mutate a cloned request for a specific proxy if a mutator is configured.
    fn apply_mutator(&self, proxy_url: &str, req: reqwest::Request) -> reqwest::Request {
        if let Some(m) = &self.request_mutator {
            m.mutate(proxy_url, req)
        } else {
            req
        }
    }

    /// Increment the per-proxy request count and retire the node if the limit is reached.
    fn track_request(&self, proxy_url: &str) {
        if let Some(max) = self.config.max_requests_per_proxy {
            let count = self.proxy_manager.increment_request_count(proxy_url);
            if count >= max {
                self.proxy_manager.set_state(proxy_url, ProxyState::Retired);
                info!(
                    proxy = %proxy_url,
                    count,
                    limit = max,
                    "proxy retired: request limit reached"
                );
            }
        }
    }

    /// Attempt to solve a WAF challenge on `proxy_url` and retry the task.
    ///
    /// Returns `true` if the retry succeeded (task is done), `false` otherwise.
    async fn handle_challenge(&self, proxy_url: &str, body: &[u8], task: &mut TaskEntry) -> bool {
        let resolver = match &self.challenge_resolver {
            Some(r) => Arc::clone(r),
            None => return false,
        };

        // Serialize challenge solves per proxy node.
        let sem = self.proxy_manager.challenge_lock(proxy_url);
        let _permit = match sem.try_acquire() {
            Ok(p) => p,
            Err(_) => {
                // Another coroutine is already solving for this proxy; skip.
                return false;
            }
        };

        debug!(proxy = %proxy_url, "attempting WAF challenge solve");
        let cookie_val = resolver.resolve(proxy_url, body).await;

        let Some(cookie_val) = cookie_val else {
            debug!(proxy = %proxy_url, "challenge solve failed; marking proxy dead");
            self.proxy_manager.set_state(proxy_url, ProxyState::Dead);
            return false;
        };

        // Store the resolved cookie and retry once with the same proxy.
        self.proxy_manager.set_cookie(proxy_url, cookie_val);

        let req = match task.request.try_clone() {
            Some(r) => r,
            None => return false,
        };
        let req = self.apply_mutator(proxy_url, req);

        let client = match self.proxy_manager.get_client(proxy_url) {
            Ok(c) => c,
            Err(_) => return false,
        };

        let proxy_timeout = self.config.proxy_timeout;
        let start = Instant::now();
        let outcome = match tokio::time::timeout(proxy_timeout, client.execute(req)).await {
            Ok(Ok(response)) => response,
            Ok(Err(e)) => {
                debug!(proxy = %proxy_url, error = %e, "challenge retry request error");
                self.health.record_failure(proxy_url, &task.host);
                return false;
            }
            Err(_) => {
                debug!(proxy = %proxy_url, "challenge retry timed out");
                self.health.record_failure(proxy_url, &task.host);
                return false;
            }
        };
        let latency_ms = start.elapsed().as_secs_f64() * 1000.0;

        let status = outcome.status();
        let headers = outcome.headers().clone();
        let retry_body = outcome.bytes().await.unwrap_or_default();
        let verdict = self.classifier.classify(status, &headers, &retry_body);

        if verdict == BodyVerdict::Success {
            self.health
                .record_success(proxy_url, &task.host, latency_ms);
            if let Some(tx) = task.result_tx.take() {
                let _ = tx.send(ScatterResponse {
                    status,
                    headers,
                    body: retry_body,
                });
            }
            self.task_pool.mark_completed();
            self.throughput.record();
            true
        } else {
            self.health.record_failure(proxy_url, &task.host);
            false
        }
    }

    async fn process_task(&self, mut task: TaskEntry) {
        let host = task.host.clone();
        let active_proxies = self.proxy_manager.get_active_proxies();

        let mut available = Vec::new();
        let mut skipped_rate_limit = 0u64;
        let mut skipped_cooldown = 0u64;

        for proxy in active_proxies {
            if !self.rate_limiter.is_available(&proxy, &host) {
                skipped_rate_limit += 1;
                continue;
            }
            if self.is_in_cooldown(&proxy, &host) {
                skipped_cooldown += 1;
                continue;
            }
            available.push(proxy);
        }

        for _ in 0..skipped_rate_limit {
            self.task_pool.mark_skipped_rate_limit();
        }
        for _ in 0..skipped_cooldown {
            self.task_pool.mark_skipped_cooldown();
        }

        if available.is_empty() {
            self.task_pool.mark_zero_available();
            warn!(
                host = %host,
                attempt = task.attempts + 1,
                skipped_rate_limit,
                skipped_cooldown,
                "no proxy currently available; delaying task"
            );
            task.attempts += 1;
            task.last_error = format!(
                "zero available proxies (rate_limit_skips={skipped_rate_limit}, cooldown_skips={skipped_cooldown})"
            );
            self.task_pool
                .push_delayed(task, Duration::from_millis(100));
            return;
        }

        let has_untested = available
            .iter()
            .any(|p| self.health.total_samples_for_proxy(p) == 0);
        let avg_success_rate = self.health.avg_success_rate_for_host(&host);
        let k = if has_untested {
            available
                .len()
                .min(self.config.max_concurrent_per_request.max(3))
        } else {
            adaptive_k(
                available.len(),
                avg_success_rate,
                self.config.max_concurrent_per_request,
            )
        }
        .max(1);

        let mut scored: Vec<(String, f64)> = available
            .iter()
            .map(|proxy| {
                let s = compute_score(&self.health, proxy, &host);
                (proxy.clone(), s)
            })
            .collect();
        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        let candidates: Vec<String> = scored.into_iter().take(k).map(|(proxy, _)| proxy).collect();

        let mut join_set: JoinSet<(String, RaceOutcome, f64)> = JoinSet::new();
        let mut actual_candidates: Vec<String> = Vec::new();
        let mut skipped_no_permit = 0u64;

        for proxy_url in &candidates {
            let req = match task.request.try_clone() {
                Some(r) => r,
                None => {
                    self.task_pool.mark_failed();
                    return;
                }
            };

            let req = self.apply_mutator(proxy_url, req);

            let client = match self.proxy_manager.get_client(proxy_url) {
                Ok(c) => c,
                Err(_) => continue,
            };

            let permit = match self.semaphore.clone().acquire_owned().await {
                Ok(p) => p,
                Err(_) => {
                    skipped_no_permit += 1;
                    continue;
                }
            };

            self.rate_limiter.mark(proxy_url, &host);
            let proxy_timeout = self.config.proxy_timeout;
            let proxy_url_owned = proxy_url.clone();
            actual_candidates.push(proxy_url.clone());

            join_set.spawn(async move {
                let start = Instant::now();
                let outcome = match tokio::time::timeout(proxy_timeout, client.execute(req)).await {
                    Ok(Ok(response)) => RaceOutcome::Response(response),
                    Ok(Err(e)) => RaceOutcome::RequestError(e.to_string()),
                    Err(_) => RaceOutcome::Timeout,
                };
                let latency_ms = start.elapsed().as_secs_f64() * 1000.0;
                drop(permit);
                (proxy_url_owned, outcome, latency_ms)
            });
        }

        for _ in 0..skipped_no_permit {
            self.task_pool.mark_skipped_no_permit();
        }

        if join_set.is_empty() {
            warn!(
                host = %host,
                attempt = task.attempts + 1,
                candidate_count = candidates.len(),
                skipped_no_permit,
                "no candidate request launched; requeueing task"
            );
            task.attempts += 1;
            task.last_error =
                format!("no candidate launched (skipped_no_permit={skipped_no_permit})");
            self.task_pool.push_delayed(task, Duration::from_millis(25));
            return;
        }

        let mut success = false;
        let mut last_error = String::new();

        'race: while let Some(result) = join_set.join_next().await {
            let (proxy_url, outcome, latency_ms) = match result {
                Ok(v) => v,
                Err(e) => {
                    last_error = format!("task join error: {e}");
                    continue;
                }
            };

            match outcome {
                RaceOutcome::Response(response) => {
                    let status = response.status();
                    let headers = response.headers().clone();
                    let body = response.bytes().await.unwrap_or_default();
                    let verdict = self.classifier.classify(status, &headers, &body);

                    match verdict {
                        BodyVerdict::Success => {
                            self.health.record_success(&proxy_url, &host, latency_ms);
                            join_set.abort_all();
                            debug!(
                                host = %host,
                                proxy = %proxy_url,
                                attempt = task.attempts + 1,
                                latency_ms = latency_ms as u64,
                                "task completed"
                            );

                            if let Some(tx) = task.result_tx.take() {
                                let _ = tx.send(ScatterResponse {
                                    status,
                                    headers,
                                    body,
                                });
                            }
                            self.task_pool.mark_completed();
                            self.throughput.record();
                            success = true;
                            break 'race;
                        }
                        BodyVerdict::Challenge => {
                            // Abort remaining racers, solve the challenge on the same
                            // proxy node, and retry once. We break regardless of outcome
                            // so only one challenge solve runs per task attempt.
                            join_set.abort_all();
                            while join_set.join_next().await.is_some() {}
                            if self.handle_challenge(&proxy_url, &body, &mut task).await {
                                success = true;
                            } else {
                                last_error = format!("challenge unresolved (proxy={proxy_url})");
                            }
                            break 'race;
                        }
                        BodyVerdict::ProxyBlocked => {
                            self.health.record_failure(&proxy_url, &host);
                            last_error = format!("proxy blocked (status={status})");
                        }
                        BodyVerdict::TargetError => {
                            last_error = format!("target error (status={status})");
                        }
                    }
                }
                RaceOutcome::RequestError(err) => {
                    self.health.record_failure(&proxy_url, &host);
                    last_error = err;
                }
                RaceOutcome::Timeout => {
                    self.health.record_failure(&proxy_url, &host);
                    last_error = "proxy timeout".into();
                }
            }
        }

        if success {
            while join_set.join_next().await.is_some() {}
        } else {
            task.attempts += 1;
            task.last_error = last_error;
            debug!(host = %host, attempt = task.attempts, reason = %task.last_error, "task requeued");
            self.task_pool.push_back(task);
        }

        // Track request counts and evict zero-rate proxies.
        for proxy_url in &actual_candidates {
            self.track_request(proxy_url);

            let samples = self.health.total_samples_for_proxy(proxy_url);
            if samples >= self.config.eviction_min_samples as u32 {
                let global_rate = self.health.global_success_rate_for_proxy(proxy_url);
                if global_rate == 0.0 {
                    self.proxy_manager.set_state(proxy_url, ProxyState::Dead);
                    debug!(host = %host, proxy = %proxy_url, samples = samples, "proxy dead | global success_rate=0%");
                }
            }
        }
    }
}