//! Rada Working Memory Manager.
//!
//! Manages active working memory blocks by prioritizing items relevant to active goals
//! and decaying less active blocks over time to prevent context overflow.

use mimirs_core::{BlockId, DensityMemory, Goal, MemoryBlock};
use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};

/// Rada Working Memory Manager.
pub struct Rada {
    alpha: f32,
    beta: f32,
    /// Timestamp for when each block was last accessed/updated.
    last_accessed: HashMap<BlockId, u64>,
}

impl Rada {
    /// Creates a new Rada Working Memory Manager.
    pub fn new(alpha: f32, beta: f32) -> Self {
        Self {
            alpha,
            beta,
            last_accessed: HashMap::new(),
        }
    }

    /// Record an access timestamp for a block to track elapsed time.
    pub fn touch(&mut self, id: BlockId) {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        self.last_accessed.insert(id, now);
    }

    /// Calculates the relevance score for a given block relative to active goals and elapsed time.
    pub fn calculate_relevance(
        &self,
        block_id: BlockId,
        block_rho: &DensityMemory,
        active_goals: &[Goal],
        goal_rhos: &HashMap<String, DensityMemory>,
    ) -> f32 {
        let mut max_overlap = 0.0f32;

        // Calculate overlap with the highest-matching active goal
        for goal in active_goals {
            if let Some(goal_rho) = goal_rhos.get(&goal.description) {
                let overlap = block_rho.overlap(goal_rho);
                if overlap > max_overlap {
                    max_overlap = overlap;
                }
            }
        }

        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        let last_time = self.last_accessed.get(&block_id).cloned().unwrap_or(now);
        let elapsed = (now - last_time) as f32;

        // Relevance = alpha * overlap - beta * elapsed_time
        self.alpha * max_overlap - self.beta * elapsed
    }

    /// Sorts blocks by relevance, keeping the most relevant ones under the token/character limit.
    /// Returns a tuple of (active_blocks, inactive_blocks).
    pub fn swap_context(
        &mut self,
        blocks: &mut [MemoryBlock],
        block_rhos: &HashMap<BlockId, DensityMemory>,
        active_goals: &[Goal],
        goal_rhos: &HashMap<String, DensityMemory>,
        total_limit: usize,
    ) -> (Vec<MemoryBlock>, Vec<MemoryBlock>) {
        let mut scored_blocks: Vec<(&MemoryBlock, f32)> = blocks
            .iter()
            .map(|b| {
                if b.pinned {
                    (b, f32::MAX)
                } else if let Some(rho) = block_rhos.get(&b.id) {
                    let score = self.calculate_relevance(b.id, rho, active_goals, goal_rhos);
                    (b, score)
                } else {
                    (b, f32::MIN)
                }
            })
            .collect();

        // Sort descending by relevance
        scored_blocks.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        let mut active = Vec::new();
        let mut inactive = Vec::new();

        let mut current_size = 0;
        for (block, _score) in scored_blocks {
            if current_size + block.value.len() <= total_limit || block.pinned {
                active.push(block.clone());
                current_size += block.value.len();
            } else {
                inactive.push(block.clone());
            }
        }

        (active, inactive)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use mimirs_core::{BlockId, Goal, GoalId, MemoryBlock};

    #[test]
    fn test_rada_swap_context() {
        let mut rada = Rada::new(10.0, 0.1);

        let id_a = BlockId::new();
        let id_b = BlockId::new();

        let mut block_a = MemoryBlock::default();
        block_a.id = id_a;
        block_a.value = "Short block".to_string();

        let mut block_b = MemoryBlock::default();
        block_b.id = id_b;
        block_b.value = "Long block content that exceeds limits".to_string();

        let mut blocks = vec![block_a.clone(), block_b.clone()];

        // Create mock densities (dimension 384) with non-zero norms
        let v_a = vec![1.0f32; 384];
        let mut v_b = vec![0.0f32; 384];
        v_b[0] = 1.0f32; // Make it non-zero so from_pure succeeds
        let goal_v = vec![1.0f32; 384];

        let rho_a = DensityMemory::from_pure(&v_a).unwrap();
        let rho_b = DensityMemory::from_pure(&v_b).unwrap();
        let rho_goal = DensityMemory::from_pure(&goal_v).unwrap();

        let mut block_rhos = HashMap::new();
        block_rhos.insert(id_a, rho_a);
        block_rhos.insert(id_b, rho_b);

        let active_goals = vec![Goal {
            id: GoalId::new(),
            description: "Task target".to_string(),
            priority: 1.0,
            persistence: 1.0,
            completion_score: 0.0,
            dependencies: Vec::new(),
        }];

        let mut goal_rhos = HashMap::new();
        goal_rhos.insert("Task target".to_string(), rho_goal);

        // Limit is 15 chars, so only block_a fits
        let (active, inactive) =
            rada.swap_context(&mut blocks, &block_rhos, &active_goals, &goal_rhos, 15);

        assert_eq!(active.len(), 1);
        assert_eq!(active[0].id, id_a);
        assert_eq!(inactive.len(), 1);
        assert_eq!(inactive[0].id, id_b);
    }
}
