//! Daemon for processing batched requests with per-model concurrency control.
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::Duration;

use metrics::{counter, gauge, histogram};
use tokio::task::JoinSet;

use opentelemetry::trace::TraceContextExt;
use tracing::Instrument;
use tracing_opentelemetry::OpenTelemetrySpanExt;

use crate::FusilladeError;
use crate::batch::BatchId;
use crate::error::Result;
use crate::http::{HttpClient, HttpResponse};
use crate::manager::{DaemonStorage, Storage};
use crate::request::{DaemonId, RequestCompletionResult};

pub mod transitions;
pub mod types;

pub use types::{
    AnyDaemonRecord, DaemonData, DaemonRecord, DaemonState, DaemonStats, DaemonStatus, Dead,
    Initializing, Running,
};

/// Predicate function to determine if a response should be retried.
///
/// Takes an HTTP response and returns true if the request should be retried.
pub type ShouldRetryFn = Arc<dyn Fn(&HttpResponse) -> bool + Send + Sync>;

/// Default retry predicate: retry on server errors (5xx), rate limits (429), timeouts (408), and not found (404).
pub fn default_should_retry(response: &HttpResponse) -> bool {
    response.status >= 500
        || response.status == 429
        || response.status == 408
        || response.status == 404
}

/// Default function for creating the should_retry Arc
fn default_should_retry_fn() -> ShouldRetryFn {
    Arc::new(default_should_retry)
}

/// Default model escalations (empty map)
fn default_model_escalations() -> Arc<dashmap::DashMap<String, ModelEscalationConfig>> {
    Arc::new(dashmap::DashMap::new())
}

/// Default escalation threshold (15 minutes)
/// This should be greater than the processing timeout (10 minutes) to allow
/// a processing request to fall back to pending before escalation kicks in.
fn default_escalation_threshold_seconds() -> i64 {
    900
}

/// Model-based escalation configuration for routing requests to a different model
/// at claim time when approaching SLA deadline.
///
/// When a request is claimed with less than `escalation_threshold_seconds` remaining
/// before batch expiry, it will be routed to the `escalation_model` instead of the
/// original model. The batch API key automatically has access to escalation models
/// in the onwards routing cache (no separate API key needed).
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ModelEscalationConfig {
    /// The model to escalate to (e.g., "o1-preview" for requests using "gpt-4")
    pub escalation_model: String,

    /// Time threshold in seconds - escalate when time remaining before batch expiry
    /// is less than this value. Default: 900 (15 minutes)
    #[serde(default = "default_escalation_threshold_seconds")]
    pub escalation_threshold_seconds: i64,
}

/// Configuration for the daemon.
#[derive(Clone, serde::Serialize, serde::Deserialize)]
pub struct DaemonConfig {
    /// Maximum number of requests to claim in each iteration
    pub claim_batch_size: usize,

    /// Per-model concurrency limits (shared, can be updated dynamically).
    ///
    /// This map is the authoritative set of models the daemon will process.
    /// Models not present in this map will not be claimed. The caller (e.g.
    /// the control layer) is responsible for populating this with all known
    /// models and their concurrency limits.
    pub model_concurrency_limits: Arc<dashmap::DashMap<String, usize>>,

    /// Per-model escalation configurations for SLA-based model switching
    /// Maps model name -> escalation config (e.g., "gpt-4" -> "o1-preview")
    /// When a request is escalated, it's routed to the escalation_model by the control layer
    #[serde(skip, default = "default_model_escalations")]
    pub model_escalations: Arc<dashmap::DashMap<String, ModelEscalationConfig>>,

    /// How long to sleep between claim iterations
    pub claim_interval_ms: u64,

    /// Maximum number of retry attempts before giving up.
    pub max_retries: Option<u32>,

    /// Stop retrying (including escalations) this many milliseconds before the batch expires.
    ///
    /// - **Negative** (e.g., -300000 = -5 min): Retry for a buffer window AFTER SLA deadline (recommended)
    /// - **Zero**: Stop exactly at SLA deadline
    /// - **Positive**: Stop BEFORE SLA deadline (terminates retryable errors within the SLA deadline, avoid!)
    /// - **None**: No deadline awareness
    ///
    /// Default: 0 (stop retrying/escalating on SLA deadline)
    pub stop_before_deadline_ms: Option<i64>,

    /// Base backoff duration in milliseconds (will be exponentially increased)
    pub backoff_ms: u64,

    /// Factor by which the backoff_ms is increased with each retry
    pub backoff_factor: u64,

    /// Maximum backoff time in milliseconds
    pub max_backoff_ms: u64,

    /// Timeout for receiving response headers (connect + time-to-first-token) in milliseconds.
    /// This should be generous enough to cover slow model inference starts.
    /// Default: 86,400,000 (24 hours).
    pub first_chunk_timeout_ms: u64,

    /// Timeout for receiving the next chunk of response body in milliseconds.
    /// Once the server starts streaming, each inter-chunk gap must be shorter
    /// than this value or the request is considered stalled.
    /// Default: 86,400,000 (24 hours).
    pub chunk_timeout_ms: u64,

    /// Timeout for the entire response body in milliseconds.
    /// Catches slow-drip responses that never trip the per-chunk timeout
    /// but take an unreasonable total time.
    /// Default: 86,400,000 (24 hours).
    pub body_timeout_ms: u64,

    /// Interval for logging daemon status (requests in flight) in milliseconds
    /// Set to None to disable periodic status logging
    pub status_log_interval_ms: Option<u64>,

    /// Interval for sending heartbeats to update daemon status in database (milliseconds)
    pub heartbeat_interval_ms: u64,

    /// Predicate function to determine if a response should be retried.
    /// Defaults to retrying 5xx, 429, 408, and 404 status codes.
    #[serde(skip, default = "default_should_retry_fn")]
    pub should_retry: ShouldRetryFn,

    /// Maximum time a request can stay in "claimed" state before being unclaimed
    /// and returned to pending (milliseconds). This handles daemon crashes.
    pub claim_timeout_ms: u64,

    /// Maximum time a request can stay in "processing" state before being unclaimed
    /// and returned to pending (milliseconds). This handles daemon crashes during execution.
    pub processing_timeout_ms: u64,

    /// Time after a daemon's last heartbeat before its requests are considered
    /// orphaned and returned to pending (milliseconds). Should be significantly
    /// larger than `heartbeat_interval_ms` to avoid reclaiming from live daemons
    /// that are merely slow. Also reclaims from daemons explicitly marked dead.
    /// Default: 30,000 (30 seconds, 6× the default heartbeat interval).
    pub stale_daemon_threshold_ms: u64,

    /// Maximum number of stale requests to unclaim in a single poll cycle.
    /// Limits database load when many requests become stale simultaneously (e.g., daemon crash).
    pub unclaim_batch_size: usize,

    /// Interval for polling batches to perform finalization and check for cancellations (milliseconds).
    ///
    /// This polling loop serves two purposes:
    /// 1. **Batch Finalization**: Fetches active batches and triggers lazy finalization
    ///    (computing completion timestamps when all requests reach terminal states)
    /// 2. **Cancellation Detection**: Checks if any active batches have been cancelled
    ///    and aborts their in-flight requests
    ///
    /// Default: 5000ms (5 seconds)
    pub cancellation_poll_interval_ms: u64,

    /// Batch table column names to include as request headers.
    /// These values are sent as `x-fusillade-batch-{column}` headers with each request.
    /// Example: ["id", "created_by", "endpoint"] produces headers like:
    ///   - x-fusillade-batch-id
    ///   - x-fusillade-batch-created-by
    ///   - x-fusillade-batch-endpoint
    #[serde(default = "default_batch_metadata_fields")]
    pub batch_metadata_fields: Vec<String>,

    /// Interval for running the orphaned row purge task (milliseconds).
    /// Deletes orphaned request_templates and requests whose parent file/batch
    /// has been soft-deleted, for right-to-erasure compliance.
    /// Set to 0 to disable purging. Default: 3,600,000 (1 hour).
    pub purge_interval_ms: u64,

    /// Maximum number of orphaned rows to delete per purge iteration.
    /// Each iteration deletes up to this many requests and this many request_templates.
    /// Default: 1000.
    pub purge_batch_size: i64,

    /// Throttle delay between consecutive purge batches within a single drain
    /// cycle (milliseconds). Prevents sustained high DB load when many orphans
    /// exist. Default: 100.
    pub purge_throttle_ms: u64,

    /// Interval for emitting per-user throughput metrics via structured logs (milliseconds).
    /// Set to None to disable per-user throughput logging.
    /// Default: 60_000 (1 minute).
    pub throughput_log_interval_ms: Option<u64>,

    /// Request paths that should use SSE streaming.
    ///
    /// When a request's path matches one of these entries, an `X-Fusillade-Stream`
    /// header is sent with the request and the response is read as an SSE stream,
    /// then reassembled into the equivalent non-streaming JSON format.
    ///
    /// The upstream proxy (onwards) is responsible for injecting `"stream": true`
    /// and `"stream_options": {"include_usage": true}` into the request body
    /// when it sees this header.
    ///
    /// Example: `vec!["/v1/chat/completions", "/v1/completions"]`
    #[serde(default)]
    pub streamable_endpoints: Vec<String>,

    /// Weight controlling how much SLA urgency influences claim scheduling (0.0–1.0).
    ///
    /// Blends per-user fairness with batch deadline urgency when ordering claims:
    /// - `0.0`: Pure user-fairness (users with fewer in-flight requests go first)
    /// - `1.0`: Pure deadline urgency (batches closest to expiry go first)
    /// - `0.3`: Recommended starting point
    ///
    /// Default: `0.0` (backward-compatible, pure user-fairness).
    #[serde(default)]
    pub urgency_weight: f64,
}

fn default_batch_metadata_fields() -> Vec<String> {
    vec![
        "id".to_string(),
        "endpoint".to_string(),
        "created_at".to_string(),
        "completion_window".to_string(),
    ]
}

impl Default for DaemonConfig {
    fn default() -> Self {
        Self {
            claim_batch_size: 100,
            model_concurrency_limits: Arc::new(dashmap::DashMap::new()),
            model_escalations: default_model_escalations(),
            claim_interval_ms: 1000,
            max_retries: Some(1000),
            stop_before_deadline_ms: Some(0),
            backoff_ms: 1000,
            backoff_factor: 2,
            max_backoff_ms: 10000,
            first_chunk_timeout_ms: 540_000,    // 9 minutes
            chunk_timeout_ms: 540_000,          // 9 minutes
            body_timeout_ms: 60_000,            // 1 minute
            status_log_interval_ms: Some(2000), // Log every 2 seconds by default
            heartbeat_interval_ms: 5000,        // Heartbeat every 5 seconds by default
            should_retry: Arc::new(default_should_retry),
            claim_timeout_ms: 60000,             // 1 minute
            processing_timeout_ms: 600000,       // 10 minutes
            stale_daemon_threshold_ms: 30_000,   // 30 seconds (6× heartbeat interval)
            unclaim_batch_size: 100,             // Unclaim up to 100 stale requests per poll
            cancellation_poll_interval_ms: 5000, // Poll every 5 seconds by default
            batch_metadata_fields: default_batch_metadata_fields(),
            purge_interval_ms: 600_000, // 10 minutes
            purge_batch_size: 1000,
            purge_throttle_ms: 100,
            throughput_log_interval_ms: Some(60_000), // Log per-user throughput every minute
            streamable_endpoints: Vec::new(),
            urgency_weight: 0.0,
        }
    }
}

/// Per-user throughput counters, reset after each emission cycle.
struct UserThroughputStats {
    completed: AtomicU64,
    failed: AtomicU64,
}

/// Daemon that processes batched requests.
///
/// The daemon continuously claims pending requests from storage, enforces
/// per-model concurrency limits, and dispatches requests for execution.
pub struct Daemon<S, H>
where
    S: Storage + DaemonStorage,
    H: HttpClient,
{
    daemon_id: DaemonId,
    storage: Arc<S>,
    http_client: Arc<H>,
    config: DaemonConfig,
    requests_in_flight: Arc<dashmap::DashMap<String, AtomicUsize>>,
    /// Per-user in-flight request counts across all models, used to prioritise
    /// users with fewer active requests during claim (per-user fair scheduling).
    user_requests_in_flight: Arc<dashmap::DashMap<String, AtomicUsize>>,
    /// Per-user throughput counters for periodic OTel emission.
    user_throughput: Arc<dashmap::DashMap<String, UserThroughputStats>>,
    requests_processed: Arc<AtomicU64>,
    requests_failed: Arc<AtomicU64>,
    shutdown_token: tokio_util::sync::CancellationToken,
    /// Map of batch_id -> cancellation token for batch-level cancellation
    /// All requests in a batch share the same cancellation token
    cancellation_tokens: Arc<dashmap::DashMap<BatchId, tokio_util::sync::CancellationToken>>,
}

impl<S, H> Daemon<S, H>
where
    S: Storage + DaemonStorage + 'static,
    H: HttpClient + 'static,
{
    /// Create a new daemon.
    pub fn new(
        storage: Arc<S>,
        http_client: Arc<H>,
        config: DaemonConfig,
        shutdown_token: tokio_util::sync::CancellationToken,
    ) -> Self {
        Self {
            daemon_id: DaemonId::from(uuid::Uuid::new_v4()),
            storage,
            http_client,
            config,
            requests_in_flight: Arc::new(dashmap::DashMap::new()),
            user_requests_in_flight: Arc::new(dashmap::DashMap::new()),
            user_throughput: Arc::new(dashmap::DashMap::new()),
            requests_processed: Arc::new(AtomicU64::new(0)),
            requests_failed: Arc::new(AtomicU64::new(0)),
            shutdown_token,
            cancellation_tokens: Arc::new(dashmap::DashMap::new()),
        }
    }

    /// Run the daemon loop.
    ///
    /// This continuously claims and processes requests until an error occurs
    /// or the task is cancelled.
    ///
    /// The daemon periodically polls for cancelled batches and aborts in-flight requests.
    #[tracing::instrument(name = "fusillade.daemon.run", skip(self), fields(daemon_id = %self.daemon_id))]
    pub async fn run(self: Arc<Self>) -> Result<()> {
        tracing::info!("Daemon starting main processing loop");

        // Register daemon in database
        let daemon_record = DaemonRecord {
            data: DaemonData {
                id: self.daemon_id,
                hostname: types::get_hostname(),
                pid: types::get_pid(),
                version: types::get_version(),
                config_snapshot: serde_json::to_value(&self.config)
                    .expect("Failed to serialize daemon config"),
            },
            state: Initializing {
                started_at: chrono::Utc::now(),
            },
        };

        let running_record = daemon_record.start(self.storage.as_ref()).await?;
        tracing::info!("Daemon registered in database");

        // Spawn periodic heartbeat task
        let storage = self.storage.clone();
        let requests_in_flight = self.requests_in_flight.clone();
        let requests_processed = self.requests_processed.clone();
        let requests_failed = self.requests_failed.clone();
        let daemon_id = self.daemon_id;
        let heartbeat_interval_ms = self.config.heartbeat_interval_ms;
        let shutdown_signal = self.shutdown_token.clone();

        let heartbeat_handle = tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_millis(heartbeat_interval_ms));
            let mut daemon_record = running_record;

            loop {
                tokio::select! {
                    _ = interval.tick() => {
                        let stats = DaemonStats {
                            requests_processed: requests_processed.load(Ordering::Relaxed),
                            requests_failed: requests_failed.load(Ordering::Relaxed),
                            requests_in_flight: requests_in_flight.iter().map(|e| e.value().load(Ordering::Relaxed)).sum(),
                        };

                        // Clone the record so we preserve it if heartbeat fails
                        let current = daemon_record.clone();
                        let heartbeat_start = std::time::Instant::now();
                        match current.heartbeat(stats, storage.as_ref()).await {
                            Ok(updated) => {
                                histogram!("fusillade_heartbeat_duration_seconds")
                                    .record(heartbeat_start.elapsed().as_secs_f64());
                                daemon_record = updated;
                                tracing::trace!(
                                    daemon_id = %daemon_id,
                                    "Heartbeat sent"
                                );
                            }
                            Err(e) => {
                                histogram!("fusillade_heartbeat_duration_seconds")
                                    .record(heartbeat_start.elapsed().as_secs_f64());
                                counter!("fusillade_heartbeat_failures_total").increment(1);
                                tracing::error!(
                                    daemon_id = %daemon_id,
                                    error = %e,
                                    "Failed to send heartbeat"
                                );
                                // daemon_record stays unchanged on error
                            }
                        }
                    }
                    _ = shutdown_signal.cancelled() => {
                        // Mark daemon as dead on shutdown
                        tracing::info!("Shutting down heartbeat task");
                        if let Err(e) = daemon_record.shutdown(storage.as_ref()).await {
                            tracing::error!(
                                daemon_id = %daemon_id,
                                error = %e,
                                "Failed to mark daemon as dead during shutdown"
                            );
                        }
                        break;
                    }
                }
            }
        });

        // Spawn periodic status logging task if configured
        if let Some(interval_ms) = self.config.status_log_interval_ms {
            let requests_in_flight = self.requests_in_flight.clone();
            let daemon_id = self.daemon_id;
            tokio::spawn(async move {
                let mut interval = tokio::time::interval(Duration::from_millis(interval_ms));
                loop {
                    interval.tick().await;
                    let count: usize = requests_in_flight
                        .iter()
                        .map(|e| e.value().load(Ordering::Relaxed))
                        .sum();
                    tracing::debug!(
                        daemon_id = %daemon_id,
                        requests_in_flight = count,
                        "Daemon status"
                    );
                }
            });
        }

        // Spawn periodic per-user throughput emission task if configured
        if let Some(interval_ms) = self.config.throughput_log_interval_ms {
            let user_throughput = self.user_throughput.clone();
            let user_requests_in_flight = self.user_requests_in_flight.clone();
            let daemon_id = self.daemon_id;
            let shutdown_token = self.shutdown_token.clone();
            tokio::spawn(async move {
                let mut interval = tokio::time::interval(Duration::from_millis(interval_ms));
                // Skip the immediate first tick to avoid a near-zero window on the first emission
                interval.tick().await;
                let mut last_emission = std::time::Instant::now();

                loop {
                    tokio::select! {
                        _ = interval.tick() => {
                            let elapsed = last_emission.elapsed();
                            let window_secs = elapsed.as_secs_f64();

                            // Atomically read-and-reset each user's counters
                            let mut users_to_remove = Vec::new();
                            for entry in user_throughput.iter() {
                                let user_id = entry.key();
                                let completed = entry.value().completed.swap(0, Ordering::Relaxed);
                                let failed = entry.value().failed.swap(0, Ordering::Relaxed);

                                if completed > 0 || failed > 0 {
                                    let in_flight = user_requests_in_flight
                                        .get(user_id)
                                        .map(|e| e.value().load(Ordering::Relaxed))
                                        .unwrap_or(0);
                                    let throughput_rpm = if window_secs > 0.0 {
                                        (completed + failed) as f64 / window_secs * 60.0
                                    } else {
                                        0.0
                                    };

                                    tracing::info!(
                                        daemon_id = %daemon_id,
                                        user = %user_id,
                                        completed = completed,
                                        failed = failed,
                                        in_flight = in_flight,
                                        throughput_rpm = format!("{throughput_rpm:.1}"),
                                        window_seconds = format!("{window_secs:.1}"),
                                        "fusillade.user_throughput"
                                    );
                                } else {
                                    // No activity — mark for eviction
                                    users_to_remove.push(user_id.clone());
                                }
                            }

                            // Evict inactive users to prevent unbounded map growth
                            for user_id in users_to_remove {
                                user_throughput.remove(&user_id);
                            }

                            last_emission = std::time::Instant::now();
                        }
                        _ = shutdown_token.cancelled() => {
                            tracing::debug!(
                                daemon_id = %daemon_id,
                                "Shutting down per-user throughput emission task"
                            );
                            break;
                        }
                    }
                }
            });
        }

        // Spawn periodic batch polling task for finalization and cancellation detection
        // This serves two purposes in one efficient loop:
        // 1. Triggers lazy finalization by fetching batches (computes completion timestamps)
        // 2. Detects cancelled batches and aborts their in-flight requests
        let cancellation_tokens = self.cancellation_tokens.clone();
        let storage = self.storage.clone();
        let shutdown_token = self.shutdown_token.clone();
        let cancellation_poll_interval_ms = self.config.cancellation_poll_interval_ms;

        tokio::spawn(async move {
            let mut interval =
                tokio::time::interval(Duration::from_millis(cancellation_poll_interval_ms));
            tracing::info!(
                interval_ms = cancellation_poll_interval_ms,
                "Batch polling started (finalization + cancellation detection)"
            );

            loop {
                tokio::select! {
                    _ = interval.tick() => {
                        // Get all active batch IDs we're currently processing
                        let active_batch_ids: Vec<BatchId> = cancellation_tokens
                            .iter()
                            .map(|entry| *entry.key())
                            .collect();

                        if active_batch_ids.is_empty() {
                            continue;
                        }

                        let poll_start = std::time::Instant::now();
                        gauge!("fusillade_cancellation_poll_batches_checked")
                            .set(active_batch_ids.len() as f64);

                        // Single bulk query to find which active batches have been cancelled
                        match storage.get_cancelled_batch_ids(&active_batch_ids).await {
                            Ok(cancelled_ids) => {
                                for batch_id in cancelled_ids {
                                    if let Some(entry) = cancellation_tokens.get(&batch_id) {
                                        entry.value().cancel();
                                        counter!("fusillade_batches_cancelled_total").increment(1);
                                        tracing::info!(batch_id = %batch_id, "Cancelled all requests in batch");
                                        drop(entry);
                                        cancellation_tokens.remove(&batch_id);
                                    }
                                }
                            }
                            Err(e) => {
                                counter!("fusillade_cancellation_poll_errors_total").increment(1);
                                tracing::warn!(
                                    error = %e,
                                    "Failed to check batch cancellation status"
                                );
                            }
                        }

                        histogram!("fusillade_cancellation_poll_duration_seconds")
                            .record(poll_start.elapsed().as_secs_f64());
                    }
                    _ = shutdown_token.cancelled() => {
                        tracing::info!("Shutting down batch polling");
                        break;
                    }
                }
            }
        });

        // Spawn periodic purge task for orphaned rows (right-to-erasure compliance)
        if self.config.purge_interval_ms > 0 {
            let storage = self.storage.clone();
            let shutdown_token = self.shutdown_token.clone();
            let purge_interval_ms = self.config.purge_interval_ms;
            let purge_batch_size = self.config.purge_batch_size;
            let purge_throttle_ms = self.config.purge_throttle_ms;

            tokio::spawn(async move {
                tracing::info!(
                    interval_ms = purge_interval_ms,
                    batch_size = purge_batch_size,
                    throttle_ms = purge_throttle_ms,
                    "Orphaned row purge task started"
                );

                loop {
                    // Sleep for the configured interval (interruptible by shutdown)
                    tokio::select! {
                        _ = tokio::time::sleep(Duration::from_millis(purge_interval_ms)) => {},
                        _ = shutdown_token.cancelled() => {
                            tracing::info!("Shutting down purge task");
                            break;
                        }
                    }

                    // Drain orphaned rows in batches
                    loop {
                        match storage.purge_orphaned_rows(purge_batch_size).await {
                            Ok(0) => break,
                            Ok(deleted) => {
                                counter!("fusillade_rows_purged_total").increment(deleted);
                                tracing::debug!(deleted, "Purged orphaned rows");
                                // Throttle between batches to avoid sustained DB load
                                tokio::select! {
                                    _ = tokio::time::sleep(Duration::from_millis(purge_throttle_ms)) => {},
                                    _ = shutdown_token.cancelled() => {
                                        tracing::info!("Shutting down purge task during drain");
                                        return;
                                    }
                                }
                            }
                            Err(e) => {
                                counter!("fusillade_purge_errors_total").increment(1);
                                tracing::error!(error = %e, "Failed to purge orphaned rows");
                                break;
                            }
                        }
                    }
                }
            });
        }

        let mut join_set: JoinSet<Result<()>> = JoinSet::new();

        let run_result = loop {
            // Check for shutdown signal
            if self.shutdown_token.is_cancelled() {
                tracing::info!("Shutdown signal received, stopping daemon");
                break Ok(());
            }

            // Poll for completed tasks (non-blocking)
            while let Some(result) = join_set.try_join_next() {
                match result {
                    Ok(Ok(())) => {
                        tracing::trace!("Task completed successfully");
                    }
                    Ok(Err(e)) => {
                        tracing::error!(error = %e, "Task failed");
                    }
                    Err(join_error) => {
                        counter!("fusillade_task_panics_total").increment(1);
                        tracing::error!(error = %join_error, "Task panicked");
                    }
                }
            }

            tracing::trace!("Sleeping before claiming");
            tokio::select! {
                _ = tokio::time::sleep(Duration::from_millis(self.config.claim_interval_ms)) => {},
                _ = self.shutdown_token.cancelled() => {
                    tracing::info!("Shutdown signal received, stopping daemon");
                    break Ok(());
                }
            }
            // Compute available capacity per model: configured limit minus in-flight.
            let available_capacity: HashMap<String, usize> = self
                .config
                .model_concurrency_limits
                .iter()
                .filter_map(|entry| {
                    let model = entry.key().clone();
                    let limit = *entry.value();
                    let in_flight = self
                        .requests_in_flight
                        .get(&model)
                        .map(|e| e.value().load(Ordering::Relaxed))
                        .unwrap_or(0);
                    let available = limit.saturating_sub(in_flight);
                    if available > 0 {
                        Some((model, available))
                    } else {
                        None
                    }
                })
                .collect();

            if available_capacity.is_empty() {
                tracing::trace!("No capacity available for any model, skipping claim");
                continue;
            }

            // Record available capacity before claiming
            let total_capacity: usize = available_capacity.values().sum();
            gauge!("fusillade_claim_capacity").set(total_capacity as f64);

            // Snapshot per-user in-flight counts for fair scheduling.
            // Users with fewer active requests get prioritised in the claim query.
            // Only include users with count > 0 to avoid unbounded array growth
            // (missing users COALESCE to 0 in the SQL query).
            let user_active_counts: HashMap<String, usize> = self
                .user_requests_in_flight
                .iter()
                .filter_map(|entry| {
                    let count = entry.value().load(Ordering::Relaxed);
                    if count > 0 {
                        Some((entry.key().clone(), count))
                    } else {
                        None
                    }
                })
                .collect();

            // Claim a batch of pending requests (guaranteed ≤ available capacity per model)
            let claim_start = std::time::Instant::now();
            let mut claimed = self
                .storage
                .claim_requests(
                    self.config.claim_batch_size,
                    self.daemon_id,
                    &available_capacity,
                    &user_active_counts,
                )
                .await?;
            histogram!("fusillade_claim_duration_seconds")
                .record(claim_start.elapsed().as_secs_f64());

            // Record claim metrics
            histogram!("fusillade_claim_size").record(claimed.len() as f64);

            tracing::debug!(
                claimed_count = claimed.len(),
                "Claimed requests from storage"
            );

            // Route requests to escalated models if time is running low
            // This replaces the old SLA racing system with a simpler approach:
            // at claim time, we check if there's enough time remaining and route
            // to the escalated model if below threshold
            for request in &mut claimed {
                if let Some(config) = self.config.model_escalations.get(&request.data.model) {
                    let time_remaining = request.state.batch_expires_at - chrono::Utc::now();
                    if time_remaining.num_seconds() < config.escalation_threshold_seconds {
                        let original_model = request.data.model.clone();
                        request.data.model = config.escalation_model.clone();

                        // Update the model field in the request body JSON
                        if let Ok(mut json) =
                            serde_json::from_str::<serde_json::Value>(&request.data.body)
                            && let Some(obj) = json.as_object_mut()
                        {
                            obj.insert(
                                "model".to_string(),
                                serde_json::Value::String(config.escalation_model.clone()),
                            );
                            if let Ok(new_body) = serde_json::to_string(&json) {
                                request.data.body = new_body;
                            }
                        }

                        // No API key swap needed - batch API keys automatically have access
                        // to escalation models in the onwards routing cache
                        counter!("fusillade_requests_routed_to_escalation_total", "original_model" => original_model.clone(), "escalation_model" => config.escalation_model.clone()).increment(1);
                        tracing::info!(
                            request_id = %request.data.id,
                            original_model = %original_model,
                            escalation_model = %config.escalation_model,
                            time_remaining_seconds = time_remaining.num_seconds(),
                            threshold_seconds = config.escalation_threshold_seconds,
                            "Routing request to escalation model due to time pressure"
                        );
                    }
                }
            }

            // Group requests by model for better concurrency control visibility
            let mut by_model: HashMap<String, Vec<_>> = HashMap::new();
            for request in claimed {
                let model = request.data.model.clone();
                by_model.entry(model).or_default().push(request);
            }

            tracing::debug!(
                models = by_model.len(),
                total_requests = by_model.values().map(|v| v.len()).sum::<usize>(),
                "Grouped requests by model"
            );

            // Dispatch requests
            for (model, requests) in by_model {
                tracing::debug!(model = %model, count = requests.len(), "Processing requests for model");

                for request in requests {
                    let request_id = request.data.id;
                    let batch_id = request.data.batch_id;

                    tracing::trace!(
                        request_id = %request_id,
                        batch_id = %batch_id,
                        model = %model,
                        "Spawning processing task"
                    );

                    // Spawn a processing task
                    let model_clone = model.clone();
                    let user_id = request.data.created_by.clone();
                    let completion_window = request
                        .data
                        .batch_metadata
                        .get("completion_window")
                        .cloned()
                        .unwrap_or_default();
                    let storage = self.storage.clone();
                    let http_client = (*self.http_client).clone();
                    let retry_config = (&self.config).into();
                    let requests_in_flight = self.requests_in_flight.clone();
                    let user_throughput = self.user_throughput.clone();
                    let user_requests_in_flight = self.user_requests_in_flight.clone();
                    let requests_processed = self.requests_processed.clone();
                    let requests_failed = self.requests_failed.clone();
                    let should_retry = self.config.should_retry.clone();
                    let shutdown_token = self.shutdown_token.clone();
                    let cancellation_tokens = self.cancellation_tokens.clone();

                    // Get or create a cancellation token for this batch
                    // All requests in a batch share the same token
                    let batch_cancellation_token =
                        cancellation_tokens.entry(batch_id).or_default().clone();

                    // Increment per-model in-flight counter and gauge
                    requests_in_flight
                        .entry(model_clone.clone())
                        .or_default()
                        .fetch_add(1, Ordering::Relaxed);
                    gauge!("fusillade_requests_in_flight", "model" => model_clone.clone())
                        .increment(1.0);

                    // Increment per-user in-flight counter for fair scheduling
                    user_requests_in_flight
                        .entry(user_id.clone())
                        .or_default()
                        .fetch_add(1, Ordering::Relaxed);
                    gauge!("fusillade_user_requests_in_flight", "user" => user_id.clone(), "completion_window" => completion_window.clone())
                        .increment(1.0);

                    let process_span = tracing::info_span!(
                        parent: tracing::Span::none(),
                        "fusillade.process_request",
                        trace_id = tracing::field::Empty,
                        otel.name = "fusillade.process_request",
                        request_id = %request_id,
                        batch_id = %batch_id,
                        model = %model,
                        outcome = tracing::field::Empty,
                    );

                    join_set.spawn(async move {
                        // Record trace_id from OTel context
                        let span = tracing::Span::current();
                        let sc = span.context().span().span_context().clone();
                        if sc.is_valid() {
                            span.record("trace_id", tracing::field::display(sc.trace_id()));
                        }

                        // Track processing start time for duration metrics
                        let processing_start = std::time::Instant::now();

                        // Ensure we decrement the per-model and per-user counters when this task completes
                        let model_for_guard = model_clone.clone();
                        let user_for_guard = user_id.clone();
                        let cw_for_guard = completion_window.clone();
                        let in_flight_for_guard = requests_in_flight.clone();
                        let user_in_flight_for_guard = user_requests_in_flight.clone();
                        let _guard = scopeguard::guard((), move |_| {
                            if let Some(counter) = in_flight_for_guard.get(&model_for_guard) {
                                counter.value().fetch_sub(1, Ordering::Relaxed);
                            }
                            gauge!("fusillade_requests_in_flight", "model" => model_for_guard).decrement(1.0);
                            gauge!("fusillade_user_requests_in_flight", "user" => user_for_guard.clone(), "completion_window" => cw_for_guard).decrement(1.0);
                            if let Some(counter) = user_in_flight_for_guard.get(&user_for_guard) {
                                let prev = counter.value().fetch_sub(1, Ordering::Relaxed);
                                drop(counter);
                                if prev == 1 {
                                    // Counter reached zero — remove to prevent unbounded map growth
                                    user_in_flight_for_guard.remove(&user_for_guard);
                                }
                            }
                        });

                        // Duration span covering the claimed state (from claim to process start)
                        let retry_attempt = request.state.retry_attempt;
                        let processing = async {
                            tracing::debug!("Sending batch request to inference endpoint");
                            request.process(
                                http_client,
                                storage.as_ref()
                            ).await
                        }.instrument(tracing::info_span!(
                            "fusillade.state.claimed",
                            otel.name = "fusillade.state.claimed",
                            request_id = %request_id,
                            daemon_id = %daemon_id,
                            retry_attempt,
                        )).await?;

                        // Capture retry attempt count before completion (not preserved in Completed state)
                        let retry_attempt_at_completion = processing.state.retry_attempt;
                        // Capture batch expiry time for SLA missed completion tracking
                        let batch_expires_at = processing.state.batch_expires_at;

                        // Duration span covering the processing state (HTTP in-flight)
                        let completion_result = async {
                            let cancellation = async {
                                tokio::select! {
                                    _ = batch_cancellation_token.cancelled() => {
                                        crate::request::transitions::CancellationReason::User
                                    }
                                    _ = shutdown_token.cancelled() => {
                                        crate::request::transitions::CancellationReason::Shutdown
                                    }
                                }
                            };

                            processing.complete(storage.as_ref(), |response| {
                                (should_retry)(response)
                            }, cancellation).await
                        }.instrument(tracing::info_span!(
                            "fusillade.state.processing",
                            otel.name = "fusillade.state.processing",
                            request_id = %request_id,
                            retry_attempt = retry_attempt_at_completion
                        )).await;

                        match completion_result {
                            Ok(RequestCompletionResult::Completed(completed)) => {
                                tracing::Span::current().record("outcome", "completed");
                                requests_processed.fetch_add(1, Ordering::Relaxed);
                                user_throughput.entry(user_id.clone()).or_insert_with(|| UserThroughputStats {
                                    completed: AtomicU64::new(0),
                                    failed: AtomicU64::new(0),
                                }).completed.fetch_add(1, Ordering::Relaxed);
                                counter!("fusillade_requests_completed_total", "model" => model_clone.clone(), "status" => "success").increment(1);
                                counter!("fusillade_user_requests_completed_total", "user" => user_id.clone(), "status" => "success", "completion_window" => completion_window.clone()).increment(1);
                                histogram!("fusillade_request_duration_seconds", "model" => model_clone.clone(), "status" => "success")
                                    .record(processing_start.elapsed().as_secs_f64());
                                // Record how many retries it took to succeed (0 = first attempt succeeded)
                                histogram!("fusillade_retry_attempts_on_success", "model" => model_clone.clone())
                                    .record(retry_attempt_at_completion as f64);

                                // Track requests completing after SLA
                                if completed.state.completed_at > batch_expires_at {
                                    counter!("fusillade_requests_completed_after_sla_total", "model" => model_clone.clone(), "status" => "success").increment(1);
                                    tracing::warn!(
                                        request_id = %request_id,
                                        batch_id = %batch_id,
                                        "Request completed successfully after SLA"
                                    );
                                }

                                tracing::debug!(request_id = %request_id, retry_attempts = retry_attempt_at_completion, "Request completed successfully");
                            }
                            Ok(RequestCompletionResult::Failed(failed)) => {
                                tracing::Span::current().record("outcome", "failed");
                                let retry_attempt = failed.state.retry_attempt;

                                // Check if this is a retriable error using the FailureReason
                                if failed.state.reason.is_retriable() {
                                    tracing::warn!(
                                        request_id = %request_id,
                                        retry_attempt,
                                        error = %failed.state.reason.to_error_message(),
                                        "Request failed with retriable error, attempting retry"
                                    );

                                    // Attempt to retry
                                    match failed.can_retry(retry_attempt, retry_config) {
                                        Ok(pending) => {
                                            // Can retry - persist as Pending
                                            if let Err(e) = storage.persist(&pending).await {
                                                tracing::error!(
                                                    request_id = %request_id,
                                                    batch_id = %batch_id,
                                                    retry_attempt = pending.state.retry_attempt,
                                                    error = %e,
                                                    "Failed to persist retry — request orphaned in processing state"
                                                );
                                                return Err(e);
                                            }
                                            counter!(
                                                "fusillade_requests_retried_total",
                                                "model" => model_clone.clone(),
                                                "attempt" => (retry_attempt + 1).to_string()
                                            ).increment(1);
                                            tracing::info!(
                                                request_id = %request_id,
                                                batch_id = %batch_id,
                                                retry_attempt = retry_attempt + 1,
                                                "request.retry_persisted"
                                            );
                                        }
                                        Err(failed) => {
                                            // No retries left - persist as Failed (terminal)
                                            if let Err(e) = storage.persist(&*failed).await {
                                                tracing::error!(
                                                    request_id = %request_id,
                                                    batch_id = %batch_id,
                                                    retry_attempt,
                                                    error = %e,
                                                    "Failed to persist terminal failure — request orphaned in processing state"
                                                );
                                                return Err(e);
                                            }
                                            user_throughput.entry(user_id.clone()).or_insert_with(|| UserThroughputStats {
                                                completed: AtomicU64::new(0),
                                                failed: AtomicU64::new(0),
                                            }).failed.fetch_add(1, Ordering::Relaxed);
                                            requests_failed.fetch_add(1, Ordering::Relaxed);
                                            counter!("fusillade_requests_completed_total", "model" => model_clone.clone(), "status" => "failed", "reason" => failed.state.reason.metric_label(), "status_code" => failed.state.reason.status_code_label()).increment(1);
                                            counter!("fusillade_user_requests_completed_total", "user" => user_id.clone(), "status" => "failed", "completion_window" => completion_window.clone()).increment(1);
                                            histogram!("fusillade_request_duration_seconds", "model" => model_clone.clone(), "status" => "failed")
                                                .record(processing_start.elapsed().as_secs_f64());

                                            // Track requests completing after SLA
                                            if failed.state.failed_at > batch_expires_at {
                                                counter!("fusillade_requests_completed_after_sla_total", "model" => model_clone.clone(), "status" => "failed").increment(1);
                                                tracing::warn!(
                                                    request_id = %request_id,
                                                    batch_id = %batch_id,
                                                    "Request failed permanently after SLA"
                                                );
                                            }

                                            tracing::warn!(
                                                request_id = %request_id,
                                                batch_id = %batch_id,
                                                retry_attempt,
                                                failure_reason = %failed.state.reason.metric_label(),
                                                error = %failed.state.reason.to_error_message(),
                                                "request.terminal_failure"
                                            );
                                        }
                                    }
                                } else {
                                    requests_failed.fetch_add(1, Ordering::Relaxed);
                                    user_throughput.entry(user_id.clone()).or_insert_with(|| UserThroughputStats {
                                        completed: AtomicU64::new(0),
                                        failed: AtomicU64::new(0),
                                    }).failed.fetch_add(1, Ordering::Relaxed);
                                    counter!("fusillade_requests_completed_total", "model" => model_clone.clone(), "status" => "failed", "reason" => failed.state.reason.metric_label(), "status_code" => failed.state.reason.status_code_label()).increment(1);
                                    counter!("fusillade_user_requests_completed_total", "user" => user_id.clone(), "status" => "failed", "completion_window" => completion_window.clone()).increment(1);
                                    histogram!("fusillade_request_duration_seconds", "model" => model_clone.clone(), "status" => "failed")
                                        .record(processing_start.elapsed().as_secs_f64());

                                    // Track requests completing after SLA
                                    if failed.state.failed_at > batch_expires_at {
                                        counter!("fusillade_requests_completed_after_sla_total", "model" => model_clone.clone(), "status" => "failed").increment(1);
                                        tracing::warn!(
                                            request_id = %request_id,
                                            batch_id = %batch_id,
                                            "Request failed with non-retriable error after SLA"
                                        );
                                    }

                                    tracing::warn!(
                                        request_id = %request_id,
                                        batch_id = %batch_id,
                                        failure_reason = %failed.state.reason.metric_label(),
                                        error = %failed.state.reason.to_error_message(),
                                        "request.terminal_failure"
                                    );
                                }
                            }
                            Ok(RequestCompletionResult::Canceled(_canceled)) => {
                                tracing::Span::current().record("outcome", "canceled");
                                counter!("fusillade_requests_completed_total", "model" => model_clone.clone(), "status" => "cancelled").increment(1);
                                counter!("fusillade_user_requests_completed_total", "user" => user_id.clone(), "status" => "cancelled", "completion_window" => completion_window.clone()).increment(1);
                                tracing::debug!(request_id = %request_id, "Request canceled by user");
                            }
                            Err(FusilladeError::Shutdown) => {
                                tracing::Span::current().record("outcome", "shutdown");
                                tracing::info!(request_id = %request_id, "Request aborted due to shutdown");
                                // Don't count as failed - request will be reclaimed
                            }
                            Err(e) => {
                                tracing::Span::current().record("outcome", "error");
                                // Unexpected error
                                tracing::error!(request_id = %request_id, error = %e, "Unexpected error processing request");
                                return Err(e);
                            }
                        }

                        // Note: We don't remove the batch cancellation token here since
                        // multiple requests in the same batch share it. Tokens are cleaned
                        // up when the daemon shuts down or batch completes.

                        Ok(())
                    }.instrument(process_span));
                }
            }
        };

        // Wait for heartbeat task to complete (it will mark daemon as dead)
        tracing::info!("Waiting for heartbeat task to complete");
        if let Err(e) = heartbeat_handle.await {
            tracing::error!(error = %e, "Heartbeat task panicked");
        }

        run_result
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::TestDbPools;
    use crate::http::{HttpResponse, MockHttpClient};
    use crate::manager::{DaemonExecutor, postgres::PostgresRequestManager};
    use std::time::Duration;

    #[sqlx::test]
    #[test_log::test]
    async fn test_daemon_claims_and_completes_request(pool: sqlx::PgPool) {
        // Setup: Create HTTP client with mock response
        let http_client = Arc::new(MockHttpClient::new());
        http_client.add_response(
            "POST /v1/test",
            Ok(HttpResponse {
                status: 200,
                body: r#"{"result":"success"}"#.to_string(),
            }),
        );

        // Setup: Create manager with fast claim interval (no sleeping)
        let config = DaemonConfig {
            claim_batch_size: 10,
            claim_interval_ms: 10, // Very fast for testing
            model_concurrency_limits: {
                let m = Arc::new(dashmap::DashMap::new());
                m.insert("test-model".to_string(), 10);
                m
            },

            model_escalations: Arc::new(dashmap::DashMap::new()),
            max_retries: Some(3),
            stop_before_deadline_ms: None,
            backoff_ms: 100,
            backoff_factor: 2,
            max_backoff_ms: 1000,
            first_chunk_timeout_ms: 5000,
            chunk_timeout_ms: 5000,
            body_timeout_ms: 86_400_000,
            status_log_interval_ms: None, // Disable status logging in tests
            heartbeat_interval_ms: 5000,
            should_retry: Arc::new(default_should_retry),
            claim_timeout_ms: 60000,
            processing_timeout_ms: 600000,
            stale_daemon_threshold_ms: 30_000,
            unclaim_batch_size: 100,
            batch_metadata_fields: vec![],
            cancellation_poll_interval_ms: 100, // Fast polling for tests
            purge_interval_ms: 0,               // Disabled in tests
            purge_batch_size: 1000,
            purge_throttle_ms: 100,
            streamable_endpoints: Vec::new(),
            throughput_log_interval_ms: None,
            urgency_weight: 0.0,
        };

        let manager = Arc::new(
            PostgresRequestManager::with_client(
                TestDbPools::new(pool.clone()).await.unwrap(),
                http_client.clone(),
            )
            .with_config(config),
        );

        // Setup: Create a file and batch to associate with our request
        let file_id = manager
            .create_file(
                "test-file".to_string(),
                Some("Test file".to_string()),
                vec![crate::RequestTemplateInput {
                    custom_id: None,
                    endpoint: "https://api.example.com".to_string(),
                    method: "POST".to_string(),
                    path: "/v1/test".to_string(),
                    body: r#"{"prompt":"test"}"#.to_string(),
                    model: "test-model".to_string(),
                    api_key: "test-key".to_string(),
                }],
            )
            .await
            .expect("Failed to create file");

        let batch = manager
            .create_batch(crate::batch::BatchInput {
                file_id,
                endpoint: "/v1/chat/completions".to_string(),
                completion_window: "24h".to_string(),
                metadata: None,
                created_by: None,
                api_key_id: None,
                api_key: None,
                total_requests: None,
            })
            .await
            .expect("Failed to create batch");

        // Get the created request from the batch
        let requests = manager
            .get_batch_requests(batch.id)
            .await
            .expect("Failed to get batch requests");
        assert_eq!(requests.len(), 1);
        let request_id = requests[0].id();

        // Start the daemon
        let shutdown_token = tokio_util::sync::CancellationToken::new();
        manager
            .clone()
            .run(shutdown_token.clone())
            .expect("Failed to start daemon");

        // Poll for completion (with timeout)
        let start = tokio::time::Instant::now();
        let timeout = Duration::from_secs(5);
        let mut completed = false;

        while start.elapsed() < timeout {
            let results = manager
                .get_requests(vec![request_id])
                .await
                .expect("Failed to get request");

            if let Some(Ok(any_request)) = results.first()
                && any_request.is_terminal()
            {
                if let crate::AnyRequest::Completed(req) = any_request {
                    // Verify the request was completed successfully
                    assert_eq!(req.state.response_status, 200);
                    assert_eq!(req.state.response_body, r#"{"result":"success"}"#);
                    completed = true;
                    break;
                } else {
                    panic!(
                        "Request reached terminal state but was not completed: {:?}",
                        any_request
                    );
                }
            }

            tokio::time::sleep(Duration::from_millis(50)).await;
        }

        // Stop the daemon
        shutdown_token.cancel();

        // Assert that the request completed
        assert!(
            completed,
            "Request did not complete within timeout. Check daemon processing."
        );

        // Verify HTTP client was called exactly once
        assert_eq!(http_client.call_count(), 1);
        let calls = http_client.get_calls();
        assert_eq!(calls[0].method, "POST");
        assert_eq!(calls[0].path, "/v1/test");
        assert_eq!(calls[0].api_key, "test-key");
    }

    #[sqlx::test]
    async fn test_daemon_respects_per_model_concurrency_limits(pool: sqlx::PgPool) {
        // Setup: Create HTTP client with triggered responses
        let http_client = Arc::new(MockHttpClient::new());

        // Add 5 triggered responses for our 5 requests
        let trigger1 = http_client.add_response_with_trigger(
            "POST /v1/test",
            Ok(HttpResponse {
                status: 200,
                body: r#"{"result":"1"}"#.to_string(),
            }),
        );
        let trigger2 = http_client.add_response_with_trigger(
            "POST /v1/test",
            Ok(HttpResponse {
                status: 200,
                body: r#"{"result":"2"}"#.to_string(),
            }),
        );
        let trigger3 = http_client.add_response_with_trigger(
            "POST /v1/test",
            Ok(HttpResponse {
                status: 200,
                body: r#"{"result":"3"}"#.to_string(),
            }),
        );
        let trigger4 = http_client.add_response_with_trigger(
            "POST /v1/test",
            Ok(HttpResponse {
                status: 200,
                body: r#"{"result":"4"}"#.to_string(),
            }),
        );
        let trigger5 = http_client.add_response_with_trigger(
            "POST /v1/test",
            Ok(HttpResponse {
                status: 200,
                body: r#"{"result":"5"}"#.to_string(),
            }),
        );

        // Setup: Create manager with concurrency limit of 2 for "gpt-4"
        let model_concurrency_limits = Arc::new(dashmap::DashMap::new());
        model_concurrency_limits.insert("gpt-4".to_string(), 2);

        let config = DaemonConfig {
            claim_batch_size: 10,
            claim_interval_ms: 10,
            model_concurrency_limits,

            model_escalations: Arc::new(dashmap::DashMap::new()),
            max_retries: Some(3),
            stop_before_deadline_ms: None,
            backoff_ms: 100,
            backoff_factor: 2,
            max_backoff_ms: 1000,
            first_chunk_timeout_ms: 5000,
            chunk_timeout_ms: 5000,
            body_timeout_ms: 86_400_000,
            status_log_interval_ms: None,
            heartbeat_interval_ms: 5000,
            should_retry: Arc::new(default_should_retry),
            claim_timeout_ms: 60000,
            processing_timeout_ms: 600000,
            stale_daemon_threshold_ms: 30_000,
            unclaim_batch_size: 100,
            batch_metadata_fields: vec![],
            cancellation_poll_interval_ms: 100, // Fast polling for tests
            purge_interval_ms: 0,               // Disabled in tests
            purge_batch_size: 1000,
            purge_throttle_ms: 100,
            streamable_endpoints: Vec::new(),
            throughput_log_interval_ms: None,
            urgency_weight: 0.0,
        };

        let manager = Arc::new(
            PostgresRequestManager::with_client(
                TestDbPools::new(pool.clone()).await.unwrap(),
                http_client.clone(),
            )
            .with_config(config),
        );

        // Setup: Create a file with 5 templates, all using "gpt-4"
        let file_id = manager
            .create_file(
                "test-file".to_string(),
                Some("Test concurrency limits".to_string()),
                vec![
                    crate::RequestTemplateInput {
                        custom_id: None,
                        endpoint: "https://api.example.com".to_string(),
                        method: "POST".to_string(),
                        path: "/v1/test".to_string(),
                        body: r#"{"prompt":"test1"}"#.to_string(),
                        model: "gpt-4".to_string(),
                        api_key: "test-key".to_string(),
                    },
                    crate::RequestTemplateInput {
                        custom_id: None,
                        endpoint: "https://api.example.com".to_string(),
                        method: "POST".to_string(),
                        path: "/v1/test".to_string(),
                        body: r#"{"prompt":"test2"}"#.to_string(),
                        model: "gpt-4".to_string(),
                        api_key: "test-key".to_string(),
                    },
                    crate::RequestTemplateInput {
                        custom_id: None,
                        endpoint: "https://api.example.com".to_string(),
                        method: "POST".to_string(),
                        path: "/v1/test".to_string(),
                        body: r#"{"prompt":"test3"}"#.to_string(),
                        model: "gpt-4".to_string(),
                        api_key: "test-key".to_string(),
                    },
                    crate::RequestTemplateInput {
                        custom_id: None,
                        endpoint: "https://api.example.com".to_string(),
                        method: "POST".to_string(),
                        path: "/v1/test".to_string(),
                        body: r#"{"prompt":"test4"}"#.to_string(),
                        model: "gpt-4".to_string(),
                        api_key: "test-key".to_string(),
                    },
                    crate::RequestTemplateInput {
                        custom_id: None,
                        endpoint: "https://api.example.com".to_string(),
                        method: "POST".to_string(),
                        path: "/v1/test".to_string(),
                        body: r#"{"prompt":"test5"}"#.to_string(),
                        model: "gpt-4".to_string(),
                        api_key: "test-key".to_string(),
                    },
                ],
            )
            .await
            .expect("Failed to create file");

        let batch = manager
            .create_batch(crate::batch::BatchInput {
                file_id,
                endpoint: "/v1/chat/completions".to_string(),
                completion_window: "24h".to_string(),
                metadata: None,
                created_by: None,
                api_key_id: None,
                api_key: None,
                total_requests: None,
            })
            .await
            .expect("Failed to create batch");

        // Start the daemon
        let shutdown_token = tokio_util::sync::CancellationToken::new();
        manager
            .clone()
            .run(shutdown_token.clone())
            .expect("Failed to start daemon");

        // Wait for exactly 2 requests to be in-flight (respecting concurrency limit)
        let start = tokio::time::Instant::now();
        let timeout = Duration::from_secs(2);
        let mut reached_limit = false;

        while start.elapsed() < timeout {
            let in_flight = http_client.in_flight_count();
            if in_flight == 2 {
                reached_limit = true;
                break;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        assert!(
            reached_limit,
            "Expected exactly 2 requests in-flight, got {}",
            http_client.in_flight_count()
        );

        // Verify exactly 2 are in-flight (not more)
        tokio::time::sleep(Duration::from_millis(100)).await;
        assert_eq!(
            http_client.in_flight_count(),
            2,
            "Concurrency limit violated: more than 2 requests in-flight"
        );

        // Trigger completion of first request
        trigger1.send(()).unwrap();

        // Wait for the third request to start
        let start = tokio::time::Instant::now();
        let timeout = Duration::from_secs(2);
        let mut third_started = false;

        while start.elapsed() < timeout {
            if http_client.call_count() >= 3 {
                third_started = true;
                break;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        assert!(
            third_started,
            "Third request should have started after first completed"
        );

        // Verify still only 2 in-flight
        assert_eq!(
            http_client.in_flight_count(),
            2,
            "Should maintain concurrency limit of 2"
        );

        // Complete remaining requests to clean up
        trigger2.send(()).unwrap();
        trigger3.send(()).unwrap();
        trigger4.send(()).unwrap();
        trigger5.send(()).unwrap();

        // Wait for all requests to complete
        let start = tokio::time::Instant::now();
        let timeout = Duration::from_secs(5);
        let mut all_completed = false;

        while start.elapsed() < timeout {
            let status = manager
                .get_batch_status(batch.id)
                .await
                .expect("Failed to get batch status");

            if status.completed_requests == 5 {
                all_completed = true;
                break;
            }
            tokio::time::sleep(Duration::from_millis(50)).await;
        }

        // Stop the daemon
        shutdown_token.cancel();

        assert!(all_completed, "All 5 requests should have completed");

        // Verify all 5 HTTP calls were made
        assert_eq!(http_client.call_count(), 5);
    }

    #[sqlx::test]
    async fn test_daemon_retries_failed_requests(pool: sqlx::PgPool) {
        // Setup: Create HTTP client with failing responses, then success
        let http_client = Arc::new(MockHttpClient::new());

        // First attempt: fails with 500
        http_client.add_response(
            "POST /v1/test",
            Ok(HttpResponse {
                status: 500,
                body: r#"{"error":"internal error"}"#.to_string(),
            }),
        );

        // Second attempt: fails with 503
        http_client.add_response(
            "POST /v1/test",
            Ok(HttpResponse {
                status: 503,
                body: r#"{"error":"service unavailable"}"#.to_string(),
            }),
        );

        // Third attempt: succeeds
        http_client.add_response(
            "POST /v1/test",
            Ok(HttpResponse {
                status: 200,
                body: r#"{"result":"success after retries"}"#.to_string(),
            }),
        );

        // Setup: Create manager with fast backoff for testing
        let config = DaemonConfig {
            claim_batch_size: 10,
            claim_interval_ms: 10,
            model_concurrency_limits: {
                let m = Arc::new(dashmap::DashMap::new());
                m.insert("test-model".to_string(), 10);
                m
            },

            model_escalations: Arc::new(dashmap::DashMap::new()),
            max_retries: Some(5),
            stop_before_deadline_ms: None,
            backoff_ms: 10, // Very fast backoff for testing
            backoff_factor: 2,
            max_backoff_ms: 100,
            first_chunk_timeout_ms: 5000,
            chunk_timeout_ms: 5000,
            body_timeout_ms: 86_400_000,
            status_log_interval_ms: None,
            heartbeat_interval_ms: 5000,
            should_retry: Arc::new(default_should_retry),
            claim_timeout_ms: 60000,
            processing_timeout_ms: 600000,
            stale_daemon_threshold_ms: 30_000,
            unclaim_batch_size: 100,
            batch_metadata_fields: vec![],
            cancellation_poll_interval_ms: 100,
            purge_interval_ms: 0, // Disabled in tests
            purge_batch_size: 1000,
            purge_throttle_ms: 100,
            streamable_endpoints: Vec::new(),
            throughput_log_interval_ms: None,
            urgency_weight: 0.0,
        };

        let manager = Arc::new(
            PostgresRequestManager::with_client(
                TestDbPools::new(pool.clone()).await.unwrap(),
                http_client.clone(),
            )
            .with_config(config),
        );

        // Setup: Create a file and batch
        let file_id = manager
            .create_file(
                "test-file".to_string(),
                Some("Test retry logic".to_string()),
                vec![crate::RequestTemplateInput {
                    custom_id: None,
                    endpoint: "https://api.example.com".to_string(),
                    method: "POST".to_string(),
                    path: "/v1/test".to_string(),
                    body: r#"{"prompt":"test"}"#.to_string(),
                    model: "test-model".to_string(),
                    api_key: "test-key".to_string(),
                }],
            )
            .await
            .expect("Failed to create file");

        let batch = manager
            .create_batch(crate::batch::BatchInput {
                file_id,
                endpoint: "/v1/chat/completions".to_string(),
                completion_window: "24h".to_string(),
                metadata: None,
                created_by: None,
                api_key_id: None,
                api_key: None,
                total_requests: None,
            })
            .await
            .expect("Failed to create batch");

        let requests = manager
            .get_batch_requests(batch.id)
            .await
            .expect("Failed to get batch requests");
        assert_eq!(requests.len(), 1);
        let request_id = requests[0].id();

        // Start the daemon
        let shutdown_token = tokio_util::sync::CancellationToken::new();
        manager
            .clone()
            .run(shutdown_token.clone())
            .expect("Failed to start daemon");

        // Poll for completion (with timeout)
        let start = tokio::time::Instant::now();
        let timeout = Duration::from_secs(5);
        let mut completed = false;

        while start.elapsed() < timeout {
            let results = manager
                .get_requests(vec![request_id])
                .await
                .expect("Failed to get request");

            if let Some(Ok(any_request)) = results.first()
                && let crate::AnyRequest::Completed(req) = any_request
            {
                // Verify the request eventually completed successfully
                assert_eq!(req.state.response_status, 200);
                assert_eq!(
                    req.state.response_body,
                    r#"{"result":"success after retries"}"#
                );
                completed = true;
                break;
            }

            tokio::time::sleep(Duration::from_millis(50)).await;
        }

        // Stop the daemon
        shutdown_token.cancel();

        assert!(completed, "Request should have completed after retries");

        // Verify the request was attempted 3 times (2 failures + 1 success)
        assert_eq!(
            http_client.call_count(),
            3,
            "Expected 3 HTTP calls (2 failed attempts + 1 success)"
        );
    }

    #[sqlx::test]
    async fn test_daemon_dynamically_updates_concurrency_limits(pool: sqlx::PgPool) {
        // Setup: Create HTTP client with triggered responses
        let http_client = Arc::new(MockHttpClient::new());

        // Add 10 triggered responses
        let mut triggers = vec![];
        for i in 1..=10 {
            let trigger = http_client.add_response_with_trigger(
                "POST /v1/test",
                Ok(HttpResponse {
                    status: 200,
                    body: format!(r#"{{"result":"{}"}}"#, i),
                }),
            );
            triggers.push(trigger);
        }

        // Setup: Start with concurrency limit of 2 for "gpt-4"
        let model_concurrency_limits = Arc::new(dashmap::DashMap::new());
        model_concurrency_limits.insert("gpt-4".to_string(), 2);

        let config = DaemonConfig {
            claim_batch_size: 10,
            claim_interval_ms: 10,
            model_concurrency_limits: model_concurrency_limits.clone(),

            model_escalations: Arc::new(dashmap::DashMap::new()),
            max_retries: Some(3),
            stop_before_deadline_ms: None,
            backoff_ms: 100,
            backoff_factor: 2,
            max_backoff_ms: 1000,
            first_chunk_timeout_ms: 5000,
            chunk_timeout_ms: 5000,
            body_timeout_ms: 86_400_000,
            status_log_interval_ms: None,
            heartbeat_interval_ms: 5000,
            should_retry: Arc::new(default_should_retry),
            claim_timeout_ms: 60000,
            processing_timeout_ms: 600000,
            stale_daemon_threshold_ms: 30_000,
            unclaim_batch_size: 100,
            batch_metadata_fields: vec![],
            cancellation_poll_interval_ms: 100,
            purge_interval_ms: 0, // Disabled in tests
            purge_batch_size: 1000,
            purge_throttle_ms: 100,
            streamable_endpoints: Vec::new(),
            throughput_log_interval_ms: None,
            urgency_weight: 0.0,
        };

        let manager = Arc::new(
            PostgresRequestManager::with_client(
                TestDbPools::new(pool.clone()).await.unwrap(),
                http_client.clone(),
            )
            .with_config(config),
        );

        // Setup: Create a file with 10 requests, all using "gpt-4"
        let templates: Vec<_> = (1..=10)
            .map(|i| crate::RequestTemplateInput {
                custom_id: None,
                endpoint: "https://api.example.com".to_string(),
                method: "POST".to_string(),
                path: "/v1/test".to_string(),
                body: format!(r#"{{"prompt":"test{}"}}"#, i),
                model: "gpt-4".to_string(),
                api_key: "test-key".to_string(),
            })
            .collect();

        let file_id = manager
            .create_file(
                "test-file".to_string(),
                Some("Test dynamic limits".to_string()),
                templates,
            )
            .await
            .expect("Failed to create file");

        let batch = manager
            .create_batch(crate::batch::BatchInput {
                file_id,
                endpoint: "/v1/chat/completions".to_string(),
                completion_window: "24h".to_string(),
                metadata: None,
                created_by: None,
                api_key_id: None,
                api_key: None,
                total_requests: None,
            })
            .await
            .expect("Failed to create batch");

        // Start the daemon
        let shutdown_token = tokio_util::sync::CancellationToken::new();
        manager
            .clone()
            .run(shutdown_token.clone())
            .expect("Failed to start daemon");

        // Wait for exactly 2 requests to be in-flight (initial limit)
        let start = tokio::time::Instant::now();
        let timeout = Duration::from_secs(2);
        let mut reached_initial_limit = false;

        while start.elapsed() < timeout {
            let in_flight = http_client.in_flight_count();
            if in_flight == 2 {
                reached_initial_limit = true;
                break;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        assert!(
            reached_initial_limit,
            "Expected exactly 2 requests in-flight with initial limit"
        );

        // Increase the limit to 5
        model_concurrency_limits.insert("gpt-4".to_string(), 5);

        // Wait a bit for the daemon to pick up the new limit
        tokio::time::sleep(Duration::from_millis(100)).await;

        // Complete one request to free up a permit and trigger daemon to check limits
        triggers.remove(0).send(()).unwrap();
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Now we should see up to 5 requests in flight
        let start = tokio::time::Instant::now();
        let timeout = Duration::from_secs(2);
        let mut reached_new_limit = false;

        while start.elapsed() < timeout {
            let in_flight = http_client.in_flight_count();
            if in_flight >= 4 {
                // Should see at least 4-5 in flight with new limit
                reached_new_limit = true;
                break;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        assert!(
            reached_new_limit,
            "Expected more requests in-flight after limit increase, got {}",
            http_client.in_flight_count()
        );

        // Now decrease the limit to 3
        model_concurrency_limits.insert("gpt-4".to_string(), 3);

        // Complete remaining requests
        for trigger in triggers {
            trigger.send(()).unwrap();
        }

        // Wait for all requests to complete
        let start = tokio::time::Instant::now();
        let timeout = Duration::from_secs(5);
        let mut all_completed = false;

        while start.elapsed() < timeout {
            let status = manager
                .get_batch_status(batch.id)
                .await
                .expect("Failed to get batch status");

            if status.completed_requests == 10 {
                all_completed = true;
                break;
            }
            tokio::time::sleep(Duration::from_millis(50)).await;
        }

        // Stop the daemon
        shutdown_token.cancel();

        assert!(all_completed, "All 10 requests should have completed");
        assert_eq!(http_client.call_count(), 10);
    }

    #[sqlx::test]
    async fn test_deadline_aware_retry_stops_before_deadline(pool: sqlx::PgPool) {
        // Test that retries stop when approaching the deadline
        let http_client = Arc::new(MockHttpClient::new());

        // All requests will fail
        for _ in 0..20 {
            http_client.add_response(
                "POST /v1/test",
                Ok(HttpResponse {
                    status: 500,
                    body: r#"{"error":"server error"}"#.to_string(),
                }),
            );
        }

        // Use deadline-aware retry with a short completion window and short buffer
        let config = DaemonConfig {
            claim_batch_size: 10,
            claim_interval_ms: 10,
            model_concurrency_limits: {
                let m = Arc::new(dashmap::DashMap::new());
                m.insert("test-model".to_string(), 10);
                m
            },

            model_escalations: Arc::new(dashmap::DashMap::new()),
            max_retries: Some(10_000),
            stop_before_deadline_ms: Some(500), // 500ms buffer before deadline
            backoff_ms: 50,
            backoff_factor: 2,
            max_backoff_ms: 200,
            first_chunk_timeout_ms: 5000,
            chunk_timeout_ms: 5000,
            body_timeout_ms: 86_400_000,
            status_log_interval_ms: None,
            heartbeat_interval_ms: 5000,
            should_retry: Arc::new(default_should_retry),
            claim_timeout_ms: 60000,
            processing_timeout_ms: 600000,
            stale_daemon_threshold_ms: 30_000,
            unclaim_batch_size: 100,
            batch_metadata_fields: vec![],
            cancellation_poll_interval_ms: 100,
            purge_interval_ms: 0, // Disabled in tests
            purge_batch_size: 1000,
            purge_throttle_ms: 100,
            streamable_endpoints: Vec::new(),
            throughput_log_interval_ms: None,
            urgency_weight: 0.0,
        };

        let manager = Arc::new(
            PostgresRequestManager::with_client(
                TestDbPools::new(pool.clone()).await.unwrap(),
                http_client.clone(),
            )
            .with_config(config),
        );

        // Create a batch with a very short completion window (2 seconds)
        let file_id = manager
            .create_file(
                "test-file".to_string(),
                Some("Test deadline cutoff".to_string()),
                vec![crate::RequestTemplateInput {
                    custom_id: None,
                    endpoint: "https://api.example.com".to_string(),
                    method: "POST".to_string(),
                    path: "/v1/test".to_string(),
                    body: r#"{"prompt":"test"}"#.to_string(),
                    model: "test-model".to_string(),
                    api_key: "test-key".to_string(),
                }],
            )
            .await
            .expect("Failed to create file");

        let batch = manager
            .create_batch(crate::batch::BatchInput {
                file_id,
                endpoint: "/v1/chat/completions".to_string(),
                completion_window: "2s".to_string(), // Very short window
                metadata: None,
                created_by: None,
                api_key_id: None,
                api_key: None,
                total_requests: None,
            })
            .await
            .expect("Failed to create batch");

        let requests = manager
            .get_batch_requests(batch.id)
            .await
            .expect("Failed to get batch requests");
        let request_id = requests[0].id();

        // Start the daemon
        let shutdown_token = tokio_util::sync::CancellationToken::new();
        manager
            .clone()
            .run(shutdown_token.clone())
            .expect("Failed to start daemon");

        // Wait for the deadline to pass
        tokio::time::sleep(Duration::from_secs(3)).await;

        // Check the request state
        let results = manager
            .get_requests(vec![request_id])
            .await
            .expect("Failed to get request");

        shutdown_token.cancel();

        if let Some(Ok(crate::AnyRequest::Failed(failed))) = results.first() {
            // Calculate expected retry attempts:
            // - Completion window: 2000ms
            // - Buffer: 500ms
            // - Effective deadline: 1500ms
            // - Backoff sequence: 50ms, 100ms, 200ms, 200ms, 200ms, 200ms, 200ms
            // - Timeline:
            //   - Initial attempt: t=0ms (attempt 0)
            //   - Retry 1: t=50ms (attempt 1)
            //   - Retry 2: t=150ms (attempt 2)
            //   - Retry 3: t=350ms (attempt 3)
            //   - Retry 4: t=550ms (attempt 4)
            //   - Retry 5: t=750ms (attempt 5)
            //   - Retry 6: t=950ms (attempt 6)
            //   - Retry 7: t=1150ms (attempt 7)
            //   - Retry 8: t=1350ms (attempt 8)
            //   - Next would be t=1550ms - EXCEEDS 1500ms deadline
            // Expected: 8 retry attempts (9 total including initial)

            let retry_count = failed.state.retry_attempt;
            let call_count = http_client.call_count();

            // 1. Verify we stopped before too many retries (deadline constraint)
            // Allow 4-9 attempts to account for timing variations in test execution,
            // parallel test execution overhead, and query overhead from batch metadata fields
            assert!(
                (4..=9).contains(&retry_count),
                "Expected 4-9 retry attempts based on deadline and backoff calculation, got {}",
                retry_count
            );

            // 2. Verify HTTP call count matches retry attempts (1 initial + N retries)
            assert_eq!(
                call_count,
                (retry_count + 1) as usize,
                "Expected call count to match retry attempts + 1 initial attempt, got {} calls for {} retry attempts",
                call_count,
                retry_count
            );

            // 3. Verify the request actually has error details from the last attempt
            assert!(
                !failed.state.reason.to_error_message().is_empty(),
                "Expected failed request to have failure reason"
            );
        } else {
            panic!(
                "Expected request to be in Failed state, got {:?}",
                results.first()
            );
        }
    }

    #[sqlx::test]
    async fn test_retry_stops_at_deadline_when_no_limits_set(pool: sqlx::PgPool) {
        // Test that when neither max_retries nor stop_before_deadline_ms is set,
        // retries stop exactly at the deadline (no buffer)
        let http_client = Arc::new(MockHttpClient::new());

        // All requests will fail
        for _ in 0..20 {
            http_client.add_response(
                "POST /v1/test",
                Ok(HttpResponse {
                    status: 500,
                    body: r#"{"error":"server error"}"#.to_string(),
                }),
            );
        }

        // No max_retries, no stop_before_deadline_ms
        let config = DaemonConfig {
            claim_batch_size: 10,
            claim_interval_ms: 10,
            model_concurrency_limits: {
                let m = Arc::new(dashmap::DashMap::new());
                m.insert("test-model".to_string(), 10);
                m
            },

            model_escalations: Arc::new(dashmap::DashMap::new()),
            max_retries: None,             // No retry limit
            stop_before_deadline_ms: None, // No buffer - should retry until deadline
            backoff_ms: 50,
            backoff_factor: 2,
            max_backoff_ms: 200,
            first_chunk_timeout_ms: 5000,
            chunk_timeout_ms: 5000,
            body_timeout_ms: 86_400_000,
            status_log_interval_ms: None,
            heartbeat_interval_ms: 5000,
            should_retry: Arc::new(default_should_retry),
            claim_timeout_ms: 60000,
            processing_timeout_ms: 600000,
            stale_daemon_threshold_ms: 30_000,
            unclaim_batch_size: 100,
            batch_metadata_fields: vec![],
            cancellation_poll_interval_ms: 100,
            purge_interval_ms: 0, // Disabled in tests
            purge_batch_size: 1000,
            purge_throttle_ms: 100,
            streamable_endpoints: Vec::new(),
            throughput_log_interval_ms: None,
            urgency_weight: 0.0,
        };

        let manager = Arc::new(
            PostgresRequestManager::with_client(
                TestDbPools::new(pool.clone()).await.unwrap(),
                http_client.clone(),
            )
            .with_config(config),
        );

        // Create a batch with a 2 second completion window
        let file_id = manager
            .create_file(
                "test-file".to_string(),
                Some("Test no limits retry".to_string()),
                vec![crate::RequestTemplateInput {
                    custom_id: None,
                    endpoint: "https://api.example.com".to_string(),
                    method: "POST".to_string(),
                    path: "/v1/test".to_string(),
                    body: r#"{"prompt":"test"}"#.to_string(),
                    model: "test-model".to_string(),
                    api_key: "test-key".to_string(),
                }],
            )
            .await
            .expect("Failed to create file");

        let batch = manager
            .create_batch(crate::batch::BatchInput {
                file_id,
                endpoint: "/v1/chat/completions".to_string(),
                completion_window: "2s".to_string(),
                metadata: None,
                created_by: None,
                api_key_id: None,
                api_key: None,
                total_requests: None,
            })
            .await
            .expect("Failed to create batch");

        let requests = manager
            .get_batch_requests(batch.id)
            .await
            .expect("Failed to get batch requests");
        let request_id = requests[0].id();

        // Start the daemon
        let shutdown_token = tokio_util::sync::CancellationToken::new();
        manager
            .clone()
            .run(shutdown_token.clone())
            .expect("Failed to start daemon");

        // Wait for the deadline to pass
        tokio::time::sleep(Duration::from_secs(3)).await;

        // Check the request state
        let results = manager
            .get_requests(vec![request_id])
            .await
            .expect("Failed to get request");

        shutdown_token.cancel();

        if let Some(Ok(crate::AnyRequest::Failed(failed))) = results.first() {
            // Calculate expected retry attempts with NO buffer:
            // - Completion window: 2000ms
            // - Buffer: 0ms (none set)
            // - Effective deadline: 2000ms
            // - Backoff sequence: 50ms, 100ms, 200ms, 200ms, 200ms...
            // - Timeline:
            //   - Initial attempt: t=0ms (attempt 0)
            //   - Retry 1: t=50ms (attempt 1)
            //   - Retry 2: t=150ms (attempt 2)
            //   - Retry 3: t=350ms (attempt 3)
            //   - Retry 4: t=550ms (attempt 4)
            //   - Retry 5: t=750ms (attempt 5)
            //   - Retry 6: t=950ms (attempt 6)
            //   - Retry 7: t=1150ms (attempt 7)
            //   - Retry 8: t=1350ms (attempt 8)
            //   - Retry 9: t=1550ms (attempt 9)
            //   - Retry 10: t=1750ms (attempt 10)
            //   - Retry 11: t=1950ms (attempt 11)
            //   - Next would be t=2150ms - EXCEEDS 2000ms deadline
            // Expected: ~11 retry attempts (12 total including initial)
            // In reality, we will see <11 due to DB calls and CPU overhead in making requests

            let retry_count = failed.state.retry_attempt;
            let call_count = http_client.call_count();

            // 1. Verify we retried more than the buffered case (which stopped at ~8)
            //    but still stopped before too many attempts
            // Allow 6-12 attempts to account for timing variations with CI slower CI CPUs,
            // parallel test execution overhead, and query overhead from batch metadata fields
            assert!(
                (6..12).contains(&retry_count),
                "Expected 6-12 retry attempts (should retry until deadline with no buffer), got {}",
                retry_count
            );

            // 2. Verify HTTP call count matches retry attempts (1 initial + N retries)
            assert_eq!(
                call_count,
                (retry_count + 1) as usize,
                "Expected call count to match retry attempts + 1 initial attempt, got {} calls for {} retry attempts",
                call_count,
                retry_count
            );

            // 3. Verify the request has error details from the last attempt
            assert!(
                !failed.state.reason.to_error_message().is_empty(),
                "Expected failed request to have failure reason"
            );
        } else {
            panic!(
                "Expected request to be in Failed state, got {:?}",
                results.first()
            );
        }
    }

    #[sqlx::test]
    #[test_log::test]
    async fn test_batch_metadata_headers_passed_through(pool: sqlx::PgPool) {
        let http_client = crate::http::MockHttpClient::new();
        http_client.add_response(
            "POST /v1/chat/completions",
            Ok(crate::http::HttpResponse {
                status: 200,
                body: r#"{"id":"chatcmpl-123","choices":[{"message":{"content":"test"}}]}"#
                    .to_string(),
            }),
        );

        let config = DaemonConfig {
            claim_batch_size: 10,
            claim_interval_ms: 10,
            model_concurrency_limits: {
                let m = Arc::new(dashmap::DashMap::new());
                m.insert("test-model".to_string(), 10);
                m
            },

            model_escalations: Arc::new(dashmap::DashMap::new()),
            max_retries: Some(3),
            stop_before_deadline_ms: None,
            backoff_ms: 100,
            backoff_factor: 2,
            max_backoff_ms: 1000,
            first_chunk_timeout_ms: 5000,
            chunk_timeout_ms: 5000,
            body_timeout_ms: 86_400_000,
            status_log_interval_ms: None,
            heartbeat_interval_ms: 5000,
            should_retry: Arc::new(default_should_retry),
            claim_timeout_ms: 60000,
            processing_timeout_ms: 600000,
            stale_daemon_threshold_ms: 30_000,
            unclaim_batch_size: 100,
            batch_metadata_fields: vec![
                "id".to_string(),
                "endpoint".to_string(),
                "created_at".to_string(),
                "completion_window".to_string(),
            ],
            cancellation_poll_interval_ms: 100,
            purge_interval_ms: 0, // Disabled in tests
            purge_batch_size: 1000,
            purge_throttle_ms: 100,
            streamable_endpoints: Vec::new(),
            throughput_log_interval_ms: None,
            urgency_weight: 0.0,
        };

        let manager = Arc::new(
            PostgresRequestManager::with_client(
                TestDbPools::new(pool.clone()).await.unwrap(),
                Arc::new(http_client.clone()),
            )
            .with_config(config),
        );

        // Create a batch
        let file_id = manager
            .create_file(
                "test-file".to_string(),
                Some("Test batch metadata".to_string()),
                vec![crate::RequestTemplateInput {
                    custom_id: None,
                    endpoint: "https://api.example.com".to_string(),
                    method: "POST".to_string(),
                    path: "/v1/chat/completions".to_string(),
                    body: r#"{"prompt":"test"}"#.to_string(),
                    model: "test-model".to_string(),
                    api_key: "test-key".to_string(),
                }],
            )
            .await
            .expect("Failed to create file");

        let batch = manager
            .create_batch(crate::batch::BatchInput {
                file_id,
                endpoint: "/v1/chat/completions".to_string(),
                completion_window: "24h".to_string(),
                metadata: None,
                created_by: Some("test-user".to_string()),
                api_key_id: None,
                api_key: None,
                total_requests: None,
            })
            .await
            .expect("Failed to create batch");

        let requests = manager
            .get_batch_requests(batch.id)
            .await
            .expect("Failed to get batch requests");
        let request_id = requests[0].id();

        // Start the daemon
        let shutdown_token = tokio_util::sync::CancellationToken::new();
        manager
            .clone()
            .run(shutdown_token.clone())
            .expect("Failed to start daemon");

        // Wait for request to be processed
        tokio::time::sleep(Duration::from_millis(500)).await;

        shutdown_token.cancel();

        // Wait a bit for shutdown
        tokio::time::sleep(Duration::from_millis(200)).await;

        // Verify the request was completed
        let results = manager
            .get_requests(vec![request_id])
            .await
            .expect("Failed to get request");

        assert_eq!(results.len(), 1);
        assert!(
            matches!(results[0], Ok(crate::AnyRequest::Completed(_))),
            "Expected request to be completed"
        );

        // Verify batch metadata was passed to HTTP client
        let calls = http_client.get_calls();
        assert_eq!(calls.len(), 1, "Expected exactly one HTTP call");

        let call = &calls[0];
        assert_eq!(
            call.batch_metadata.len(),
            4,
            "Expected 4 batch metadata fields"
        );

        // Verify each configured field was passed through
        assert!(
            call.batch_metadata.contains_key("id"),
            "Expected batch id in metadata"
        );
        assert!(
            call.batch_metadata.contains_key("endpoint"),
            "Expected batch endpoint in metadata"
        );
        assert!(
            call.batch_metadata.contains_key("created_at"),
            "Expected batch created_at in metadata"
        );
        assert!(
            call.batch_metadata.contains_key("completion_window"),
            "Expected batch completion_window in metadata"
        );

        // Verify values are correct
        assert_eq!(
            call.batch_metadata.get("endpoint"),
            Some(&"/v1/chat/completions".to_string()),
            "Batch endpoint should match"
        );
        assert_eq!(
            call.batch_metadata.get("completion_window"),
            Some(&"24h".to_string()),
            "Completion window should match"
        );
    }

    #[sqlx::test]
    #[test_log::test]
    async fn test_batch_metadata_extracts_fields_from_json_metadata(pool: sqlx::PgPool) {
        let http_client = crate::http::MockHttpClient::new();
        http_client.add_response(
            "POST /v1/chat/completions",
            Ok(crate::http::HttpResponse {
                status: 200,
                body: r#"{"id":"chatcmpl-123","choices":[{"message":{"content":"test"}}]}"#
                    .to_string(),
            }),
        );

        // Configure batch_metadata_fields to include "request_source" which is stored
        // inside the metadata JSON, not as a direct column
        let config = DaemonConfig {
            claim_batch_size: 10,
            claim_interval_ms: 10,
            model_concurrency_limits: {
                let m = Arc::new(dashmap::DashMap::new());
                m.insert("test-model".to_string(), 10);
                m
            },

            model_escalations: Arc::new(dashmap::DashMap::new()),
            max_retries: Some(3),
            stop_before_deadline_ms: None,
            backoff_ms: 100,
            backoff_factor: 2,
            max_backoff_ms: 1000,
            first_chunk_timeout_ms: 5000,
            chunk_timeout_ms: 5000,
            body_timeout_ms: 86_400_000,
            status_log_interval_ms: None,
            heartbeat_interval_ms: 5000,
            should_retry: Arc::new(default_should_retry),
            claim_timeout_ms: 60000,
            processing_timeout_ms: 600000,
            stale_daemon_threshold_ms: 30_000,
            unclaim_batch_size: 100,
            batch_metadata_fields: vec![
                "id".to_string(),
                "endpoint".to_string(),
                "completion_window".to_string(),
                "request_source".to_string(), // This comes from metadata JSON
            ],
            cancellation_poll_interval_ms: 100,
            purge_interval_ms: 0, // Disabled in tests
            purge_batch_size: 1000,
            purge_throttle_ms: 100,
            streamable_endpoints: Vec::new(),
            throughput_log_interval_ms: None,
            urgency_weight: 0.0,
        };

        let manager = Arc::new(
            PostgresRequestManager::with_client(
                TestDbPools::new(pool.clone()).await.unwrap(),
                Arc::new(http_client.clone()),
            )
            .with_config(config),
        );

        // Create a batch with metadata containing request_source
        let file_id = manager
            .create_file(
                "test-file".to_string(),
                Some("Test metadata JSON extraction".to_string()),
                vec![crate::RequestTemplateInput {
                    custom_id: None,
                    endpoint: "https://api.example.com".to_string(),
                    method: "POST".to_string(),
                    path: "/v1/chat/completions".to_string(),
                    body: r#"{"prompt":"test"}"#.to_string(),
                    model: "test-model".to_string(),
                    api_key: "test-key".to_string(),
                }],
            )
            .await
            .expect("Failed to create file");

        let batch = manager
            .create_batch(crate::batch::BatchInput {
                file_id,
                endpoint: "/v1/chat/completions".to_string(),
                completion_window: "24h".to_string(),
                metadata: Some(serde_json::json!({
                    "request_source": "api",
                    "created_by": "user-123"
                })),
                created_by: Some("test-user".to_string()),
                api_key_id: None,
                api_key: None,
                total_requests: None,
            })
            .await
            .expect("Failed to create batch");

        let requests = manager
            .get_batch_requests(batch.id)
            .await
            .expect("Failed to get batch requests");
        let request_id = requests[0].id();

        // Start the daemon
        let shutdown_token = tokio_util::sync::CancellationToken::new();
        manager
            .clone()
            .run(shutdown_token.clone())
            .expect("Failed to start daemon");

        // Wait for request to be processed
        tokio::time::sleep(Duration::from_millis(500)).await;

        shutdown_token.cancel();

        // Wait a bit for shutdown
        tokio::time::sleep(Duration::from_millis(200)).await;

        // Verify the request was completed
        let results = manager
            .get_requests(vec![request_id])
            .await
            .expect("Failed to get request");

        assert_eq!(results.len(), 1);
        assert!(
            matches!(results[0], Ok(crate::AnyRequest::Completed(_))),
            "Expected request to be completed"
        );

        // Verify batch metadata was passed to HTTP client
        let calls = http_client.get_calls();
        assert_eq!(calls.len(), 1, "Expected exactly one HTTP call");

        let call = &calls[0];
        assert_eq!(
            call.batch_metadata.len(),
            4,
            "Expected 4 batch metadata fields (id, endpoint, completion_window, request_source)"
        );

        // Verify direct column fields
        assert!(
            call.batch_metadata.contains_key("id"),
            "Expected batch id in metadata"
        );
        assert_eq!(
            call.batch_metadata.get("endpoint"),
            Some(&"/v1/chat/completions".to_string()),
            "Batch endpoint should match"
        );

        // Verify request_source was extracted from metadata JSON
        assert_eq!(
            call.batch_metadata.get("request_source"),
            Some(&"api".to_string()),
            "request_source should be extracted from metadata JSON"
        );
    }

    #[sqlx::test]
    #[test_log::test]
    async fn test_batch_api_key_overrides_template_api_key(pool: sqlx::PgPool) {
        // Setup: Create HTTP client with mock response
        let http_client = Arc::new(MockHttpClient::new());
        http_client.add_response(
            "POST /v1/test",
            Ok(HttpResponse {
                status: 200,
                body: r#"{"result":"success"}"#.to_string(),
            }),
        );

        let config = DaemonConfig {
            claim_batch_size: 10,
            claim_interval_ms: 10,
            model_concurrency_limits: {
                let m = Arc::new(dashmap::DashMap::new());
                m.insert("test-model".to_string(), 10);
                m
            },
            model_escalations: Arc::new(dashmap::DashMap::new()),
            max_retries: Some(3),
            stop_before_deadline_ms: None,
            backoff_ms: 100,
            backoff_factor: 2,
            max_backoff_ms: 1000,
            first_chunk_timeout_ms: 5000,
            chunk_timeout_ms: 5000,
            body_timeout_ms: 86_400_000,
            status_log_interval_ms: None,
            heartbeat_interval_ms: 5000,
            should_retry: Arc::new(default_should_retry),
            claim_timeout_ms: 60000,
            processing_timeout_ms: 600000,
            stale_daemon_threshold_ms: 30_000,
            unclaim_batch_size: 100,
            batch_metadata_fields: vec![],
            cancellation_poll_interval_ms: 100,
            purge_interval_ms: 0,
            purge_batch_size: 1000,
            purge_throttle_ms: 100,
            streamable_endpoints: Vec::new(),
            throughput_log_interval_ms: None,
            urgency_weight: 0.0,
        };

        let manager = Arc::new(
            PostgresRequestManager::with_client(
                TestDbPools::new(pool.clone()).await.unwrap(),
                http_client.clone(),
            )
            .with_config(config),
        );

        // Create file with template using the file uploader's key
        let file_id = manager
            .create_file(
                "test-file".to_string(),
                Some("Test file".to_string()),
                vec![crate::RequestTemplateInput {
                    custom_id: None,
                    endpoint: "https://api.example.com".to_string(),
                    method: "POST".to_string(),
                    path: "/v1/test".to_string(),
                    body: r#"{"prompt":"test"}"#.to_string(),
                    model: "test-model".to_string(),
                    api_key: "template-key".to_string(),
                }],
            )
            .await
            .expect("Failed to create file");

        // Create batch with a different key (the batch creator's key)
        let batch = manager
            .create_batch(crate::batch::BatchInput {
                file_id,
                endpoint: "/v1/chat/completions".to_string(),
                completion_window: "24h".to_string(),
                metadata: None,
                created_by: None,
                api_key_id: None,
                api_key: Some("batch-creator-key".to_string()),
                total_requests: None,
            })
            .await
            .expect("Failed to create batch");

        let requests = manager
            .get_batch_requests(batch.id)
            .await
            .expect("Failed to get batch requests");
        assert_eq!(requests.len(), 1);
        let request_id = requests[0].id();

        // Start the daemon
        let shutdown_token = tokio_util::sync::CancellationToken::new();
        manager
            .clone()
            .run(shutdown_token.clone())
            .expect("Failed to start daemon");

        // Poll for completion
        let start = tokio::time::Instant::now();
        let timeout = Duration::from_secs(5);
        let mut completed = false;

        while start.elapsed() < timeout {
            let results = manager
                .get_requests(vec![request_id])
                .await
                .expect("Failed to get request");

            if let Some(Ok(any_request)) = results.first()
                && any_request.is_terminal()
            {
                completed = true;
                break;
            }

            tokio::time::sleep(Duration::from_millis(50)).await;
        }

        shutdown_token.cancel();
        assert!(completed, "Request did not complete within timeout");

        // Verify the daemon used the batch creator's key, not the template key
        assert_eq!(http_client.call_count(), 1);
        let calls = http_client.get_calls();
        assert_eq!(
            calls[0].api_key, "batch-creator-key",
            "Daemon should use batch creator's API key, not the template's"
        );
    }
}
