//! Databento downloader module
//!
//! Handles downloading data from Databento API with retry logic,
//! rate limiting, and cost tracking.
//!
//! Uses the Databento REST API (`/v0/timeseries.get_range`) to stream
//! historical market data in DBN (Databento Binary) format and persist
//! it to local files. Each download job gets its own sub-directory
//! under the configured output directory.

use crate::error::{AcquisitionError, AcquisitionResult};
use reqwest::Client as HttpClient;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::AsyncWriteExt;
use tokio::sync::RwLock;
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, info, warn};
use uuid::Uuid;

/// Base URL for the Databento historical API.
const DATABENTO_HIST_BASE_URL: &str = "https://hist.databento.com";

/// Databento download configuration
#[derive(Debug, Clone)]
pub struct DatabentoDownloaderConfig {
    /// Databento API key (sourced from DATABENTO_API_KEY env var by default)
    pub api_key: String,
    /// Root directory where downloaded files are stored
    pub output_dir: PathBuf,
    /// Maximum number of retry attempts per download
    pub max_retries: usize,
    /// Initial delay between retries in milliseconds (grows via exponential backoff)
    pub retry_delay_ms: u64,
    /// HTTP request timeout in seconds
    pub request_timeout_secs: u64,
    /// Optional override of the Databento historical API base URL
    pub base_url: Option<String>,
}

impl Default for DatabentoDownloaderConfig {
    fn default() -> Self {
        Self {
            api_key: std::env::var("DATABENTO_API_KEY").unwrap_or_default(),
            output_dir: PathBuf::from("data/downloads"),
            max_retries: 3,
            retry_delay_ms: 1000,
            request_timeout_secs: 300,
            base_url: None,
        }
    }
}

/// Download job specification
#[derive(Debug, Clone)]
pub struct DownloadJob {
    pub job_id: Uuid,
    pub dataset: String,
    pub symbols: Vec<String>,
    pub start_date: String,
    pub end_date: String,
    pub schema: String,
    pub priority: i32,
}

/// Progress information for an active download
#[derive(Debug, Clone)]
pub struct DownloadProgress {
    /// Bytes received so far
    pub bytes_downloaded: u64,
    /// Total bytes expected (if known from Content-Length)
    pub total_bytes: Option<u64>,
}

/// Databento API downloader
///
/// Manages HTTP communication with the Databento historical REST API,
/// including authentication, retry logic with exponential backoff,
/// per-job cancellation, and streaming downloads to disk.
pub struct DatabentoDownloader {
    config: DatabentoDownloaderConfig,
    http_client: HttpClient,
    /// Active cancellation tokens keyed by job ID
    active_jobs: Arc<RwLock<HashMap<Uuid, CancellationToken>>>,
}

impl DatabentoDownloader {
    /// Create new downloader
    pub fn new(config: DatabentoDownloaderConfig) -> Self {
        let timeout = Duration::from_secs(config.request_timeout_secs);
        let http_client = HttpClient::builder()
            .timeout(timeout)
            .pool_idle_timeout(Duration::from_secs(30))
            .pool_max_idle_per_host(4)
            .build()
            .unwrap_or_default();

        Self {
            config,
            http_client,
            active_jobs: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Return the effective API base URL
    fn base_url(&self) -> &str {
        self.config
            .base_url
            .as_deref()
            .unwrap_or(DATABENTO_HIST_BASE_URL)
    }

    /// Validate that the API key is present and non-empty
    fn validate_api_key(&self) -> AcquisitionResult<()> {
        if self.config.api_key.is_empty() {
            return Err(AcquisitionError::Authentication {
                message: "DATABENTO_API_KEY is not set or empty".to_string(),
            });
        }
        Ok(())
    }

    /// Validate job parameters before making an API request
    fn validate_job(job: &DownloadJob) -> AcquisitionResult<()> {
        if job.dataset.is_empty() {
            return Err(AcquisitionError::InvalidRequest {
                message: "dataset must not be empty".to_string(),
            });
        }
        if job.symbols.is_empty() {
            return Err(AcquisitionError::InvalidRequest {
                message: "symbols must not be empty".to_string(),
            });
        }
        if job.start_date.is_empty() || job.end_date.is_empty() {
            return Err(AcquisitionError::InvalidRequest {
                message: "start_date and end_date must not be empty".to_string(),
            });
        }
        if job.schema.is_empty() {
            return Err(AcquisitionError::InvalidRequest {
                message: "schema must not be empty".to_string(),
            });
        }
        Ok(())
    }

    /// Download data for a job.
    ///
    /// Streams the response body from the Databento `timeseries.get_range`
    /// endpoint to a local `.dbn.zst` file in `{output_dir}/{job_id}/`.
    /// Implements retry with exponential backoff and per-job cancellation.
    ///
    /// Returns the path to the downloaded file on success.
    pub async fn download(&self, job: &DownloadJob) -> AcquisitionResult<PathBuf> {
        self.validate_api_key()?;
        Self::validate_job(job)?;

        // Register cancellation token for this job
        let cancel_token = CancellationToken::new();
        {
            let mut jobs = self.active_jobs.write().await;
            jobs.insert(job.job_id, cancel_token.clone());
        }

        // Ensure the cleanup runs even on early return
        let result = self
            .download_with_retry(job, &cancel_token)
            .await;

        // Unregister the cancellation token
        {
            let mut jobs = self.active_jobs.write().await;
            jobs.remove(&job.job_id);
        }

        result
    }

    /// Inner download loop with retry and exponential backoff
    async fn download_with_retry(
        &self,
        job: &DownloadJob,
        cancel_token: &CancellationToken,
    ) -> AcquisitionResult<PathBuf> {
        let mut delay = Duration::from_millis(self.config.retry_delay_ms);

        for attempt in 0..=self.config.max_retries {
            if cancel_token.is_cancelled() {
                return Err(AcquisitionError::Internal {
                    message: format!("Download cancelled for job {}", job.job_id),
                });
            }

            match self.execute_download(job, cancel_token).await {
                Ok(path) => return Ok(path),
                Err(e) => {
                    if attempt >= self.config.max_retries {
                        error!(
                            job_id = %job.job_id,
                            attempts = attempt + 1,
                            "Download failed after all retry attempts: {}", e
                        );
                        return Err(e);
                    }

                    // Do not retry on non-transient errors
                    if Self::is_non_retryable(&e) {
                        error!(
                            job_id = %job.job_id,
                            "Non-retryable error, aborting: {}", e
                        );
                        return Err(e);
                    }

                    info!(
                        job_id = %job.job_id,
                        attempt = attempt + 1,
                        max_retries = self.config.max_retries,
                        delay_ms = delay.as_millis() as u64,
                        "Download attempt did not succeed: {}. Retrying...", e
                    );

                    tokio::select! {
                        () = tokio::time::sleep(delay) => {},
                        () = cancel_token.cancelled() => {
                            return Err(AcquisitionError::Internal {
                                message: format!("Download cancelled during retry backoff for job {}", job.job_id),
                            });
                        }
                    }

                    // Exponential backoff with factor 2, capped at 30 seconds
                    delay = (delay * 2).min(Duration::from_secs(30));
                },
            }
        }

        // Unreachable in practice, but satisfies the compiler
        Err(AcquisitionError::Internal {
            message: "Exhausted retry attempts".to_string(),
        })
    }

    /// Determine whether an error is non-retryable (e.g. auth, bad request)
    fn is_non_retryable(err: &AcquisitionError) -> bool {
        matches!(
            err,
            AcquisitionError::Authentication { .. }
                | AcquisitionError::InvalidRequest { .. }
                | AcquisitionError::Config { .. }
        )
    }

    /// Execute a single download attempt: build request, stream response to disk
    async fn execute_download(
        &self,
        job: &DownloadJob,
        cancel_token: &CancellationToken,
    ) -> AcquisitionResult<PathBuf> {
        let url = format!("{}/v0/timeseries.get_range", self.base_url());
        let symbols_joined = job.symbols.join(",");

        info!(
            job_id = %job.job_id,
            dataset = %job.dataset,
            symbols = %symbols_joined,
            schema = %job.schema,
            start = %job.start_date,
            end = %job.end_date,
            "Starting Databento download"
        );

        // Build the request body as form parameters per the Databento REST API
        let form_params = [
            ("dataset", job.dataset.as_str()),
            ("symbols", &symbols_joined),
            ("schema", job.schema.as_str()),
            ("start", job.start_date.as_str()),
            ("end", job.end_date.as_str()),
            ("encoding", "dbn"),
            ("compression", "zstd"),
            ("stype_in", "raw_symbol"),
        ];

        let response = self
            .http_client
            .post(&url)
            .basic_auth(&self.config.api_key, Option::<&str>::None)
            .form(&form_params)
            .send()
            .await?;

        let status = response.status();

        if status == reqwest::StatusCode::UNAUTHORIZED
            || status == reqwest::StatusCode::FORBIDDEN
        {
            return Err(AcquisitionError::Authentication {
                message: format!(
                    "Databento API authentication failed (HTTP {})",
                    status.as_u16()
                ),
            });
        }

        if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
            return Err(AcquisitionError::RateLimit {
                message: "Databento API rate limit exceeded".to_string(),
            });
        }

        if status == reqwest::StatusCode::BAD_REQUEST {
            let body = response.text().await.unwrap_or_default();
            return Err(AcquisitionError::InvalidRequest {
                message: format!("Databento API bad request: {}", body),
            });
        }

        if !status.is_success() {
            let body = response.text().await.unwrap_or_default();
            return Err(AcquisitionError::DatabentorAPI {
                message: format!("Databento API error (HTTP {}): {}", status.as_u16(), body),
            });
        }

        // Determine expected size from Content-Length if available
        let total_bytes = response.content_length();

        // Prepare output directory and file path
        let job_dir = self.config.output_dir.join(job.job_id.to_string());
        tokio::fs::create_dir_all(&job_dir).await.map_err(|e| {
            AcquisitionError::Internal {
                message: format!(
                    "Failed to create output directory {}: {}",
                    job_dir.display(),
                    e
                ),
            }
        })?;

        let filename = format!(
            "{}_{}_{}_{}.dbn.zst",
            job.dataset.replace('.', "_"),
            job.schema,
            job.start_date,
            job.end_date
        );
        let output_path = job_dir.join(&filename);

        debug!(
            job_id = %job.job_id,
            path = %output_path.display(),
            total_bytes = ?total_bytes,
            "Streaming response to file"
        );

        // Stream the response body to disk chunk by chunk
        let mut file = tokio::fs::File::create(&output_path)
            .await
            .map_err(|e| AcquisitionError::Internal {
                message: format!(
                    "Failed to create output file {}: {}",
                    output_path.display(),
                    e
                ),
            })?;

        let mut bytes_downloaded: u64 = 0;
        let mut stream = response.bytes_stream();
        use futures::StreamExt;

        loop {
            tokio::select! {
                chunk = stream.next() => {
                    match chunk {
                        Some(Ok(bytes)) => {
                            file.write_all(&bytes).await.map_err(|e| {
                                AcquisitionError::Internal {
                                    message: format!("Failed to write to {}: {}", output_path.display(), e),
                                }
                            })?;
                            bytes_downloaded += bytes.len() as u64;
                        },
                        Some(Err(e)) => {
                            // Clean up partial file on error
                            let _ = tokio::fs::remove_file(&output_path).await;
                            return Err(AcquisitionError::Network {
                                message: format!("Stream interrupted: {}", e),
                            });
                        },
                        None => break, // Stream finished
                    }
                }
                () = cancel_token.cancelled() => {
                    // Clean up partial file on cancellation
                    let _ = tokio::fs::remove_file(&output_path).await;
                    return Err(AcquisitionError::Internal {
                        message: format!("Download cancelled for job {}", job.job_id),
                    });
                }
            }
        }

        file.flush().await.map_err(|e| AcquisitionError::Internal {
            message: format!("Failed to flush file {}: {}", output_path.display(), e),
        })?;

        // Verify we received data
        if bytes_downloaded == 0 {
            let _ = tokio::fs::remove_file(&output_path).await;
            return Err(AcquisitionError::DataCorruption {
                message: "Received zero bytes from Databento API".to_string(),
            });
        }

        // If we know the total, verify completeness
        if let Some(expected) = total_bytes {
            if bytes_downloaded != expected {
                let _ = tokio::fs::remove_file(&output_path).await;
                return Err(AcquisitionError::DataCorruption {
                    message: format!(
                        "Incomplete download: received {} of {} bytes",
                        bytes_downloaded, expected
                    ),
                });
            }
        }

        info!(
            job_id = %job.job_id,
            bytes = bytes_downloaded,
            path = %output_path.display(),
            "Download completed successfully"
        );

        Ok(output_path)
    }

    /// Estimate cost for a download using the Databento metadata cost endpoint.
    ///
    /// Calls `GET /v0/metadata.get_cost` with the given parameters. If the API
    /// key is not configured, returns 0.0 as a fallback (cost estimation is
    /// non-critical for the download workflow).
    pub async fn estimate_cost(
        &self,
        dataset: &str,
        symbols: &[String],
        start_date: &str,
        end_date: &str,
    ) -> AcquisitionResult<f64> {
        // If no API key, return zero rather than failing the scheduling flow
        if self.config.api_key.is_empty() {
            debug!("No API key configured; returning zero cost estimate");
            return Ok(0.0);
        }

        let url = format!("{}/v0/metadata.get_cost", self.base_url());
        let symbols_joined = symbols.join(",");

        let query_params = [
            ("dataset", dataset),
            ("symbols", &symbols_joined),
            ("start", start_date),
            ("end", end_date),
            ("schema", "trades"),
            ("stype_in", "raw_symbol"),
        ];

        let response = self
            .http_client
            .get(&url)
            .basic_auth(&self.config.api_key, Option::<&str>::None)
            .query(&query_params)
            .send()
            .await
            .map_err(|e| AcquisitionError::Network {
                message: format!("Cost estimation request failed: {}", e),
            })?;

        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            warn!(
                http_status = status.as_u16(),
                "Cost estimation API returned error: {}; defaulting to 0.0",
                body
            );
            return Ok(0.0);
        }

        // Databento returns a JSON number representing cost in USD
        let cost: f64 = response.json().await.map_err(|e| {
            AcquisitionError::DatabentorAPI {
                message: format!("Failed to parse cost response: {}", e),
            }
        })?;

        debug!(cost_usd = cost, "Estimated download cost");
        Ok(cost)
    }

    /// Cancel a download by signalling its cancellation token.
    ///
    /// If the job is not currently active (already completed or never started),
    /// this is a no-op and returns `Ok(())`.
    pub async fn cancel(&self, job_id: Uuid) -> AcquisitionResult<()> {
        let jobs = self.active_jobs.read().await;
        if let Some(token) = jobs.get(&job_id) {
            info!(job_id = %job_id, "Cancelling active download");
            token.cancel();
            Ok(())
        } else {
            debug!(job_id = %job_id, "No active download found for cancellation (may have already completed)");
            Ok(())
        }
    }

    /// Check whether a job is currently active (downloading)
    pub async fn is_active(&self, job_id: &Uuid) -> bool {
        let jobs = self.active_jobs.read().await;
        jobs.contains_key(job_id)
    }

    /// Get the count of currently active downloads
    pub async fn active_download_count(&self) -> usize {
        let jobs = self.active_jobs.read().await;
        jobs.len()
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;

    fn test_config() -> DatabentoDownloaderConfig {
        DatabentoDownloaderConfig {
            api_key: String::new(),
            output_dir: PathBuf::from("/tmp/foxhunt_test_downloads"),
            max_retries: 2,
            retry_delay_ms: 10,
            request_timeout_secs: 5,
            base_url: None,
        }
    }

    fn test_job() -> DownloadJob {
        DownloadJob {
            job_id: Uuid::new_v4(),
            dataset: "GLBX.MDP3".to_string(),
            symbols: vec!["ES.FUT".to_string()],
            start_date: "2025-01-01".to_string(),
            end_date: "2025-01-02".to_string(),
            schema: "ohlcv-1m".to_string(),
            priority: 3,
        }
    }

    #[test]
    fn test_default_config() {
        let config = DatabentoDownloaderConfig::default();
        assert_eq!(config.max_retries, 3);
        assert_eq!(config.retry_delay_ms, 1000);
        assert_eq!(config.request_timeout_secs, 300);
        assert_eq!(config.output_dir, PathBuf::from("data/downloads"));
        assert!(config.base_url.is_none());
    }

    #[test]
    fn test_base_url_override() {
        let mut config = test_config();
        let downloader = DatabentoDownloader::new(config.clone());
        assert_eq!(downloader.base_url(), DATABENTO_HIST_BASE_URL);

        config.base_url = Some("https://custom-host.example.com".to_string());
        let downloader = DatabentoDownloader::new(config);
        assert_eq!(downloader.base_url(), "https://custom-host.example.com");
    }

    #[test]
    fn test_validate_api_key_empty() {
        let config = test_config();
        let downloader = DatabentoDownloader::new(config);
        let result = downloader.validate_api_key();
        assert!(result.is_err());
        let err_msg = result.err().map(|e| e.to_string()).unwrap_or_default();
        assert!(err_msg.contains("DATABENTO_API_KEY"));
    }

    #[test]
    fn test_validate_api_key_present() {
        let mut config = test_config();
        config.api_key = "db-test-key-12345".to_string();
        let downloader = DatabentoDownloader::new(config);
        assert!(downloader.validate_api_key().is_ok());
    }

    #[test]
    fn test_validate_job_valid() {
        let job = test_job();
        assert!(DatabentoDownloader::validate_job(&job).is_ok());
    }

    #[test]
    fn test_validate_job_empty_dataset() {
        let mut job = test_job();
        job.dataset = String::new();
        let result = DatabentoDownloader::validate_job(&job);
        assert!(result.is_err());
        let err_msg = result.err().map(|e| e.to_string()).unwrap_or_default();
        assert!(err_msg.contains("dataset"));
    }

    #[test]
    fn test_validate_job_empty_symbols() {
        let mut job = test_job();
        job.symbols = Vec::new();
        let result = DatabentoDownloader::validate_job(&job);
        assert!(result.is_err());
        let err_msg = result.err().map(|e| e.to_string()).unwrap_or_default();
        assert!(err_msg.contains("symbols"));
    }

    #[test]
    fn test_validate_job_empty_dates() {
        let mut job = test_job();
        job.start_date = String::new();
        let result = DatabentoDownloader::validate_job(&job);
        assert!(result.is_err());
    }

    #[test]
    fn test_validate_job_empty_schema() {
        let mut job = test_job();
        job.schema = String::new();
        let result = DatabentoDownloader::validate_job(&job);
        assert!(result.is_err());
        let err_msg = result.err().map(|e| e.to_string()).unwrap_or_default();
        assert!(err_msg.contains("schema"));
    }

    #[test]
    fn test_is_non_retryable() {
        assert!(DatabentoDownloader::is_non_retryable(
            &AcquisitionError::Authentication {
                message: "bad key".to_string(),
            }
        ));
        assert!(DatabentoDownloader::is_non_retryable(
            &AcquisitionError::InvalidRequest {
                message: "bad params".to_string(),
            }
        ));
        assert!(DatabentoDownloader::is_non_retryable(
            &AcquisitionError::Config {
                message: "bad config".to_string(),
            }
        ));

        // Transient errors should be retryable
        assert!(!DatabentoDownloader::is_non_retryable(
            &AcquisitionError::Network {
                message: "timeout".to_string(),
            }
        ));
        assert!(!DatabentoDownloader::is_non_retryable(
            &AcquisitionError::DatabentorAPI {
                message: "500".to_string(),
            }
        ));
        assert!(!DatabentoDownloader::is_non_retryable(
            &AcquisitionError::RateLimit {
                message: "429".to_string(),
            }
        ));
    }

    #[tokio::test]
    async fn test_download_requires_api_key() {
        let config = test_config();
        let downloader = DatabentoDownloader::new(config);
        let job = test_job();

        let result = downloader.download(&job).await;
        assert!(result.is_err());
        let err_msg = result.err().map(|e| e.to_string()).unwrap_or_default();
        assert!(err_msg.contains("DATABENTO_API_KEY"));
    }

    #[tokio::test]
    async fn test_estimate_cost_without_api_key() {
        let config = test_config();
        let downloader = DatabentoDownloader::new(config);

        // Without an API key, estimate_cost should gracefully return 0.0
        let cost = downloader
            .estimate_cost(
                "GLBX.MDP3",
                &["ES.FUT".to_string()],
                "2025-01-01",
                "2025-01-02",
            )
            .await;
        assert!(cost.is_ok());
        assert!((cost.unwrap_or(f64::NAN) - 0.0).abs() < f64::EPSILON);
    }

    #[tokio::test]
    async fn test_cancel_nonexistent_job() {
        let config = test_config();
        let downloader = DatabentoDownloader::new(config);

        // Cancelling a job that doesn't exist is a no-op
        let result = downloader.cancel(Uuid::new_v4()).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_active_download_tracking() {
        let config = test_config();
        let downloader = DatabentoDownloader::new(config);

        assert_eq!(downloader.active_download_count().await, 0);

        let job_id = Uuid::new_v4();
        assert!(!downloader.is_active(&job_id).await);
    }

    #[tokio::test]
    async fn test_download_with_mock_server() {
        // Start a mock HTTP server that returns a small DBN-like payload
        let mut server = mockito::Server::new_async().await;
        let mock_body: Vec<u8> = vec![0xDB, 0x4E, 0x01, 0x00, 0xFF, 0xFE]; // fake DBN bytes

        let mock = server
            .mock("POST", "/v0/timeseries.get_range")
            .match_header("authorization", mockito::Matcher::Regex("Basic .+".to_string()))
            .with_status(200)
            .with_body(&mock_body)
            .create_async()
            .await;

        let temp_dir = tempfile::tempdir().ok();
        let output_dir = temp_dir
            .as_ref()
            .map(|d| d.path().to_path_buf())
            .unwrap_or_else(|| PathBuf::from("/tmp/foxhunt_test_downloads"));

        let config = DatabentoDownloaderConfig {
            api_key: "db-test-key".to_string(),
            output_dir,
            max_retries: 0,
            retry_delay_ms: 10,
            request_timeout_secs: 5,
            base_url: Some(server.url()),
        };
        let downloader = DatabentoDownloader::new(config);
        let job = test_job();

        let result = downloader.download(&job).await;
        mock.assert_async().await;

        assert!(result.is_ok(), "Expected Ok, got: {:?}", result.err());
        let path = result.unwrap_or_default();
        assert!(path.exists(), "Downloaded file should exist at {:?}", path);

        // Verify the file content matches
        let content = tokio::fs::read(&path).await.unwrap_or_default();
        assert_eq!(content, mock_body);
    }

    #[tokio::test]
    async fn test_download_handles_auth_failure() {
        let mut server = mockito::Server::new_async().await;
        let mock = server
            .mock("POST", "/v0/timeseries.get_range")
            .with_status(401)
            .with_body("Unauthorized")
            .create_async()
            .await;

        let config = DatabentoDownloaderConfig {
            api_key: "bad-key".to_string(),
            output_dir: PathBuf::from("/tmp/foxhunt_test_auth"),
            max_retries: 2,
            retry_delay_ms: 10,
            request_timeout_secs: 5,
            base_url: Some(server.url()),
        };
        let downloader = DatabentoDownloader::new(config);
        let job = test_job();

        let result = downloader.download(&job).await;
        // Should fail and NOT retry (auth is non-retryable)
        mock.assert_async().await;
        assert!(result.is_err());
        let err_msg = result.err().map(|e| e.to_string()).unwrap_or_default();
        assert!(err_msg.contains("authentication") || err_msg.contains("Authentication"));
    }

    #[tokio::test]
    async fn test_download_retries_on_server_error() {
        let mut server = mockito::Server::new_async().await;

        // First two attempts return 500, third succeeds
        let mock_fail = server
            .mock("POST", "/v0/timeseries.get_range")
            .with_status(500)
            .with_body("Internal Server Error")
            .expect(2)
            .create_async()
            .await;

        let mock_success = server
            .mock("POST", "/v0/timeseries.get_range")
            .with_status(200)
            .with_body(vec![0x01, 0x02, 0x03])
            .expect(1)
            .create_async()
            .await;

        let temp_dir = tempfile::tempdir().ok();
        let output_dir = temp_dir
            .as_ref()
            .map(|d| d.path().to_path_buf())
            .unwrap_or_else(|| PathBuf::from("/tmp/foxhunt_test_retry"));

        let config = DatabentoDownloaderConfig {
            api_key: "db-test-key".to_string(),
            output_dir,
            max_retries: 3,
            retry_delay_ms: 10,
            request_timeout_secs: 5,
            base_url: Some(server.url()),
        };
        let downloader = DatabentoDownloader::new(config);
        let job = test_job();

        let result = downloader.download(&job).await;
        mock_fail.assert_async().await;
        mock_success.assert_async().await;

        assert!(result.is_ok(), "Expected success after retries, got: {:?}", result.err());
    }

    #[tokio::test]
    async fn test_cancellation_during_download() {
        let mut server = mockito::Server::new_async().await;

        // Mock a slow response (server returns data but we'll cancel before reading all)
        let _mock = server
            .mock("POST", "/v0/timeseries.get_range")
            .with_status(200)
            .with_body(vec![0xAB; 1024])
            .create_async()
            .await;

        let temp_dir = tempfile::tempdir().ok();
        let output_dir = temp_dir
            .as_ref()
            .map(|d| d.path().to_path_buf())
            .unwrap_or_else(|| PathBuf::from("/tmp/foxhunt_test_cancel"));

        let config = DatabentoDownloaderConfig {
            api_key: "db-test-key".to_string(),
            output_dir,
            max_retries: 0,
            retry_delay_ms: 10,
            request_timeout_secs: 5,
            base_url: Some(server.url()),
        };
        let downloader = Arc::new(DatabentoDownloader::new(config));
        let job = test_job();
        let job_id = job.job_id;
        let dl = downloader.clone();

        // Spawn the download in a background task
        let handle = tokio::spawn(async move {
            dl.download(&job).await
        });

        // Give it a moment to register the token, then cancel
        tokio::time::sleep(Duration::from_millis(50)).await;
        let _ = downloader.cancel(job_id).await;

        // The download should either succeed (if it completed before cancel) or be cancelled
        let result = handle.await;
        assert!(result.is_ok()); // The task itself should not panic
    }

    #[tokio::test]
    async fn test_estimate_cost_with_mock() {
        let mut server = mockito::Server::new_async().await;
        let mock = server
            .mock("GET", "/v0/metadata.get_cost")
            .match_query(mockito::Matcher::AllOf(vec![
                mockito::Matcher::UrlEncoded("dataset".to_string(), "GLBX.MDP3".to_string()),
            ]))
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body("12.50")
            .create_async()
            .await;

        let config = DatabentoDownloaderConfig {
            api_key: "db-test-key".to_string(),
            output_dir: PathBuf::from("/tmp"),
            max_retries: 0,
            retry_delay_ms: 10,
            request_timeout_secs: 5,
            base_url: Some(server.url()),
        };
        let downloader = DatabentoDownloader::new(config);

        let cost = downloader
            .estimate_cost(
                "GLBX.MDP3",
                &["ES.FUT".to_string()],
                "2025-01-01",
                "2025-01-02",
            )
            .await;

        mock.assert_async().await;
        assert!(cost.is_ok());
        assert!((cost.unwrap_or(0.0) - 12.50).abs() < f64::EPSILON);
    }
}
