scatter-proxy 0.3.0

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

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

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

/// Internal outcome of a single proxy attempt within a race.
enum RaceOutcome {
    Response(reqwest::Response),
    RequestError(String),
    Timeout,
}

/// Core scheduling loop that picks pending tasks from the [`TaskPool`], selects
/// proxies, races requests through multiple paths, and records outcomes.
///
/// Constructed via [`Scheduler::new`] and driven by [`Scheduler::run`], which
/// runs until a shutdown signal is received.
pub(crate) struct Scheduler {
    config: Arc<ScatterProxyConfig>,
    task_pool: Arc<TaskPool>,
    health: Arc<HealthTracker>,
    rate_limiter: Arc<RateLimiter>,
    circuit_breakers: Arc<CircuitBreakerManager>,
    proxy_manager: Arc<ProxyManager>,
    classifier: Arc<dyn BodyClassifier>,
    semaphore: Arc<Semaphore>,
    throughput: Arc<ThroughputTracker>,
}

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

    /// Run the scheduling loop until a shutdown signal is received.
    pub async fn run(self, mut shutdown: broadcast::Receiver<()>) {
        debug!("scheduler started");
        loop {
            tokio::select! {
                _ = shutdown.recv() => break,
                _ = self.schedule_one() => {}
            }
        }
        debug!("scheduler stopped");
    }

    /// Check if a `(proxy, host)` pair is in exponential-backoff cooldown.
    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
    }

    /// The core scheduling method — pick one task, race it through proxies, and
    /// process results.
    async fn schedule_one(&self) {
        // ── 1. Build skip set ───────────────────────────────────────────
        let cb_state = self.circuit_breakers.get_all();
        let skip_hosts: HashSet<String> = cb_state
            .iter()
            .filter(|(_, &is_open)| is_open)
            .map(|(host, _)| host.clone())
            .collect();

        // ── 2. Handle probes ────────────────────────────────────────────
        for probe_host in &skip_hosts {
            if self.circuit_breakers.should_probe(probe_host) {
                // Build a skip set that allows the probe host through.
                let mut probe_skip = skip_hosts.clone();
                probe_skip.remove(probe_host);
                if let Some(task) = self.task_pool.pick_next(&probe_skip) {
                    if task.host == *probe_host {
                        self.run_probe(task).await;
                    } else {
                        // Not the host we wanted — put it back for normal scheduling.
                        self.task_pool.push_back(task);
                    }
                }
            }
        }

        // ── 3. Pick task ────────────────────────────────────────────────
        let mut task = match self.task_pool.pick_next(&skip_hosts) {
            Some(t) => t,
            None => {
                tokio::time::sleep(Duration::from_millis(50)).await;
                return;
            }
        };

        let host = task.host.clone();

        // ── 4. Find available proxies ───────────────────────────────────
        let active_proxies = self.proxy_manager.get_active_proxies();
        let available: Vec<String> = active_proxies
            .into_iter()
            .filter(|proxy| {
                self.rate_limiter.is_available(proxy, &host) && !self.is_in_cooldown(proxy, &host)
            })
            .collect();

        // ── 5. Zero available ───────────────────────────────────────────
        if available.is_empty() {
            self.circuit_breakers.trip(&host, "zero available proxies");
            warn!(host = %host, "circuit OPEN | reason=zero available proxies");
            self.task_pool.push_back(task);
            return;
        }

        // ── 6. Compute K ───────────────────────────────────────────────
        // During cold start (any proxy untested), race more paths to collect
        // health data quickly.
        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(5))
        } else {
            adaptive_k(
                available.len(),
                avg_success_rate,
                self.config.max_concurrent_per_request,
            )
        }
        .max(1);

        // ── 7. Select top-K by score ────────────────────────────────────
        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();

        // ── 8. Race ────────────────────────────────────────────────────
        let mut join_set: JoinSet<(String, RaceOutcome, f64)> = JoinSet::new();
        let mut actual_candidates: Vec<String> = Vec::new();

        for proxy_url in &candidates {
            // Clone the request — if the body isn't cloneable we cannot retry
            // this task at all: release the slot and let the handle resolve to
            // 502 via channel close.
            let req = match task.request.try_clone() {
                Some(r) => r,
                None => {
                    self.task_pool.mark_failed();
                    return; // drops task → tx closes → caller gets 502
                }
            };

            // Obtain a reqwest::Client configured for this proxy.
            let client = match self.proxy_manager.get_client(proxy_url) {
                Ok(c) => c,
                Err(_) => continue,
            };

            // Acquire a global in-flight permit.
            let permit = match self.semaphore.clone().try_acquire_owned() {
                Ok(p) => p,
                Err(_) => continue,
            };

            // Mark rate-limiter so no other task reuses this (proxy, host) too soon.
            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)
            });
        }

        // All candidates were skipped (client error / no permits) — requeue.
        if join_set.is_empty() {
            self.task_pool.push_back(task);
            return;
        }

        // ── 9. Process results ─────────────────────────────────────────
        let mut success = false;
        let mut last_error = String::new();

        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);
                            self.circuit_breakers.record_success(&host);

                            // Cancel every other in-flight attempt.
                            join_set.abort_all();

                            if let Some(tx) = task.result_tx.take() {
                                let _ = tx.send(ScatterResponse {
                                    status,
                                    headers,
                                    body,
                                });
                            }
                            self.task_pool.mark_completed();
                            self.throughput.record();

                            debug!(
                                host = %host,
                                winner = %proxy_url,
                                latency_ms = latency_ms,
                                attempt = task.attempts + 1,
                                "task done"
                            );

                            success = true;
                            break;
                        }
                        BodyVerdict::ProxyBlocked => {
                            self.health.record_failure(&proxy_url, &host);
                            last_error = format!("proxy blocked (status={status})");
                        }
                        BodyVerdict::TargetError => {
                            self.circuit_breakers.record_target_error(&host);
                            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();
                }
            }
        }

        // Drain any remaining aborted tasks so the JoinSet is fully cleaned up.
        if success {
            while join_set.join_next().await.is_some() {}
        }

        // ── 10. All failed — always requeue ────────────────────────────
        if !success {
            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);
        }

        // ── 11. Eviction check ──────────────────────────────────────────
        for proxy_url in &actual_candidates {
            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!(
                        proxy = %proxy_url,
                        samples = samples,
                        "proxy dead | global success_rate=0%"
                    );
                }
            }
        }
    }

    /// Run a single-proxy probe for a circuit-broken host to check whether the
    /// target has recovered.
    async fn run_probe(&self, mut task: TaskEntry) {
        let host = task.host.clone();

        // Pick one available proxy.
        let active_proxies = self.proxy_manager.get_active_proxies();
        let proxy_url = match active_proxies
            .iter()
            .find(|p| self.rate_limiter.is_available(p, &host) && !self.is_in_cooldown(p, &host))
        {
            Some(p) => p.clone(),
            None => {
                self.task_pool.push_back(task);
                return;
            }
        };

        // Clone request — unrecoverable if body not cloneable.
        let req = match task.request.try_clone() {
            Some(r) => r,
            None => {
                self.task_pool.mark_failed();
                return; // drops task → tx closes → caller gets 502
            }
        };

        // Get client.
        let client = match self.proxy_manager.get_client(&proxy_url) {
            Ok(c) => c,
            Err(_) => {
                self.task_pool.push_back(task);
                return;
            }
        };

        // Acquire semaphore permit.
        let permit = match self.semaphore.clone().try_acquire_owned() {
            Ok(p) => p,
            Err(_) => {
                self.task_pool.push_back(task);
                return;
            }
        };

        self.rate_limiter.mark(&proxy_url, &host);

        // Execute probe request.
        let start = Instant::now();
        let outcome =
            match tokio::time::timeout(self.config.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);

        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);
                        self.circuit_breakers.record_success(&host);

                        if let Some(tx) = task.result_tx.take() {
                            let _ = tx.send(ScatterResponse {
                                status,
                                headers,
                                body,
                            });
                        }
                        self.task_pool.mark_completed();
                        self.throughput.record();

                        debug!(
                            host = %host,
                            winner = %proxy_url,
                            latency_ms = latency_ms,
                            attempt = task.attempts + 1,
                            "task done"
                        );
                    }
                    BodyVerdict::ProxyBlocked => {
                        self.health.record_failure(&proxy_url, &host);
                        task.attempts += 1;
                        task.last_error = format!("proxy blocked (status={status})");
                        self.task_pool.push_back(task);
                    }
                    BodyVerdict::TargetError => {
                        self.circuit_breakers.record_target_error(&host);
                        task.attempts += 1;
                        task.last_error = format!("target error (status={status})");
                        self.task_pool.push_back(task);
                    }
                }
            }
            RaceOutcome::RequestError(err) => {
                self.health.record_failure(&proxy_url, &host);
                task.attempts += 1;
                task.last_error = err;
                self.task_pool.push_back(task);
            }
            RaceOutcome::Timeout => {
                self.health.record_failure(&proxy_url, &host);
                task.attempts += 1;
                task.last_error = "proxy timeout".into();
                self.task_pool.push_back(task);
            }
        }

        // Eviction check for the probe proxy.
        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!(
                    proxy = %proxy_url,
                    samples = samples,
                    "proxy dead | global success_rate=0%"
                );
            }
        }
    }
}