// Copyright 2025 The Drasi Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![allow(clippy::unwrap_used)]
use std::{
    collections::{BTreeMap, HashMap, HashSet},
    hash::{Hash, Hasher},
    ops::Bound,
    sync::Arc,
};

use async_stream::stream;
use async_trait::async_trait;
use hashers::jenkins::spooky_hash::SpookyHasher;
use tokio::sync::RwLock;

use crate::{
    interface::{ElementArchiveIndex, ElementIndex, ElementResult, ElementStream, IndexError},
    models::{
        Element, ElementMetadata, ElementPropertyMap, ElementReference, ElementTimestamp,
        ElementValue, QueryJoin, QueryJoinKey, TimestampBound, TimestampRange,
    },
    path_solver::match_path::MatchPath,
};

#[allow(clippy::type_complexity)]
pub struct InMemoryElementIndex {
    elements: Arc<RwLock<HashMap<ElementReference, Arc<Element>>>>,

    slot_affinity: Arc<RwLock<HashMap<ElementReference, HashSet<usize>>>>,

    element_by_slot: Arc<RwLock<HashSet<(usize, ElementReference)>>>,
    element_by_slot_in: Arc<RwLock<HashMap<(usize, ElementReference), HashSet<ElementReference>>>>,
    element_by_slot_out: Arc<RwLock<HashMap<(usize, ElementReference), HashSet<ElementReference>>>>,

    element_archive: Arc<RwLock<HashMap<ElementReference, ElementArchive>>>,
    archive_enabled: bool,

    join_spec_by_label: Arc<RwLock<HashMap<Arc<str>, Vec<(Arc<QueryJoin>, Vec<usize>)>>>>,

    // [(join_label, field_value)] => [QueryJoinKey] => ElementReference[]
    partial_joins:
        Arc<RwLock<HashMap<(String, u64), HashMap<QueryJoinKey, HashSet<ElementReference>>>>>,
}

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

impl InMemoryElementIndex {
    pub fn new() -> Self {
        Self {
            elements: Arc::new(RwLock::new(HashMap::new())),
            slot_affinity: Arc::new(RwLock::new(HashMap::new())),
            element_by_slot: Arc::new(RwLock::new(HashSet::new())),
            element_by_slot_in: Arc::new(RwLock::new(HashMap::new())),
            element_by_slot_out: Arc::new(RwLock::new(HashMap::new())),
            element_archive: Arc::new(RwLock::new(HashMap::new())),
            join_spec_by_label: Arc::new(RwLock::new(HashMap::new())),
            partial_joins: Arc::new(RwLock::new(HashMap::new())),
            archive_enabled: false,
        }
    }

    pub fn enable_archive(&mut self) {
        self.archive_enabled = true;
    }

    async fn clear_slot_affinity(&self, element: &Element) -> Result<(), IndexError> {
        let mut af_guard = self.slot_affinity.write().await;
        // Remove the entry rather than draining it in place, otherwise an empty
        // HashSet is retained for every element that is ever deleted, leaking
        // memory proportional to total element churn. Callers that update a live
        // element (`set_element`) re-insert affinity via `set_slot_affinity`.
        let affinity = match af_guard.remove(element.get_reference()) {
            Some(affinity) => affinity,
            // No recorded affinity means there is nothing to clear.
            None => return Ok(()),
        };

        match element {
            Element::Node {
                metadata,
                properties: _,
            } => {
                let mut guard = self.element_by_slot.write().await;

                for old_slot in affinity {
                    guard.remove(&(old_slot, metadata.reference.clone()));
                }
            }
            Element::Relation {
                metadata,
                in_node,
                out_node,
                properties: _,
            } => {
                let mut in_guard = self.element_by_slot_in.write().await;
                let mut out_guard = self.element_by_slot_out.write().await;

                for old_slot in affinity {
                    let in_key = (old_slot, in_node.clone());
                    if let Some(set) = in_guard.get_mut(&in_key) {
                        set.remove(&metadata.reference);
                        if set.is_empty() {
                            in_guard.remove(&in_key);
                        }
                    }

                    let out_key = (old_slot, out_node.clone());
                    if let Some(set) = out_guard.get_mut(&out_key) {
                        set.remove(&metadata.reference);
                        if set.is_empty() {
                            out_guard.remove(&out_key);
                        }
                    }
                }
            }
        };

        Ok(())
    }

    async fn set_slot_affinity(
        &self,
        element: &Element,
        slots: &Vec<usize>,
    ) -> Result<(), IndexError> {
        let mut af_guard = self.slot_affinity.write().await;
        let affinity = af_guard.entry(element.get_reference().clone()).or_default();

        match element {
            Element::Node {
                metadata,
                properties: _,
            } => {
                let mut guard = self.element_by_slot.write().await;

                for slot in slots {
                    guard.insert((*slot, metadata.reference.clone()));
                    affinity.insert(*slot);
                }
            }
            Element::Relation {
                metadata,
                in_node,
                out_node,
                properties: _,
            } => {
                let mut in_guard = self.element_by_slot_in.write().await;
                let mut out_guard = self.element_by_slot_out.write().await;

                for slot in slots {
                    affinity.insert(*slot);

                    in_guard
                        .entry((*slot, in_node.clone()))
                        .or_default()
                        .insert(metadata.reference.clone());

                    out_guard
                        .entry((*slot, out_node.clone()))
                        .or_default()
                        .insert(metadata.reference.clone());
                }
            }
        };

        Ok(())
    }

    async fn update_source_joins(
        &self,
        old_element: Option<Arc<Element>>,
        new_element: Arc<Element>,
    ) -> Result<(), IndexError> {
        match new_element.as_ref() {
            Element::Node {
                metadata,
                properties,
            } => {
                let join_spec_by_label = self.join_spec_by_label.read().await;
                for (label, joins) in join_spec_by_label.iter() {
                    if !metadata.labels.contains(label) {
                        continue;
                    }

                    for (qj, slots) in joins {
                        for qjk in &qj.keys {
                            if qjk.label.as_str() != label.as_ref() {
                                continue;
                            }

                            match properties.get(&qjk.property) {
                                Some(p) => {
                                    if let Some(old_element) = &old_element {
                                        if let Element::Node {
                                            metadata: _old_metadata,
                                            properties: old_properties,
                                        } = old_element.as_ref()
                                        {
                                            if let Some(old_p) = old_properties.get(&qjk.property) {
                                                if old_p == p {
                                                    continue;
                                                }
                                            }
                                        }
                                    }

                                    let value_hash = get_value_hash(p);
                                    let mut partial_joins_guard = self.partial_joins.write().await;
                                    let partial_joins = partial_joins_guard
                                        .entry((qj.id.clone(), value_hash))
                                        .or_default();

                                    let did_insert = partial_joins
                                        .entry(qjk.clone())
                                        .or_default()
                                        .insert(metadata.reference.clone());

                                    let mut elements_to_set = Vec::new();
                                    let mut values_to_delete = Vec::new();

                                    if did_insert {
                                        //remove old partial joins
                                        if let Some(old_element) = &old_element {
                                            if let Element::Node {
                                                metadata: _old_metadata,
                                                properties: old_properties,
                                            } = old_element.as_ref()
                                            {
                                                if let Some(old_p) =
                                                    old_properties.get(&qjk.property)
                                                {
                                                    values_to_delete.push(old_p);
                                                }
                                            }
                                        }

                                        //find matching counterparts
                                        for qjk2 in &qj.keys {
                                            if qjk == qjk2 {
                                                continue;
                                            }

                                            if let Some(others) = partial_joins.get(qjk2) {
                                                for other in others {
                                                    let in_out = Element::Relation {
                                                        metadata: ElementMetadata {
                                                            reference: get_join_virtual_ref(
                                                                new_element.get_reference(),
                                                                other,
                                                            ),
                                                            labels: Arc::from([Arc::from(
                                                                qj.id.clone(),
                                                            )]),
                                                            effective_from: new_element
                                                                .get_effective_from(),
                                                        },
                                                        in_node: new_element
                                                            .get_reference()
                                                            .clone(),
                                                        out_node: other.clone(),
                                                        properties: ElementPropertyMap::new(),
                                                    };

                                                    let out_in = Element::Relation {
                                                        metadata: ElementMetadata {
                                                            reference: get_join_virtual_ref(
                                                                other,
                                                                new_element.get_reference(),
                                                            ),
                                                            labels: Arc::from([Arc::from(
                                                                qj.id.clone(),
                                                            )]),
                                                            effective_from: new_element
                                                                .get_effective_from(),
                                                        },
                                                        in_node: other.clone(),
                                                        out_node: new_element
                                                            .get_reference()
                                                            .clone(),
                                                        properties: ElementPropertyMap::new(),
                                                    };

                                                    elements_to_set.push((in_out, slots.clone()));
                                                    elements_to_set.push((out_in, slots.clone()));
                                                }
                                            }
                                        }
                                    }

                                    drop(partial_joins_guard);

                                    if let Some(old_element) = &old_element {
                                        for val_to_delete in values_to_delete {
                                            self.delete_source_join(
                                                old_element.get_reference(),
                                                qj,
                                                qjk,
                                                val_to_delete,
                                            )
                                            .await?;
                                        }
                                    }

                                    for (element, slots) in elements_to_set {
                                        self.set_element(&element, &slots).await?;
                                    }
                                }
                                None => continue,
                            }
                        }
                    }
                }
            }
            _ => return Ok(()),
        }

        Ok(())
    }

    async fn delete_source_joins(&self, old_element: Arc<Element>) -> Result<(), IndexError> {
        match old_element.as_ref() {
            Element::Node {
                metadata,
                properties,
            } => {
                let join_spec_by_label = self.join_spec_by_label.read().await;
                for (label, joins) in join_spec_by_label.iter() {
                    if !metadata.labels.contains(label) {
                        continue;
                    }

                    for (qj, _slots) in joins {
                        for qjk in &qj.keys {
                            if qjk.label.as_str() != label.as_ref() {
                                continue;
                            }

                            match properties.get(&qjk.property) {
                                Some(p) => {
                                    self.delete_source_join(
                                        old_element.get_reference(),
                                        qj,
                                        qjk,
                                        p,
                                    )
                                    .await?;
                                }
                                None => continue,
                            }
                        }
                    }
                }
            }
            _ => return Ok(()),
        }

        Ok(())
    }

    async fn delete_source_join(
        &self,
        old_element: &ElementReference,
        query_join: &QueryJoin,
        join_key: &QueryJoinKey,
        value: &ElementValue,
    ) -> Result<(), IndexError> {
        let mut elements_to_delete = Vec::new();
        let value_hash = get_value_hash(value);
        let outer_key = (query_join.id.clone(), value_hash);
        let mut partial_joins_guard = self.partial_joins.write().await;

        // Look the entry up rather than inserting it. Using `entry(..).or_default()`
        // on the delete path would create an empty outer container that is never
        // reclaimed, causing unbounded growth for high-cardinality/high-churn joins.
        let Some(partial_joins) = partial_joins_guard.get_mut(&outer_key) else {
            log::warn!(
                "delete_source_join: no partial joins for query_join {} (value_hash {}) - \
                 element {:?} may not have been fully registered",
                query_join.id,
                value_hash,
                old_element
            );
            return Ok(());
        };

        let did_remove = match partial_joins.get_mut(join_key) {
            Some(element_set) => {
                let removed = element_set.remove(old_element);
                // Drop the now-empty set so emptied containers don't accumulate.
                if element_set.is_empty() {
                    partial_joins.remove(join_key);
                }
                removed
            }
            None => {
                log::warn!(
                    "delete_source_join: join_key {:?} not found for element {:?} in query_join {} - \
                     element may not have been fully registered",
                    join_key,
                    old_element,
                    query_join.id
                );
                false
            }
        };

        if did_remove {
            for qjk2 in &query_join.keys {
                if join_key == qjk2 {
                    continue;
                }

                if let Some(others) = partial_joins.get(qjk2) {
                    for other in others {
                        let in_out = get_join_virtual_ref(old_element, other);
                        let out_in = get_join_virtual_ref(other, old_element);

                        elements_to_delete.push(in_out);
                        elements_to_delete.push(out_in);
                    }
                }
            }
        }

        // Drop the outer (join_id, value_hash) entry once its inner map is empty.
        let outer_now_empty = partial_joins.is_empty();
        if outer_now_empty {
            partial_joins_guard.remove(&outer_key);
        }

        drop(partial_joins_guard);
        for element_to_delete in elements_to_delete {
            self.delete_element(&element_to_delete).await?;
        }

        Ok(())
    }
}

#[async_trait]
impl ElementIndex for InMemoryElementIndex {
    async fn get_slot_element_by_ref(
        &self,
        slot: usize,
        element_ref: &ElementReference,
    ) -> Result<Option<Arc<Element>>, IndexError> {
        let guard = self.element_by_slot.read().await;
        if guard.contains(&(slot, element_ref.clone())) {
            let guard = self.elements.read().await;
            match guard.get(element_ref) {
                None => Ok(None),
                Some(element) => Ok(Some(element.clone())),
            }
        } else {
            Ok(None)
        }
    }

    async fn get_slot_elements_by_inbound(
        &self,
        slot: usize,
        inbound_ref: &ElementReference,
    ) -> Result<ElementStream, IndexError> {
        let guard = self.element_by_slot_in.read().await;
        match guard.get(&(slot, inbound_ref.clone())) {
            None => Ok(Box::pin(tokio_stream::empty::<ElementResult>())),
            Some(element_refs) => {
                let mut result = Vec::new();
                let guard = self.elements.read().await;
                for element_ref in element_refs {
                    match guard.get(element_ref) {
                        None => (),
                        Some(element) => result.push(element.clone()),
                    }
                }

                Ok(element_vec_to_stream(result))
            }
        }
    }

    async fn get_slot_elements_by_outbound(
        &self,
        slot: usize,
        outbound_ref: &ElementReference,
    ) -> Result<ElementStream, IndexError> {
        let guard = self.element_by_slot_out.read().await;
        match guard.get(&(slot, outbound_ref.clone())) {
            None => Ok(Box::pin(tokio_stream::empty::<ElementResult>())),
            Some(element_refs) => {
                let mut result = Vec::new();
                let guard = self.elements.read().await;
                for element_ref in element_refs {
                    match guard.get(element_ref) {
                        None => (),
                        Some(element) => result.push(element.clone()),
                    }
                }
                Ok(element_vec_to_stream(result))
            }
        }
    }

    async fn get_element(
        &self,
        element_ref: &ElementReference,
    ) -> Result<Option<Arc<Element>>, IndexError> {
        let guard = self.elements.read().await;
        match guard.get(element_ref) {
            None => Ok(None),
            Some(element) => Ok(Some(element.clone())),
        }
    }

    async fn set_element(
        &self,
        element: &Element,
        slot_affinity: &Vec<usize>,
    ) -> Result<(), IndexError> {
        let mut guard = self.elements.write().await;

        let new_element = Arc::new(element.clone());
        let old_element = guard.insert(element.get_reference().clone(), new_element.clone());

        drop(guard);

        if let Some(prev) = &old_element {
            self.clear_slot_affinity(prev).await?;
        }
        self.set_slot_affinity(new_element.as_ref(), slot_affinity)
            .await?;
        self.update_source_joins(old_element, new_element.clone())
            .await?;

        if self.archive_enabled {
            let mut guard = self.element_archive.write().await;
            let archive = guard
                .entry(new_element.get_reference().clone())
                .or_insert_with(ElementArchive::new);

            archive.insert(new_element.clone());
        }

        Ok(())
    }

    async fn delete_element(&self, element_ref: &ElementReference) -> Result<(), IndexError> {
        let mut guard = self.elements.write().await;
        let old_element = guard.remove(element_ref);

        drop(guard);
        if let Some(old_element) = old_element {
            self.clear_slot_affinity(&old_element).await?;
            self.delete_source_joins(old_element).await?;
        }

        Ok(())
    }

    async fn clear(&self) -> Result<(), IndexError> {
        let mut guard = self.elements.write().await;
        guard.clear();

        let mut guard = self.slot_affinity.write().await;
        guard.clear();

        let mut guard = self.element_by_slot.write().await;
        guard.clear();

        let mut guard = self.element_by_slot_in.write().await;
        guard.clear();

        let mut guard = self.element_by_slot_out.write().await;
        guard.clear();

        let mut guard = self.partial_joins.write().await;
        guard.clear();

        Ok(())
    }

    async fn set_joins(&self, match_path: &MatchPath, joins: &Vec<Arc<QueryJoin>>) {
        let joins_by_label = extract_join_spec_by_label(match_path, joins);
        let mut join_spec_by_label = self.join_spec_by_label.write().await;
        join_spec_by_label.clone_from(&joins_by_label);
    }
}

#[async_trait]
impl ElementArchiveIndex for InMemoryElementIndex {
    async fn get_element_as_at(
        &self,
        element_ref: &ElementReference,
        time: ElementTimestamp,
    ) -> Result<Option<Arc<Element>>, IndexError> {
        if !self.archive_enabled {
            return Err(IndexError::ArchiveNotEnabled);
        }

        let guard = self.element_archive.read().await;
        match guard.get(element_ref) {
            None => Ok(None),
            Some(archive) => match archive.get(time) {
                None => Ok(None),
                Some(element) => Ok(Some(element)),
            },
        }
    }

    async fn get_element_versions(
        &self,
        element_ref: &ElementReference,
        range: TimestampRange<ElementTimestamp>,
    ) -> Result<ElementStream, IndexError> {
        if !self.archive_enabled {
            return Err(IndexError::ArchiveNotEnabled);
        }

        let from = range.from;
        let to = range.to;
        let guard = self.element_archive.read().await;
        match guard.get(element_ref) {
            None => Ok(Box::pin(tokio_stream::empty::<ElementResult>())),
            Some(archive) => match from {
                TimestampBound::Included(from) => archive.get_range(from, to),
                TimestampBound::StartFromPrevious(from) => {
                    let initial_start = archive.retrieve_previous_start_timestamp(from).unwrap();
                    archive.get_range(initial_start, to)
                }
            },
        }
    }

    async fn clear(&self) -> Result<(), IndexError> {
        if !self.archive_enabled {
            return Err(IndexError::NotSupported);
        }

        let mut guard = self.element_archive.write().await;
        guard.clear();
        Ok(())
    }
}

struct ElementArchive {
    data: BTreeMap<ElementTimestamp, Arc<Element>>,
}

impl ElementArchive {
    fn new() -> Self {
        ElementArchive {
            data: BTreeMap::new(),
        }
    }

    fn insert(&mut self, element: Arc<Element>) {
        self.data.insert(element.get_effective_from(), element);
    }

    fn retrieve_previous_start_timestamp(
        &self,
        time: ElementTimestamp,
    ) -> Option<ElementTimestamp> {
        match self
            .data
            .range((Bound::Included(&0), Bound::Included(&time)))
            .map(|x| *x.0)
            .next_back()
        {
            None => Some(0),
            Some(v) => Some(v),
        }
    }

    fn get(&self, time: ElementTimestamp) -> Option<Arc<Element>> {
        let mut cur = self.data.range((Bound::Unbounded, Bound::Included(&time)));
        if let Some((_, element)) = cur.next_back() {
            Some(element.clone())
        } else {
            None
        }
    }

    fn get_range(
        &self,
        from: ElementTimestamp,
        to: ElementTimestamp,
    ) -> Result<ElementStream, IndexError> {
        let data: Vec<Arc<Element>> = self
            .data
            .range((Bound::Included(&from), Bound::Included(&to)))
            .map(|x| x.1.clone())
            .collect();

        Ok(element_vec_to_stream(data))
    }
}

fn extract_join_spec_by_label(
    match_path: &MatchPath,
    joins: &Vec<Arc<QueryJoin>>,
) -> HashMap<Arc<str>, Vec<(Arc<QueryJoin>, Vec<usize>)>> {
    let mut result: HashMap<Arc<str>, Vec<(Arc<QueryJoin>, Vec<usize>)>> = HashMap::new();

    for join in joins {
        let mut slots = Vec::new();
        for (slot_num, slot) in match_path.slots.iter().enumerate() {
            if slot.spec.labels.contains(&join.id.as_str().into()) {
                slots.push(slot_num);
            }
        }
        if slots.is_empty() {
            continue;
        }

        for jk in &join.keys {
            result
                .entry(Arc::from(jk.label.clone()))
                .or_default()
                .push((join.clone(), slots.clone()));
        }
    }

    result
}

fn get_join_virtual_ref(ref1: &ElementReference, ref2: &ElementReference) -> ElementReference {
    let new_id = format!("{}:{}", ref1.element_id, ref2.element_id);
    ElementReference::new("&join", new_id.as_str())
}

fn element_vec_to_stream(elements: Vec<Arc<Element>>) -> ElementStream {
    Box::pin(stream! {
        for e in elements {
            yield Ok(e);
        }
    })
}

fn get_value_hash(value: &ElementValue) -> u64 {
    let mut hasher = SpookyHasher::default();
    value.hash(&mut hasher);
    hasher.finish()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::interface::ElementIndex;
    use crate::path_solver::match_path::{MatchPath, MatchPathSlot, SlotElementSpec};
    use std::collections::HashSet;

    /// Helper function to create a simple MatchPath with one slot containing a label
    fn create_match_path_with_label(label: &str) -> MatchPath {
        MatchPath {
            slots: vec![MatchPathSlot {
                spec: SlotElementSpec {
                    annotation: None,
                    labels: vec![Arc::from(label)],
                    predicates: vec![],
                },
                in_slots: vec![],
                out_slots: vec![],
                paths: HashSet::from([0]),
                optional: false,
            }],
            optional_paths: HashSet::new(),
        }
    }

    /// Helper function to create a test node element
    fn create_test_node(source_id: &str, element_id: &str, label: &str) -> Element {
        Element::Node {
            metadata: ElementMetadata {
                reference: ElementReference::new(source_id, element_id),
                labels: Arc::from([Arc::from(label)]),
                effective_from: 0,
            },
            properties: ElementPropertyMap::new(),
        }
    }

    /// Helper function to create a test relation element
    fn create_test_relation(
        source_id: &str,
        element_id: &str,
        label: &str,
        in_id: &str,
        out_id: &str,
    ) -> Element {
        Element::Relation {
            metadata: ElementMetadata {
                reference: ElementReference::new(source_id, element_id),
                labels: Arc::from([Arc::from(label)]),
                effective_from: 0,
            },
            in_node: ElementReference::new(source_id, in_id),
            out_node: ElementReference::new(source_id, out_id),
            properties: ElementPropertyMap::new(),
        }
    }

    /// Helper function to create a test node with a property
    fn create_test_node_with_property(
        source_id: &str,
        element_id: &str,
        label: &str,
        prop_name: &str,
        prop_value: i64,
    ) -> Element {
        let mut properties = ElementPropertyMap::new();
        properties.insert(prop_name, ElementValue::Integer(prop_value));
        Element::Node {
            metadata: ElementMetadata {
                reference: ElementReference::new(source_id, element_id),
                labels: Arc::from([Arc::from(label)]),
                effective_from: 0,
            },
            properties,
        }
    }

    #[tokio::test]
    async fn test_basic_element_operations() {
        let index = InMemoryElementIndex::new();

        // Insert an element
        let element = create_test_node("source1", "node1", "Person");
        index.set_element(&element, &vec![0]).await.unwrap();

        // Retrieve the element
        let retrieved = index
            .get_element(&ElementReference::new("source1", "node1"))
            .await
            .unwrap();
        assert!(retrieved.is_some());

        // Delete the element
        index
            .delete_element(&ElementReference::new("source1", "node1"))
            .await
            .unwrap();

        // Verify it's gone
        let retrieved = index
            .get_element(&ElementReference::new("source1", "node1"))
            .await
            .unwrap();
        assert!(retrieved.is_none());
    }

    #[tokio::test]
    async fn test_delete_element_not_in_index() {
        let index = InMemoryElementIndex::new();

        // Deleting an element that doesn't exist should succeed
        let result = index
            .delete_element(&ElementReference::new("source1", "nonexistent"))
            .await;
        assert!(result.is_ok());
    }

    /// Test that deleting an element works correctly when joins are configured
    /// but the element was inserted BEFORE joins were set up.
    ///
    /// This tests the fix for the panic in delete_source_join when
    /// partial_joins doesn't contain the expected join key.
    ///
    /// Scenario:
    /// 1. Insert element with join-related label and property
    /// 2. Configure joins AFTER the element exists (simulating hot config change)
    /// 3. Delete the element - this should NOT panic
    #[tokio::test]
    async fn test_delete_element_with_joins_configured_after_insert() {
        let index = InMemoryElementIndex::new();

        // Step 1: Insert an element with a property that matches a join key
        // At this point, NO joins are configured, so partial_joins remains empty
        let element =
            create_test_node_with_property("source1", "person1", "Person", "company_id", 42);
        index.set_element(&element, &vec![0]).await.unwrap();

        // Step 2: Configure joins AFTER the element exists
        // The join_spec_by_label will be populated, but partial_joins for this
        // element's join key will NOT be (since element was inserted before joins existed)
        let join = Arc::new(QueryJoin {
            id: "company_join".to_string(),
            keys: vec![QueryJoinKey {
                label: "Person".to_string(),
                property: "company_id".to_string(),
            }],
        });

        // Create a MatchPath where the slot has the join ID as a label
        // (This is how extract_join_spec_by_label determines which slots use a join)
        let match_path = create_match_path_with_label("company_join");

        index.set_joins(&match_path, &vec![join]).await;

        // Step 3: Delete the element
        // BEFORE THE FIX: This would panic at line 374 with:
        //   "called `Option::unwrap()` on a `None` value"
        // AFTER THE FIX: This should complete successfully
        let result = index
            .delete_element(&ElementReference::new("source1", "person1"))
            .await;

        assert!(
            result.is_ok(),
            "delete_element should not panic when partial_joins is missing the join key"
        );

        // Verify the element is actually deleted
        let retrieved = index
            .get_element(&ElementReference::new("source1", "person1"))
            .await
            .unwrap();
        assert!(retrieved.is_none());
    }

    /// Test that normal join operations still work correctly after the fix.
    /// This ensures the defensive check doesn't break the happy path.
    #[tokio::test]
    async fn test_join_operations_work_correctly() {
        let index = InMemoryElementIndex::new();

        // Configure joins FIRST (normal path)
        let join = Arc::new(QueryJoin {
            id: "company_join".to_string(),
            keys: vec![
                QueryJoinKey {
                    label: "Person".to_string(),
                    property: "company_id".to_string(),
                },
                QueryJoinKey {
                    label: "Company".to_string(),
                    property: "id".to_string(),
                },
            ],
        });

        let match_path = create_match_path_with_label("company_join");
        index.set_joins(&match_path, &vec![join]).await;

        // Insert a Person with company_id
        let person =
            create_test_node_with_property("source1", "person1", "Person", "company_id", 42);
        index.set_element(&person, &vec![0]).await.unwrap();

        // Insert a Company with matching id
        let company = create_test_node_with_property("source1", "company1", "Company", "id", 42);
        index.set_element(&company, &vec![0]).await.unwrap();

        // At this point, virtual join relations should have been created
        // Delete the person - this should clean up the join relations
        let result = index
            .delete_element(&ElementReference::new("source1", "person1"))
            .await;

        assert!(result.is_ok(), "Normal join deletion should work correctly");
    }

    /// Test updating an element's join property value.
    /// When a property value changes, the old join key entry should be removed
    /// and a new one created.
    #[tokio::test]
    async fn test_update_element_join_property() {
        let index = InMemoryElementIndex::new();

        // Configure joins
        let join = Arc::new(QueryJoin {
            id: "company_join".to_string(),
            keys: vec![QueryJoinKey {
                label: "Person".to_string(),
                property: "company_id".to_string(),
            }],
        });

        let match_path = create_match_path_with_label("company_join");
        index.set_joins(&match_path, &vec![join]).await;

        // Insert element with initial property value
        let person1 =
            create_test_node_with_property("source1", "person1", "Person", "company_id", 42);
        index.set_element(&person1, &vec![0]).await.unwrap();

        // Update element with new property value
        let person1_updated =
            create_test_node_with_property("source1", "person1", "Person", "company_id", 99);
        let result = index.set_element(&person1_updated, &vec![0]).await;

        assert!(
            result.is_ok(),
            "Updating join property should work correctly"
        );

        // Delete should still work
        let result = index
            .delete_element(&ElementReference::new("source1", "person1"))
            .await;

        assert!(result.is_ok(), "Deletion after update should work");
    }

    /// Test clearing the index clears all join-related state
    #[tokio::test]
    async fn test_clear_clears_partial_joins() {
        let index = InMemoryElementIndex::new();

        // Configure joins
        let join = Arc::new(QueryJoin {
            id: "company_join".to_string(),
            keys: vec![QueryJoinKey {
                label: "Person".to_string(),
                property: "company_id".to_string(),
            }],
        });

        let match_path = create_match_path_with_label("company_join");
        index.set_joins(&match_path, &vec![join]).await;

        // Insert some elements
        let person =
            create_test_node_with_property("source1", "person1", "Person", "company_id", 42);
        index.set_element(&person, &vec![0]).await.unwrap();

        // Clear the index (use fully qualified path to disambiguate from ElementArchiveIndex::clear)
        ElementIndex::clear(&index).await.unwrap();

        // Verify element is gone
        let retrieved = index
            .get_element(&ElementReference::new("source1", "person1"))
            .await
            .unwrap();
        assert!(retrieved.is_none());
    }

    /// Test that delete_source_join does not panic when join_key was never registered.
    #[tokio::test]
    async fn test_delete_source_join_missing_key_no_panic() {
        let index = InMemoryElementIndex::new();

        // Create element reference that was never inserted
        let element_ref = ElementReference::new("test_source", "element_123");

        // Create a join key that was never registered
        let query_join = QueryJoin {
            id: "test_join".to_string(),
            keys: vec![QueryJoinKey {
                label: "TestLabel".to_string(),
                property: "id".to_string(),
            }],
        };

        let join_key = &query_join.keys[0];
        let value = ElementValue::String("test_value".into());

        // This should NOT panic - previously it would panic on unwrap()
        let result = index
            .delete_source_join(&element_ref, &query_join, join_key, &value)
            .await;

        // Should complete successfully (no-op) without panic
        assert!(result.is_ok());

        // The delete path must NOT create an outer entry for a missing key,
        // otherwise emptied/absent containers leak and grow unbounded (issue #641).
        {
            let guard = index.partial_joins.read().await;
            assert!(
                guard.is_empty(),
                "delete_source_join on a missing key must not create partial_joins entries"
            );
        }
    }

    /// Test that delete_source_join correctly removes an element that was previously registered.
    #[tokio::test]
    async fn test_delete_source_join_registered_key() {
        let index = InMemoryElementIndex::new();

        let element_ref = ElementReference::new("test_source", "element_456");
        let value = ElementValue::String("shared_key".into());
        let value_hash = get_value_hash(&value);

        let query_join = QueryJoin {
            id: "test_join".to_string(),
            keys: vec![QueryJoinKey {
                label: "TestLabel".to_string(),
                property: "shared_id".to_string(),
            }],
        };

        let join_key = &query_join.keys[0];

        // Manually register the element in partial_joins
        {
            let mut guard = index.partial_joins.write().await;
            let partial_joins = guard
                .entry((query_join.id.clone(), value_hash))
                .or_default();
            partial_joins
                .entry(join_key.clone())
                .or_default()
                .insert(element_ref.clone());
        }

        // Verify element was registered
        {
            let guard = index.partial_joins.read().await;
            let partial_joins = guard.get(&(query_join.id.clone(), value_hash)).unwrap();
            assert!(partial_joins.get(join_key).unwrap().contains(&element_ref));
        }

        // Delete should succeed
        let result = index
            .delete_source_join(&element_ref, &query_join, join_key, &value)
            .await;
        assert!(result.is_ok());

        // Verify the emptied containers were pruned: removing the last element for
        // this (join_id, value_hash) should drop both the inner set and the outer
        // entry rather than leaving empty containers behind (see issue #641).
        {
            let guard = index.partial_joins.read().await;
            assert!(
                guard.get(&(query_join.id.clone(), value_hash)).is_none(),
                "outer partial_joins entry should be pruned once its last element is removed"
            );
        }
    }

    /// Test that after delete_source_join on a missing key, the index remains in a valid state.
    #[tokio::test]
    async fn test_delete_source_join_missing_key_index_integrity() {
        let index = InMemoryElementIndex::new();

        // Insert a different element with a different key
        let existing_ref = ElementReference::new("test_source", "existing_element");
        let existing_value = ElementValue::String("existing_key".into());
        let existing_hash = get_value_hash(&existing_value);

        let query_join = QueryJoin {
            id: "integrity_join".to_string(),
            keys: vec![QueryJoinKey {
                label: "Label".to_string(),
                property: "prop".to_string(),
            }],
        };

        let join_key = &query_join.keys[0];

        // Register the existing element
        {
            let mut guard = index.partial_joins.write().await;
            let partial_joins = guard
                .entry((query_join.id.clone(), existing_hash))
                .or_default();
            partial_joins
                .entry(join_key.clone())
                .or_default()
                .insert(existing_ref.clone());
        }

        // Try to delete a different element with a different value hash (missing key)
        let missing_ref = ElementReference::new("test_source", "missing_element");
        let missing_value = ElementValue::String("nonexistent_key".into());

        let result = index
            .delete_source_join(&missing_ref, &query_join, join_key, &missing_value)
            .await;
        assert!(result.is_ok());

        // Verify existing element is still intact
        {
            let guard = index.partial_joins.read().await;
            let partial_joins = guard.get(&(query_join.id.clone(), existing_hash)).unwrap();
            assert!(partial_joins.get(join_key).unwrap().contains(&existing_ref));
        }
    }

    /// Regression test for the `slot_affinity` memory leak: after an element is
    /// deleted, no empty `slot_affinity` entry should be retained for it.
    #[tokio::test]
    async fn test_delete_element_prunes_slot_affinity() {
        let index = InMemoryElementIndex::new();

        let element_ref = ElementReference::new("source1", "node1");
        let element = create_test_node("source1", "node1", "Person");

        index.set_element(&element, &vec![0, 1]).await.unwrap();

        // A live element should have a slot_affinity entry.
        {
            let guard = index.slot_affinity.read().await;
            assert!(guard.contains_key(&element_ref));
        }

        index.delete_element(&element_ref).await.unwrap();

        // After deletion the entry must be removed entirely, not left empty.
        {
            let guard = index.slot_affinity.read().await;
            assert!(
                !guard.contains_key(&element_ref),
                "slot_affinity should not retain an entry for a deleted element"
            );
        }
    }

    /// High-churn scenario: repeatedly creating and deleting distinct elements
    /// must not cause `slot_affinity` to grow without bound.
    #[tokio::test]
    async fn test_element_churn_does_not_grow_slot_affinity() {
        let index = InMemoryElementIndex::new();

        for i in 0..1000 {
            let element_id = format!("node{i}");
            let element_ref = ElementReference::new("source1", &element_id);
            let element = create_test_node("source1", &element_id, "Person");

            index.set_element(&element, &vec![0]).await.unwrap();
            index.delete_element(&element_ref).await.unwrap();
        }

        let guard = index.slot_affinity.read().await;
        assert!(
            guard.is_empty(),
            "slot_affinity should be empty after all churned elements are deleted, but had {} entries",
            guard.len()
        );
    }

    /// Updating a live element must keep its `slot_affinity` entry intact, since
    /// `set_element` clears then re-inserts affinity.
    #[tokio::test]
    async fn test_update_element_retains_slot_affinity() {
        let index = InMemoryElementIndex::new();

        let element_ref = ElementReference::new("source1", "node1");
        let element = create_test_node("source1", "node1", "Person");

        index.set_element(&element, &vec![0]).await.unwrap();

        // Update the same element with a different slot affinity.
        index.set_element(&element, &vec![1, 2]).await.unwrap();

        let guard = index.slot_affinity.read().await;
        let affinity = guard
            .get(&element_ref)
            .expect("live element should retain a slot_affinity entry after update");
        assert_eq!(affinity, &HashSet::from([1, 2]));
    }

    /// The early-return `None` path in `clear_slot_affinity` must be a no-op:
    /// deleting an element twice (or one that was never inserted) should return
    /// `Ok(())` without panicking and must not create a `slot_affinity` entry.
    #[tokio::test]
    async fn test_double_delete_no_slot_affinity_entry() {
        let index = InMemoryElementIndex::new();

        let element_ref = ElementReference::new("source1", "node1");
        let element = create_test_node("source1", "node1", "Person");

        index.set_element(&element, &vec![0]).await.unwrap();

        // First delete removes the element and its affinity entry.
        index.delete_element(&element_ref).await.unwrap();

        // Second delete hits the `None` early-return path and must be a no-op.
        let result = index.delete_element(&element_ref).await;
        assert!(result.is_ok(), "double delete should not panic");

        let guard = index.slot_affinity.read().await;
        assert!(
            !guard.contains_key(&element_ref),
            "double delete must not leave a slot_affinity entry"
        );
    }

    /// Deleting a relation must prune its `slot_affinity` entry as well as the
    /// `element_by_slot_in` / `element_by_slot_out` reverse indexes.
    #[tokio::test]
    async fn test_delete_relation_prunes_slot_affinity_and_reverse_indexes() {
        let index = InMemoryElementIndex::new();

        // Insert the endpoint nodes and a relation between them.
        let in_ref = ElementReference::new("source1", "in1");
        let out_ref = ElementReference::new("source1", "out1");
        let relation_ref = ElementReference::new("source1", "rel1");

        index
            .set_element(&create_test_node("source1", "in1", "Person"), &vec![0])
            .await
            .unwrap();
        index
            .set_element(&create_test_node("source1", "out1", "Person"), &vec![0])
            .await
            .unwrap();

        let relation = create_test_relation("source1", "rel1", "KNOWS", "in1", "out1");
        index.set_element(&relation, &vec![1]).await.unwrap();

        // While live, the relation should be present in all three maps.
        {
            let af_guard = index.slot_affinity.read().await;
            assert!(af_guard.contains_key(&relation_ref));

            let in_guard = index.element_by_slot_in.read().await;
            assert!(in_guard
                .get(&(1, in_ref.clone()))
                .is_some_and(|set| set.contains(&relation_ref)));

            let out_guard = index.element_by_slot_out.read().await;
            assert!(out_guard
                .get(&(1, out_ref.clone()))
                .is_some_and(|set| set.contains(&relation_ref)));
        }

        index.delete_element(&relation_ref).await.unwrap();

        // After deletion no map should retain a reference to the relation.
        {
            let af_guard = index.slot_affinity.read().await;
            assert!(
                !af_guard.contains_key(&relation_ref),
                "slot_affinity should not retain an entry for a deleted relation"
            );

            let in_guard = index.element_by_slot_in.read().await;
            assert!(!in_guard
                .get(&(1, in_ref.clone()))
                .is_some_and(|set| set.contains(&relation_ref)));

            let out_guard = index.element_by_slot_out.read().await;
            assert!(!out_guard
                .get(&(1, out_ref.clone()))
                .is_some_and(|set| set.contains(&relation_ref)));
        }
    }

    /// High-churn variant for relations: repeatedly creating and deleting
    /// distinct relations must not cause `slot_affinity` to grow without bound.
    #[tokio::test]
    async fn test_relation_churn_does_not_grow_slot_affinity() {
        let index = InMemoryElementIndex::new();

        for i in 0..1000 {
            let rel_id = format!("rel{i}");
            let relation_ref = ElementReference::new("source1", &rel_id);
            let relation = create_test_relation("source1", &rel_id, "KNOWS", "in1", "out1");

            index.set_element(&relation, &vec![0]).await.unwrap();
            index.delete_element(&relation_ref).await.unwrap();
        }

        let guard = index.slot_affinity.read().await;
        assert!(
            guard.is_empty(),
            "slot_affinity should be empty after all churned relations are deleted, but had {} entries",
            guard.len()
        );
    }

    /// Regression test for issue #641: emptied partial-join containers must be
    /// pruned so `partial_joins` does not grow unbounded for high-cardinality /
    /// high-churn join values.
    ///
    /// Registers two elements under two distinct values, deletes the last element
    /// for one value, and asserts that value's inner `HashSet` and outer entry are
    /// both dropped while the other value remains intact.
    #[tokio::test]
    async fn test_delete_source_join_prunes_emptied_containers() {
        let index = InMemoryElementIndex::new();

        let query_join = QueryJoin {
            id: "churn_join".to_string(),
            keys: vec![QueryJoinKey {
                label: "Label".to_string(),
                property: "prop".to_string(),
            }],
        };
        let join_key = &query_join.keys[0];

        // Two elements registered under two different values (two outer entries).
        let ref_a = ElementReference::new("test_source", "elem_a");
        let value_a = ElementValue::String("value_a".into());
        let hash_a = get_value_hash(&value_a);

        let ref_b = ElementReference::new("test_source", "elem_b");
        let value_b = ElementValue::String("value_b".into());
        let hash_b = get_value_hash(&value_b);

        {
            let mut guard = index.partial_joins.write().await;
            guard
                .entry((query_join.id.clone(), hash_a))
                .or_default()
                .entry(join_key.clone())
                .or_default()
                .insert(ref_a.clone());
            guard
                .entry((query_join.id.clone(), hash_b))
                .or_default()
                .entry(join_key.clone())
                .or_default()
                .insert(ref_b.clone());
        }

        assert_eq!(index.partial_joins.read().await.len(), 2);

        // Remove the only element for value_a.
        index
            .delete_source_join(&ref_a, &query_join, join_key, &value_a)
            .await
            .unwrap();

        let guard = index.partial_joins.read().await;

        // value_a's outer entry (and its now-empty inner set) must be gone.
        assert!(
            guard.get(&(query_join.id.clone(), hash_a)).is_none(),
            "outer entry for the emptied value should be pruned"
        );

        // value_b must be untouched.
        let remaining = guard
            .get(&(query_join.id.clone(), hash_b))
            .expect("outer entry for the still-populated value must remain");
        assert!(remaining.get(join_key).unwrap().contains(&ref_b));
        assert_eq!(guard.len(), 1, "only the emptied value should be pruned");
    }

    #[tokio::test]
    async fn test_get_element_as_at_archive_disabled() {
        let index = InMemoryElementIndex::new();
        let result = index
            .get_element_as_at(&ElementReference::new("source1", "node1"), 0)
            .await;
        assert!(matches!(result, Err(IndexError::ArchiveNotEnabled)));
    }

    #[tokio::test]
    async fn test_get_element_versions_archive_disabled() {
        let index = InMemoryElementIndex::new();
        let result = index
            .get_element_versions(
                &ElementReference::new("source1", "node1"),
                TimestampRange {
                    from: TimestampBound::Included(0),
                    to: 100,
                },
            )
            .await;
        assert!(matches!(result, Err(IndexError::ArchiveNotEnabled)));
    }

    #[tokio::test]
    async fn test_clear_archive_disabled_returns_not_supported() {
        let index = InMemoryElementIndex::new();
        let result = ElementArchiveIndex::clear(&index).await;
        assert!(matches!(result, Err(IndexError::NotSupported)));
    }

    #[tokio::test]
    async fn test_archive_reads_with_archive_enabled() {
        let mut index = InMemoryElementIndex::new();
        index.enable_archive();

        let element = create_test_node("source1", "node1", "Person");
        index.set_element(&element, &vec![0]).await.unwrap();

        let as_at = index
            .get_element_as_at(&ElementReference::new("source1", "node1"), 10)
            .await
            .unwrap();
        assert!(as_at.is_some());

        let versions = index
            .get_element_versions(
                &ElementReference::new("source1", "node1"),
                TimestampRange {
                    from: TimestampBound::Included(0),
                    to: 100,
                },
            )
            .await;
        assert!(versions.is_ok());
    }

    /// Regression test for issue #642: deleting a relation must prune the now-empty
    /// `(slot, node)` entries from `element_by_slot_in` / `element_by_slot_out` instead
    /// of leaving empty `HashSet`s behind, which would leak memory over time.
    #[tokio::test]
    async fn test_slot_maps_pruned_after_relation_delete() {
        let index = InMemoryElementIndex::new();

        // Insert two relations that share the same slot but reference distinct nodes.
        let rel1 = create_test_relation("source1", "rel1", "KNOWS", "in1", "out1");
        let rel2 = create_test_relation("source1", "rel2", "KNOWS", "in2", "out2");
        index.set_element(&rel1, &vec![0]).await.unwrap();
        index.set_element(&rel2, &vec![0]).await.unwrap();

        // Both relations should be indexed by their in/out nodes.
        assert_eq!(index.element_by_slot_in.read().await.len(), 2);
        assert_eq!(index.element_by_slot_out.read().await.len(), 2);

        // Delete the first relation; its (slot, node) entries should be removed entirely,
        // not left as empty sets.
        index
            .delete_element(&ElementReference::new("source1", "rel1"))
            .await
            .unwrap();

        {
            let in_guard = index.element_by_slot_in.read().await;
            let out_guard = index.element_by_slot_out.read().await;
            assert_eq!(in_guard.len(), 1, "empty in-slot entry should be pruned");
            assert_eq!(out_guard.len(), 1, "empty out-slot entry should be pruned");
            assert!(!in_guard.contains_key(&(0, ElementReference::new("source1", "in1"))));
            assert!(!out_guard.contains_key(&(0, ElementReference::new("source1", "out1"))));
        }

        // Delete the second relation; both maps should now be empty.
        index
            .delete_element(&ElementReference::new("source1", "rel2"))
            .await
            .unwrap();

        assert!(
            index.element_by_slot_in.read().await.is_empty(),
            "element_by_slot_in should be empty after all relations are deleted"
        );
        assert!(
            index.element_by_slot_out.read().await.is_empty(),
            "element_by_slot_out should be empty after all relations are deleted"
        );
        assert!(
            index.slot_affinity.read().await.is_empty(),
            "slot_affinity should be empty after all relations are deleted"
        );
    }

    /// Regression test for the multi-slot case: a relation assigned to several slots must
    /// have every `(slot, node)` entry pruned on delete, since `clear_slot_affinity`
    /// iterates over all slots in the affinity set.
    #[tokio::test]
    async fn test_slot_maps_pruned_multi_slot() {
        let index = InMemoryElementIndex::new();

        let rel = create_test_relation("source1", "rel1", "KNOWS", "in1", "out1");
        index.set_element(&rel, &vec![0, 1, 2]).await.unwrap();

        // One (slot, node) entry per slot for both in and out maps.
        assert_eq!(index.element_by_slot_in.read().await.len(), 3);
        assert_eq!(index.element_by_slot_out.read().await.len(), 3);

        index
            .delete_element(&ElementReference::new("source1", "rel1"))
            .await
            .unwrap();

        assert!(
            index.element_by_slot_in.read().await.is_empty(),
            "all in-slot entries should be pruned across every slot"
        );
        assert!(
            index.element_by_slot_out.read().await.is_empty(),
            "all out-slot entries should be pruned across every slot"
        );
        assert!(
            index.slot_affinity.read().await.is_empty(),
            "slot_affinity should be pruned for the deleted relation"
        );
    }

    /// Ensure pruning does not drop a shared `(slot, node)` entry while another relation
    /// still references that node.
    #[tokio::test]
    async fn test_slot_maps_retain_shared_node_entry() {
        let index = InMemoryElementIndex::new();

        // Two relations sharing the same in_node and out_node.
        let rel1 = create_test_relation("source1", "rel1", "KNOWS", "in1", "out1");
        let rel2 = create_test_relation("source1", "rel2", "KNOWS", "in1", "out1");
        index.set_element(&rel1, &vec![0]).await.unwrap();
        index.set_element(&rel2, &vec![0]).await.unwrap();

        let in_key = (0, ElementReference::new("source1", "in1"));
        let out_key = (0, ElementReference::new("source1", "out1"));

        // Delete one relation: the shared entries must survive with the remaining ref.
        index
            .delete_element(&ElementReference::new("source1", "rel1"))
            .await
            .unwrap();

        {
            let in_guard = index.element_by_slot_in.read().await;
            let out_guard = index.element_by_slot_out.read().await;
            let in_set = in_guard.get(&in_key).expect("shared in entry retained");
            let out_set = out_guard.get(&out_key).expect("shared out entry retained");
            assert_eq!(in_set.len(), 1);
            assert_eq!(out_set.len(), 1);
            assert!(in_set.contains(&ElementReference::new("source1", "rel2")));
            assert!(out_set.contains(&ElementReference::new("source1", "rel2")));
        }

        // Delete the last relation: entries should now be pruned.
        index
            .delete_element(&ElementReference::new("source1", "rel2"))
            .await
            .unwrap();

        assert!(index.element_by_slot_in.read().await.is_empty());
        assert!(index.element_by_slot_out.read().await.is_empty());
    }
}
