databuild/databuild/repositories/partitions/mod.rs

use crate::*;
use crate::event_log::{BuildEventLog, BuildEventLogError, Result};
use crate::status_utils::list_response_helpers;
use std::sync::Arc;
use std::collections::HashMap;
use serde::Serialize;

/// Repository for querying partition data from the build event log
pub struct PartitionsRepository {
    event_log: Arc<dyn BuildEventLog>,
}

/// Summary of a partition's current state and history
#[derive(Debug, Clone, Serialize)]
pub struct PartitionInfo {
    pub partition_ref: String,
    pub current_status: PartitionStatus,
    pub last_updated: i64,
    pub builds_count: usize,
    pub last_successful_build: Option<String>,
    pub invalidation_count: usize,
}

/// Detailed partition status with timeline
#[derive(Debug, Clone, Serialize)]
pub struct PartitionStatusEvent {
    pub timestamp: i64,
    pub status: PartitionStatus,
    pub message: String,
    pub build_request_id: String,
    pub job_run_id: Option<String>,
}

impl PartitionsRepository {
    /// Create a new PartitionsRepository
    pub fn new(event_log: Arc<dyn BuildEventLog>) -> Self {
        Self { event_log }
    }

    /// List all partitions with their current status
    /// 
    /// Returns a list of all partitions that have been referenced in the build event log,
    /// along with their current status and summary information.
    pub async fn list(&self, limit: Option<usize>) -> Result<Vec<PartitionInfo>> {
        // Get all partition events from the event log
        let events = self.event_log.get_events_in_range(0, i64::MAX).await?;
        
        let mut partition_data: HashMap<String, Vec<PartitionStatusEvent>> = HashMap::new();
        
        // Collect all partition events
        for event in events {
            if let Some(build_event::EventType::PartitionEvent(p_event)) = &event.event_type {
                if let Some(partition_ref) = &p_event.partition_ref {
                    let status = match p_event.status_code {
                        1 => PartitionStatus::PartitionRequested,
                        2 => PartitionStatus::PartitionAnalyzed,
                        3 => PartitionStatus::PartitionBuilding,
                        4 => PartitionStatus::PartitionAvailable,
                        5 => PartitionStatus::PartitionFailed,
                        6 => PartitionStatus::PartitionDelegated,
                        _ => PartitionStatus::PartitionUnknown,
                    };
                    
                    let status_event = PartitionStatusEvent {
                        timestamp: event.timestamp,
                        status,
                        message: p_event.message.clone(),
                        build_request_id: event.build_request_id.clone(),
                        job_run_id: if p_event.job_run_id.is_empty() { None } else { Some(p_event.job_run_id.clone()) },
                    };
                    
                    partition_data.entry(partition_ref.str.clone())
                        .or_insert_with(Vec::new)
                        .push(status_event);
                }
            }
            
            // Also check for partition invalidation events
            if let Some(build_event::EventType::PartitionInvalidationEvent(pi_event)) = &event.event_type {
                if let Some(partition_ref) = &pi_event.partition_ref {
                    let status_event = PartitionStatusEvent {
                        timestamp: event.timestamp,
                        status: PartitionStatus::PartitionUnknown, // Invalidated
                        message: format!("Invalidated: {}", pi_event.reason),
                        build_request_id: event.build_request_id.clone(),
                        job_run_id: None,
                    };
                    
                    partition_data.entry(partition_ref.str.clone())
                        .or_insert_with(Vec::new)
                        .push(status_event);
                }
            }
        }
        
        // Convert to PartitionInfo structs
        let mut partition_infos: Vec<PartitionInfo> = partition_data.into_iter()
            .map(|(partition_ref, mut events)| {
                // Sort events by timestamp
                events.sort_by_key(|e| e.timestamp);
                
                // Get current status from latest event
                let (current_status, last_updated) = events.last()
                    .map(|e| (e.status.clone(), e.timestamp))
                    .unwrap_or((PartitionStatus::PartitionUnknown, 0));
                
                // Count builds and find last successful build
                let builds: std::collections::HashSet<String> = events.iter()
                    .map(|e| e.build_request_id.clone())
                    .collect();
                
                let last_successful_build = events.iter()
                    .rev()
                    .find(|e| e.status == PartitionStatus::PartitionAvailable)
                    .map(|e| e.build_request_id.clone());
                
                // Count invalidations
                let invalidation_count = events.iter()
                    .filter(|e| e.message.starts_with("Invalidated:"))
                    .count();
                
                PartitionInfo {
                    partition_ref,
                    current_status,
                    last_updated,
                    builds_count: builds.len(),
                    last_successful_build,
                    invalidation_count,
                }
            })
            .collect();
        
        // Sort by most recently updated
        partition_infos.sort_by(|a, b| b.last_updated.cmp(&a.last_updated));
        
        // Apply limit if specified
        if let Some(limit) = limit {
            partition_infos.truncate(limit);
        }
        
        Ok(partition_infos)
    }

    /// Show detailed information about a specific partition
    /// 
    /// Returns the complete timeline of status changes for the specified partition,
    /// including all builds that have referenced it.
    pub async fn show(&self, partition_ref: &str) -> Result<Option<(PartitionInfo, Vec<PartitionStatusEvent>)>> {
        // Get all events for this partition
        let events = self.event_log.get_partition_events(partition_ref, None).await?;
        
        if events.is_empty() {
            return Ok(None);
        }
        
        let mut status_events = Vec::new();
        let mut builds = std::collections::HashSet::new();
        
        // Process partition events
        for event in &events {
            if let Some(build_event::EventType::PartitionEvent(p_event)) = &event.event_type {
                let status = match p_event.status_code {
                    1 => PartitionStatus::PartitionRequested,
                    2 => PartitionStatus::PartitionAnalyzed,
                    3 => PartitionStatus::PartitionBuilding,
                    4 => PartitionStatus::PartitionAvailable,
                    5 => PartitionStatus::PartitionFailed,
                    6 => PartitionStatus::PartitionDelegated,
                    _ => PartitionStatus::PartitionUnknown,
                };
                
                status_events.push(PartitionStatusEvent {
                    timestamp: event.timestamp,
                    status,
                    message: p_event.message.clone(),
                    build_request_id: event.build_request_id.clone(),
                    job_run_id: if p_event.job_run_id.is_empty() { None } else { Some(p_event.job_run_id.clone()) },
                });
                
                builds.insert(event.build_request_id.clone());
            }
        }
        
        // Also check for invalidation events in all events
        let all_events = self.event_log.get_events_in_range(0, i64::MAX).await?;
        let mut invalidation_count = 0;
        
        for event in all_events {
            if let Some(build_event::EventType::PartitionInvalidationEvent(pi_event)) = &event.event_type {
                if let Some(partition) = &pi_event.partition_ref {
                    if partition.str == partition_ref {
                        status_events.push(PartitionStatusEvent {
                            timestamp: event.timestamp,
                            status: PartitionStatus::PartitionUnknown, // Invalidated
                            message: format!("Invalidated: {}", pi_event.reason),
                            build_request_id: event.build_request_id.clone(),
                            job_run_id: None,
                        });
                        invalidation_count += 1;
                    }
                }
            }
        }
        
        // Sort events by timestamp
        status_events.sort_by_key(|e| e.timestamp);
        
        // Get current status from latest event
        let (current_status, last_updated) = status_events.last()
            .map(|e| (e.status.clone(), e.timestamp))
            .unwrap_or((PartitionStatus::PartitionUnknown, 0));
        
        // Find last successful build
        let last_successful_build = status_events.iter()
            .rev()
            .find(|e| e.status == PartitionStatus::PartitionAvailable)
            .map(|e| e.build_request_id.clone());
        
        let partition_info = PartitionInfo {
            partition_ref: partition_ref.to_string(),
            current_status,
            last_updated,
            builds_count: builds.len(),
            last_successful_build,
            invalidation_count,
        };
        
        Ok(Some((partition_info, status_events)))
    }

    /// Invalidate a partition with a reason
    /// 
    /// This method uses the EventWriter to write a partition invalidation event.
    /// It validates that the partition exists before invalidating it.
    pub async fn invalidate(&self, partition_ref: &str, reason: String, build_request_id: String) -> Result<()> {
        // First check if the partition exists
        let partition_exists = self.show(partition_ref).await?.is_some();
        
        if !partition_exists {
            return Err(BuildEventLogError::QueryError(
                format!("Cannot invalidate non-existent partition: {}", partition_ref)
            ));
        }
        
        // Use EventWriter to write the invalidation event
        let event_writer = crate::event_log::writer::EventWriter::new(self.event_log.clone());
        let partition = PartitionRef { str: partition_ref.to_string() };
        
        event_writer.invalidate_partition(build_request_id, partition, reason).await
    }

    /// List partitions returning protobuf response format with dual status fields
    /// 
    /// This method provides the unified CLI/Service response format with both
    /// status codes (enum values) and status names (human-readable strings).
    pub async fn list_protobuf(&self, request: PartitionsListRequest) -> Result<PartitionsListResponse> {
        // Get legacy format data
        let partition_infos = self.list(request.limit.map(|l| l as usize)).await?;
        
        // Convert to protobuf format with dual status fields
        let partitions: Vec<PartitionSummary> = partition_infos.into_iter()
            .map(|info| {
                list_response_helpers::create_partition_summary(
                    info.partition_ref,
                    info.current_status,
                    info.last_updated,
                    info.builds_count,
                    info.invalidation_count,
                    info.last_successful_build,
                )
            })
            .collect();

        // TODO: Implement proper pagination with offset and has_more
        // For now, return simple response without full pagination support
        let total_count = partitions.len() as u32;
        let has_more = false; // This would be calculated based on actual total vs returned

        Ok(PartitionsListResponse {
            partitions,
            total_count,
            has_more,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::event_log::mock::{MockBuildEventLog, test_events};

    #[tokio::test]
    async fn test_partitions_repository_list_empty() {
        let mock_log = Arc::new(MockBuildEventLog::new().await.unwrap());
        let repo = PartitionsRepository::new(mock_log);
        
        let partitions = repo.list(None).await.unwrap();
        assert!(partitions.is_empty());
    }

    #[tokio::test]
    async fn test_partitions_repository_list_with_data() {
        let build_id = "test-build-123".to_string();
        let partition1 = PartitionRef { str: "data/users".to_string() };
        let partition2 = PartitionRef { str: "data/orders".to_string() };
        
        // Create events for multiple partitions
        let events = vec![
            test_events::build_request_received(Some(build_id.clone()), vec![partition1.clone(), partition2.clone()]),
            test_events::partition_status(Some(build_id.clone()), partition1.clone(), PartitionStatus::PartitionBuilding, None),
            test_events::partition_status(Some(build_id.clone()), partition1.clone(), PartitionStatus::PartitionAvailable, None),
            test_events::partition_status(Some(build_id.clone()), partition2.clone(), PartitionStatus::PartitionBuilding, None),
            test_events::partition_status(Some(build_id.clone()), partition2.clone(), PartitionStatus::PartitionFailed, None),
        ];
        
        let mock_log = Arc::new(MockBuildEventLog::with_events(events).await.unwrap());
        let repo = PartitionsRepository::new(mock_log);
        
        let partitions = repo.list(None).await.unwrap();
        assert_eq!(partitions.len(), 2);
        
        // Find partitions by name
        let users_partition = partitions.iter().find(|p| p.partition_ref == "data/users").unwrap();
        let orders_partition = partitions.iter().find(|p| p.partition_ref == "data/orders").unwrap();
        
        assert_eq!(users_partition.current_status, PartitionStatus::PartitionAvailable);
        assert_eq!(orders_partition.current_status, PartitionStatus::PartitionFailed);
        assert_eq!(users_partition.builds_count, 1);
        assert_eq!(orders_partition.builds_count, 1);
    }

    #[tokio::test]
    async fn test_partitions_repository_show() {
        let build_id = "test-build-456".to_string();
        let partition = PartitionRef { str: "analytics/metrics".to_string() };
        
        let events = vec![
            test_events::partition_status(Some(build_id.clone()), partition.clone(), PartitionStatus::PartitionRequested, None),
            test_events::partition_status(Some(build_id.clone()), partition.clone(), PartitionStatus::PartitionBuilding, None),
            test_events::partition_status(Some(build_id.clone()), partition.clone(), PartitionStatus::PartitionAvailable, None),
        ];
        
        let mock_log = Arc::new(MockBuildEventLog::with_events(events).await.unwrap());
        let repo = PartitionsRepository::new(mock_log);
        
        let result = repo.show(&partition.str).await.unwrap();
        assert!(result.is_some());
        
        let (info, timeline) = result.unwrap();
        assert_eq!(info.partition_ref, "analytics/metrics");
        assert_eq!(info.current_status, PartitionStatus::PartitionAvailable);
        assert_eq!(info.builds_count, 1);
        assert_eq!(timeline.len(), 3);
        
        // Verify timeline order
        assert_eq!(timeline[0].status, PartitionStatus::PartitionRequested);
        assert_eq!(timeline[1].status, PartitionStatus::PartitionBuilding);
        assert_eq!(timeline[2].status, PartitionStatus::PartitionAvailable);
    }

    #[tokio::test]
    async fn test_partitions_repository_show_nonexistent() {
        let mock_log = Arc::new(MockBuildEventLog::new().await.unwrap());
        let repo = PartitionsRepository::new(mock_log);
        
        let result = repo.show("nonexistent/partition").await.unwrap();
        assert!(result.is_none());
    }

    #[tokio::test]
    async fn test_partitions_repository_invalidate() {
        let build_id = "test-build-789".to_string();
        let partition = PartitionRef { str: "temp/data".to_string() };
        
        // Start with an existing partition
        let events = vec![
            test_events::partition_status(Some(build_id.clone()), partition.clone(), PartitionStatus::PartitionAvailable, None),
        ];
        
        let mock_log = Arc::new(MockBuildEventLog::with_events(events).await.unwrap());
        let repo = PartitionsRepository::new(mock_log.clone());
        
        // Invalidate the partition
        repo.invalidate(&partition.str, "Test invalidation".to_string(), build_id.clone()).await.unwrap();
        
        // Verify the invalidation was recorded
        // Note: This test demonstrates the pattern, but the MockBuildEventLog would need
        // to be enhanced to properly store invalidation events for full verification
        
        // Try to invalidate a non-existent partition
        let result = repo.invalidate("nonexistent/partition", "Should fail".to_string(), build_id).await;
        assert!(result.is_err());
    }
}