Fundamental_Analysis/services/yfinance-provider-service/src/state.rs

use std::sync::Arc;

use dashmap::DashMap;
use uuid::Uuid;
use common_contracts::observability::{TaskProgress, ObservabilityTaskStatus};
use common_contracts::workflow_harness::TaskState;

use crate::config::AppConfig;
use crate::yfinance::YFinanceDataProvider;

#[derive(Clone)]
pub struct AppState {
    pub tasks: Arc<DashMap<Uuid, TaskProgress>>,
    pub config: Arc<AppConfig>,
    pub yfinance_provider: Arc<YFinanceDataProvider>,
}

impl AppState {
    pub fn new(config: AppConfig) -> Self {
        let provider = Arc::new(YFinanceDataProvider::new(config.yfinance_enabled));

        Self {
            tasks: Arc::new(DashMap::new()),
            config: Arc::new(config),
            yfinance_provider: provider,
        }
    }
}

// Implement TaskState trait for AppState
#[async_trait::async_trait]
impl TaskState for AppState {
    fn update_status(&self, task_id: Uuid, status: ObservabilityTaskStatus, progress: u8, details: String) {
        if let Some(mut task) = self.tasks.get_mut(&task_id) {
            task.status = status;
            task.progress_percent = progress;
            task.details = details;
        }
    }

    fn fail_task(&self, task_id: Uuid, error: String) {
        if let Some(mut task) = self.tasks.get_mut(&task_id) {
            task.status = ObservabilityTaskStatus::Failed;
            task.details = format!("Failed: {}", error);
        }
    }

    fn complete_task(&self, task_id: Uuid, details: String) {
        if let Some(mut task) = self.tasks.get_mut(&task_id) {
            task.status = ObservabilityTaskStatus::Completed;
            task.progress_percent = 100;
            task.details = details;
        }
    }

    fn get_nats_addr(&self) -> String {
        self.config.nats_addr.clone()
    }

    fn get_persistence_url(&self) -> String {
        self.config.data_persistence_service_url.clone()
    }
}