use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::collections::{BTreeMap, HashMap};

#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct HNSWIndex {
    dim: usize,
    vectors: HashMap<u64, Vec<f32>>,
}

impl HNSWIndex {
    pub fn new(dim: usize, _m: usize, _ef_construction: usize) -> Self {
        Self {
            dim,
            vectors: HashMap::new(),
        }
    }

    pub fn insert(&mut self, id: u64, vector: &[f32]) {
        if vector.len() != self.dim {
            return;
        }
        self.vectors.insert(id, vector.to_vec());
    }

    pub fn search(&self, query: &[f32], k: usize, _ef: usize) -> Vec<(u64, f32)> {
        if query.len() != self.dim || k == 0 {
            return Vec::new();
        }

        let mut scores: Vec<(u64, f32)> = self
            .vectors
            .iter()
            .map(|(id, vector)| (*id, cosine_similarity(query, vector)))
            .collect();

        scores.sort_by(|left, right| {
            right
                .1
                .partial_cmp(&left.1)
                .unwrap_or(Ordering::Equal)
        });
        scores.truncate(k);
        scores
    }

    pub fn len(&self) -> usize {
        self.vectors.len()
    }
}

#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct AgentIndex {
    agent_id: [u8; 16],
    episodic: HNSWIndex,
    semantic: HNSWIndex,
    procedural: Option<HNSWIndex>,
}

#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct AgentIndexManager {
    indexes: HashMap<[u8; 16], AgentIndex>,
}

impl AgentIndexManager {
    pub fn new() -> Self {
        Self {
            indexes: HashMap::new(),
        }
    }

    pub fn get_or_create(&mut self, agent_id: [u8; 16], dim: usize) -> &mut AgentIndex {
        self.indexes.entry(agent_id).or_insert_with(|| AgentIndex {
            agent_id,
            episodic: HNSWIndex::new(dim, 16, 200),
            semantic: HNSWIndex::new(dim, 16, 200),
            procedural: None,
        })
    }

    pub fn insert_episodic(&mut self, agent_id: [u8; 16], id: u64, vector: &[f32]) {
        let index = self.get_or_create(agent_id, vector.len());
        index.episodic.insert(id, vector);
    }

    pub fn insert_semantic(&mut self, agent_id: [u8; 16], id: u64, vector: &[f32]) {
        let index = self.get_or_create(agent_id, vector.len());
        index.semantic.insert(id, vector);
    }

    pub fn search_episodic(
        &self,
        agent_id: &[u8; 16],
        query: &[f32],
        k: usize,
        ef: usize,
    ) -> Vec<(u64, f32)> {
        self.indexes
            .get(agent_id)
            .map(|index| index.episodic.search(query, k, ef))
            .unwrap_or_default()
    }

    pub fn search_semantic(
        &self,
        agent_id: &[u8; 16],
        query: &[f32],
        k: usize,
        ef: usize,
    ) -> Vec<(u64, f32)> {
        self.indexes
            .get(agent_id)
            .map(|index| index.semantic.search(query, k, ef))
            .unwrap_or_default()
    }

    pub fn agent_count(&self) -> usize {
        self.indexes.len()
    }

    pub fn total_vectors(&self) -> usize {
        self.indexes
            .values()
            .map(|index| {
                index.episodic.len()
                    + index.semantic.len()
                    + index.procedural.as_ref().map_or(0, |p| p.len())
            })
            .sum()
    }
}

#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct AnomalyDetector {
    centroid: Vec<f32>,
    sample_count: usize,
    running_sum: Vec<f32>,
}

impl AnomalyDetector {
    pub fn new(dim: usize) -> Self {
        Self {
            centroid: vec![0.0; dim],
            sample_count: 0,
            running_sum: vec![0.0; dim],
        }
    }

    pub fn update(&mut self, embedding: &[f32]) {
        if embedding.len() != self.centroid.len() {
            return;
        }
        self.sample_count += 1;
        for (index, value) in embedding.iter().enumerate() {
            self.running_sum[index] += value;
            self.centroid[index] = self.running_sum[index] / self.sample_count as f32;
        }
    }

    pub fn is_anomalous(&self, embedding: &[f32], threshold: f32) -> (bool, f32) {
        if embedding.len() != self.centroid.len() {
            return (false, 0.0);
        }
        if self.sample_count < 50 {
            return (false, 0.0);
        }
        let distance = cosine_distance(embedding, &self.centroid);
        (distance > threshold, distance)
    }

    pub fn centroid(&self) -> &[f32] {
        &self.centroid
    }

    pub fn sample_count(&self) -> usize {
        self.sample_count
    }
}

#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct TemporalIndex {
    entries: BTreeMap<u64, Vec<u64>>,
}

impl TemporalIndex {
    pub fn new() -> Self {
        Self {
            entries: BTreeMap::new(),
        }
    }

    pub fn insert(&mut self, timestamp_ms: u64, memory_id: u64) {
        self.entries
            .entry(timestamp_ms)
            .or_insert_with(Vec::new)
            .push(memory_id);
    }

    pub fn range_query(&self, start_ms: u64, end_ms: u64) -> Vec<u64> {
        self.entries
            .range(start_ms..end_ms)
            .flat_map(|(_, ids)| ids.iter().copied())
            .collect()
    }

    pub fn count_in_range(&self, start_ms: u64, end_ms: u64) -> usize {
        self.entries
            .range(start_ms..end_ms)
            .map(|(_, ids)| ids.len())
            .sum()
    }

    pub fn nearest_before(&self, timestamp_ms: u64) -> Option<(u64, Vec<u64>)> {
        self.entries
            .range(..timestamp_ms)
            .next_back()
            .map(|(timestamp, ids)| (*timestamp, ids.clone()))
    }

    pub fn len(&self) -> usize {
        self.entries.values().map(|ids| ids.len()).sum()
    }
}

fn cosine_similarity(left: &[f32], right: &[f32]) -> f32 {
    let left_norm = left.iter().map(|v| v * v).sum::<f32>().sqrt();
    let right_norm = right.iter().map(|v| v * v).sum::<f32>().sqrt();
    if left_norm <= 1e-12 || right_norm <= 1e-12 {
        return 0.0;
    }
    let dot = left
        .iter()
        .zip(right.iter())
        .map(|(l, r)| l * r)
        .sum::<f32>();
    (dot / (left_norm * right_norm)).clamp(-1.0, 1.0)
}

fn cosine_distance(left: &[f32], right: &[f32]) -> f32 {
    1.0 - cosine_similarity(left, right)
}

#[cfg(feature = "python-ffi")]
mod pyffi {
    use super::{AgentIndexManager, AnomalyDetector, TemporalIndex};
    use pyo3::prelude::*;

    fn parse_agent_id(agent_id: Vec<u8>) -> [u8; 16] {
        let mut aid = [0u8; 16];
        for (index, byte) in agent_id.iter().take(16).enumerate() {
            aid[index] = *byte;
        }
        aid
    }

    #[pyclass]
    pub struct RustAgentIndex {
        inner: AgentIndexManager,
    }

    #[pymethods]
    impl RustAgentIndex {
        #[new]
        fn new() -> Self {
            Self {
                inner: AgentIndexManager::new(),
            }
        }

        fn insert_episodic(&mut self, agent_id: Vec<u8>, id: u64, vector: Vec<f32>) {
            self.inner
                .insert_episodic(parse_agent_id(agent_id), id, &vector);
        }

        fn insert_semantic(&mut self, agent_id: Vec<u8>, id: u64, vector: Vec<f32>) {
            self.inner
                .insert_semantic(parse_agent_id(agent_id), id, &vector);
        }

        fn search_episodic(
            &self,
            agent_id: Vec<u8>,
            query: Vec<f32>,
            k: usize,
            ef: usize,
        ) -> Vec<(u64, f32)> {
            let aid = parse_agent_id(agent_id);
            self.inner.search_episodic(&aid, &query, k, ef)
        }

        fn search_semantic(
            &self,
            agent_id: Vec<u8>,
            query: Vec<f32>,
            k: usize,
            ef: usize,
        ) -> Vec<(u64, f32)> {
            let aid = parse_agent_id(agent_id);
            self.inner.search_semantic(&aid, &query, k, ef)
        }

        fn agent_count(&self) -> usize {
            self.inner.agent_count()
        }

        fn total_vectors(&self) -> usize {
            self.inner.total_vectors()
        }
    }

    #[pyclass]
    pub struct RustAnomalyDetector {
        inner: AnomalyDetector,
    }

    #[pymethods]
    impl RustAnomalyDetector {
        #[new]
        fn new(dim: usize) -> Self {
            Self {
                inner: AnomalyDetector::new(dim),
            }
        }

        fn update(&mut self, embedding: Vec<f32>) {
            self.inner.update(&embedding);
        }

        fn is_anomalous(&self, embedding: Vec<f32>, threshold: f32) -> (bool, f32) {
            self.inner.is_anomalous(&embedding, threshold)
        }

        fn sample_count(&self) -> usize {
            self.inner.sample_count()
        }
    }

    #[pyclass]
    pub struct RustTemporalIndex {
        inner: TemporalIndex,
    }

    #[pymethods]
    impl RustTemporalIndex {
        #[new]
        fn new() -> Self {
            Self {
                inner: TemporalIndex::new(),
            }
        }

        fn insert(&mut self, timestamp_ms: u64, memory_id: u64) {
            self.inner.insert(timestamp_ms, memory_id);
        }

        fn range_query(&self, start_ms: u64, end_ms: u64) -> Vec<u64> {
            self.inner.range_query(start_ms, end_ms)
        }

        fn count_in_range(&self, start_ms: u64, end_ms: u64) -> usize {
            self.inner.count_in_range(start_ms, end_ms)
        }

        fn len(&self) -> usize {
            self.inner.len()
        }
    }

    #[pymodule]
    fn chronos_core_v07(_py: Python, module: &Bound<PyModule>) -> PyResult<()> {
        module.add_class::<RustAgentIndex>()?;
        module.add_class::<RustAnomalyDetector>()?;
        module.add_class::<RustTemporalIndex>()?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::{AgentIndexManager, AnomalyDetector, TemporalIndex};

    #[test]
    fn agent_index_manager_insert_and_search() {
        let mut manager = AgentIndexManager::new();
        let agent = [1u8; 16];
        manager.insert_episodic(agent, 1, &[1.0, 0.0, 0.0]);
        manager.insert_semantic(agent, 2, &[0.0, 1.0, 0.0]);

        let episodic = manager.search_episodic(&agent, &[1.0, 0.0, 0.0], 5, 50);
        let semantic = manager.search_semantic(&agent, &[0.0, 1.0, 0.0], 5, 50);

        assert_eq!(manager.agent_count(), 1);
        assert_eq!(manager.total_vectors(), 2);
        assert_eq!(episodic.first().map(|item| item.0), Some(1));
        assert_eq!(semantic.first().map(|item| item.0), Some(2));
    }

    #[test]
    fn anomaly_detector_thresholds() {
        let mut detector = AnomalyDetector::new(3);
        for _ in 0..60 {
            detector.update(&[1.0, 1.0, 1.0]);
        }
        let (flag, distance) = detector.is_anomalous(&[1.0, 1.0, 1.0], 0.4);
        assert!(!flag);
        assert!(distance <= 0.01);
    }

    #[test]
    fn temporal_index_range_and_nearest() {
        let mut index = TemporalIndex::new();
        index.insert(100, 1);
        index.insert(200, 2);
        index.insert(300, 3);

        assert_eq!(index.range_query(100, 301), vec![1, 2, 3]);
        assert_eq!(index.count_in_range(100, 250), 2);
        assert_eq!(index.nearest_before(220).map(|x| x.0), Some(200));
    }
}
