foxhunt/services/data_acquisition_service/src/service.rs

//! gRPC service implementation for Data Acquisition Service
//!
//! Implements the DataAcquisitionService proto interface with:
//! - Download job scheduling and management
//! - Integration with Databento downloader
//! - Data validation pipeline
//! - Automatic MinIO upload
//! - Job persistence and status tracking

use crate::downloader::{DatabentoDownloader, DatabentoDownloaderConfig};
use crate::proto::data_acquisition_service_server::DataAcquisitionService;
use crate::proto::*;
use crate::uploader::{MinIOUploader, MinIOUploaderConfig};
use crate::validator::{DataValidator, ValidationConfig};
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use futures::Stream;
use tokio::sync::RwLock;
use tonic::{Request, Response, Status};
use uuid::Uuid;
use tracing::instrument;

/// In-memory job store backed by a `HashMap<String, DownloadJobDetails>`.
///
/// Acceptable for the current MVP: the data-acquisition service manages a small
/// number of concurrent download jobs (typically <100) that are short-lived.
/// Jobs are created, run, and either complete or fail within minutes. In-memory
/// storage keeps the service dependency-free (no PostgreSQL / SQLite required)
/// and simplifies deployment. If durability across restarts becomes a requirement
/// (e.g., multi-hour downloads or job auditing), migrate to the shared PostgreSQL
/// instance used by other services.
type JobStore = Arc<RwLock<std::collections::HashMap<String, DownloadJobDetails>>>;

/// Data Acquisition Service implementation
pub struct DataAcquisitionServiceImpl {
    downloader: Arc<DatabentoDownloader>,
    _uploader: Arc<MinIOUploader>,
    _validator: Arc<DataValidator>,
    job_store: JobStore,
}

impl DataAcquisitionServiceImpl {
    /// Create new service instance
    pub fn new() -> Self {
        Self {
            downloader: Arc::new(DatabentoDownloader::new(
                DatabentoDownloaderConfig::default(),
            )),
            _uploader: Arc::new(MinIOUploader::new(MinIOUploaderConfig::default())),
            _validator: Arc::new(DataValidator::new(ValidationConfig::default())),
            job_store: Arc::new(RwLock::new(std::collections::HashMap::new())),
        }
    }

    /// Create service with custom configuration
    pub fn with_config(
        downloader_config: DatabentoDownloaderConfig,
        uploader_config: MinIOUploaderConfig,
        validator_config: ValidationConfig,
    ) -> Self {
        Self {
            downloader: Arc::new(DatabentoDownloader::new(downloader_config)),
            _uploader: Arc::new(MinIOUploader::new(uploader_config)),
            _validator: Arc::new(DataValidator::new(validator_config)),
            job_store: Arc::new(RwLock::new(std::collections::HashMap::new())),
        }
    }
}

impl Default for DataAcquisitionServiceImpl {
    fn default() -> Self {
        Self::new()
    }
}

#[tonic::async_trait]
impl DataAcquisitionService for DataAcquisitionServiceImpl {
    #[instrument(skip_all)]
    async fn schedule_download(
        &self,
        request: Request<ScheduleDownloadRequest>,
    ) -> Result<Response<ScheduleDownloadResponse>, Status> {
        let req = request.into_inner();
        let job_id = Uuid::new_v4();

        // Estimate cost
        let estimated_cost = self
            .downloader
            .estimate_cost(&req.dataset, &req.symbols, &req.start_date, &req.end_date)
            .await
            .map_err(|e| Status::internal(format!("Cost estimation failed: {}", e)))?;

        // Create job details
        let job_details = DownloadJobDetails {
            job_id: job_id.to_string(),
            status: DownloadStatus::Pending as i32,
            dataset: req.dataset.clone(),
            symbols: req.symbols.clone(),
            start_date: req.start_date.clone(),
            end_date: req.end_date.clone(),
            schema: req.schema.clone(),
            description: req.description.clone(),
            tags: req.tags.clone(),
            priority: req.priority,
            progress_percentage: 0.0,
            bytes_downloaded: 0,
            total_bytes: 0,
            estimated_cost_usd: estimated_cost,
            actual_cost_usd: 0.0,
            records_count: 0,
            invalid_records: 0,
            data_quality_score: 0.0,
            local_path: String::new(),
            minio_path: String::new(),
            created_at: chrono::Utc::now().timestamp(),
            started_at: 0,
            completed_at: 0,
            error_message: String::new(),
            retry_count: 0,
            created_by: "system".to_string(),
        };

        // Store job
        {
            let mut store = self.job_store.write().await;
            store.insert(job_id.to_string(), job_details.clone());
        }

        // Spawn background download task
        {
            let job_store = Arc::clone(&self.job_store);
            let downloader = Arc::clone(&self.downloader);
            let job_id_str = job_id.to_string();
            let download_job = crate::downloader::DownloadJob {
                job_id,
                dataset: req.dataset.clone(),
                symbols: req.symbols.clone(),
                start_date: req.start_date.clone(),
                end_date: req.end_date.clone(),
                schema: req.schema.clone(),
                priority: req.priority as i32,
            };

            tokio::spawn(async move {
                // Mark as downloading
                {
                    let mut store = job_store.write().await;
                    if let Some(job) = store.get_mut(&job_id_str) {
                        job.status = DownloadStatus::Downloading as i32;
                        job.started_at = chrono::Utc::now().timestamp();
                    }
                }

                match downloader.download(&download_job).await {
                    Ok(path) => {
                        let mut store = job_store.write().await;
                        if let Some(job) = store.get_mut(&job_id_str) {
                            job.status = DownloadStatus::Completed as i32;
                            job.progress_percentage = 100.0;
                            job.completed_at = chrono::Utc::now().timestamp();
                            job.local_path = path.to_string_lossy().to_string();
                        }
                    }
                    Err(e) => {
                        let mut store = job_store.write().await;
                        if let Some(job) = store.get_mut(&job_id_str) {
                            job.status = DownloadStatus::Failed as i32;
                            job.completed_at = chrono::Utc::now().timestamp();
                            job.error_message = e.to_string();
                            job.retry_count += 1;
                        }
                    }
                }
            });
        }

        Ok(Response::new(ScheduleDownloadResponse {
            job_id: job_id.to_string(),
            status: DownloadStatus::Pending as i32,
            estimated_cost_usd: estimated_cost,
            message: format!("Download scheduled for {} symbols", req.symbols.len()),
        }))
    }

    #[instrument(skip_all)]
    async fn get_download_status(
        &self,
        request: Request<GetDownloadStatusRequest>,
    ) -> Result<Response<GetDownloadStatusResponse>, Status> {
        let req = request.into_inner();

        let store = self.job_store.read().await;
        let job_details = store
            .get(&req.job_id)
            .ok_or_else(|| Status::not_found(format!("Job not found: {}", req.job_id)))?
            .clone();

        Ok(Response::new(GetDownloadStatusResponse {
            job_details: Some(job_details),
        }))
    }

    #[instrument(skip_all)]
    async fn cancel_download(
        &self,
        request: Request<CancelDownloadRequest>,
    ) -> Result<Response<CancelDownloadResponse>, Status> {
        let req = request.into_inner();

        let job_id = Uuid::parse_str(&req.job_id)
            .map_err(|e| Status::invalid_argument(format!("Invalid job ID: {}", e)))?;

        self.downloader
            .cancel(job_id)
            .await
            .map_err(|e| Status::internal(format!("Cancel failed: {}", e)))?;

        // Update job status
        {
            let mut store = self.job_store.write().await;
            if let Some(job) = store.get_mut(&req.job_id) {
                job.status = DownloadStatus::Cancelled as i32;
                job.error_message = req.reason.clone();
            }
        }

        Ok(Response::new(CancelDownloadResponse {
            success: true,
            message: "Download cancelled successfully".to_string(),
        }))
    }

    #[instrument(skip_all)]
    async fn list_download_jobs(
        &self,
        request: Request<ListDownloadJobsRequest>,
    ) -> Result<Response<ListDownloadJobsResponse>, Status> {
        let req = request.into_inner();

        let store = self.job_store.read().await;
        let mut jobs: Vec<_> = store.values().cloned().collect();

        // Apply status filter
        if req.status_filter != 0 {
            jobs.retain(|j| j.status == req.status_filter);
        }

        // Apply time filter
        if req.start_time > 0 {
            jobs.retain(|j| j.created_at >= req.start_time);
        }
        if req.end_time > 0 {
            jobs.retain(|j| j.created_at <= req.end_time);
        }

        // Pagination
        let page = req.page.max(1) as usize;
        let page_size = req.page_size.clamp(1, 100) as usize;
        let start_idx = (page - 1) * page_size;
        let _end_idx = start_idx + page_size;

        let total_count = jobs.len() as u32;
        let paginated_jobs = jobs.into_iter().skip(start_idx).take(page_size);

        let job_summaries: Vec<DownloadJobSummary> = paginated_jobs
            .map(|job| DownloadJobSummary {
                job_id: job.job_id,
                status: job.status,
                dataset: job.dataset,
                symbols: job.symbols,
                start_date: job.start_date,
                end_date: job.end_date,
                progress_percentage: job.progress_percentage,
                actual_cost_usd: job.actual_cost_usd,
                created_at: job.created_at,
                completed_at: job.completed_at,
            })
            .collect();

        Ok(Response::new(ListDownloadJobsResponse {
            jobs: job_summaries,
            total_count,
            page: req.page,
            page_size: req.page_size,
        }))
    }

    #[instrument(skip_all)]
    async fn health_check(
        &self,
        _request: Request<HealthCheckRequest>,
    ) -> Result<Response<HealthCheckResponse>, Status> {
        let mut details = std::collections::HashMap::new();
        details.insert("service".to_string(), "data_acquisition".to_string());
        details.insert("version".to_string(), env!("CARGO_PKG_VERSION").to_string());

        // Check job store accessibility
        let store = self.job_store.read().await;
        let total_jobs = store.len();
        let active_jobs = store
            .values()
            .filter(|j| j.status == DownloadStatus::Downloading as i32)
            .count();
        let failed_jobs = store
            .values()
            .filter(|j| j.status == DownloadStatus::Failed as i32)
            .count();
        drop(store);

        // Check downloader active count
        let active_downloads = self.downloader.active_download_count().await;

        details.insert("total_jobs".to_string(), total_jobs.to_string());
        details.insert("active_jobs".to_string(), active_jobs.to_string());
        details.insert("failed_jobs".to_string(), failed_jobs.to_string());
        details.insert(
            "active_downloads".to_string(),
            active_downloads.to_string(),
        );
        details.insert("status".to_string(), "operational".to_string());

        Ok(Response::new(HealthCheckResponse {
            healthy: true,
            message: format!(
                "Service is healthy: {} active downloads, {} total jobs",
                active_downloads, total_jobs
            ),
            details,
        }))
    }

    type StreamDownloadStatusStream =
        Pin<Box<dyn Stream<Item = Result<ListDownloadJobsResponse, Status>> + Send>>;

    async fn stream_download_status(
        &self,
        request: Request<StreamDownloadStatusRequest>,
    ) -> Result<Response<Self::StreamDownloadStatusStream>, Status> {
        let req = request.into_inner();
        let interval_secs = if req.interval_seconds == 0 { 5 } else { req.interval_seconds };
        let store = self.job_store.clone();

        let stream = async_stream::stream! {
            let mut interval = tokio::time::interval(Duration::from_secs(u64::from(interval_secs)));
            loop {
                interval.tick().await;
                let jobs: Vec<DownloadJobSummary> = store.read().await.values().map(|job| DownloadJobSummary {
                    job_id: job.job_id.clone(),
                    status: job.status,
                    dataset: job.dataset.clone(),
                    symbols: job.symbols.clone(),
                    start_date: job.start_date.clone(),
                    end_date: job.end_date.clone(),
                    progress_percentage: job.progress_percentage,
                    actual_cost_usd: job.actual_cost_usd,
                    created_at: job.created_at,
                    completed_at: job.completed_at,
                }).collect();
                #[allow(clippy::cast_possible_truncation)]
                let total = jobs.len() as u32;
                yield Ok(ListDownloadJobsResponse {
                    jobs,
                    total_count: total,
                    page: 1,
                    page_size: total,
                });
            }
        };

        Ok(Response::new(Box::pin(stream)))
    }
}