use std::collections::HashSet;

use anyhow::Result;
use chrono::Utc;
use uuid::Uuid;

use crate::schema::{Memory, RelationType};
use crate::store::Store;

#[derive(Debug, Clone)]
pub struct QueryRequest {
    pub text: String,
    pub embedding: Vec<f32>,
    pub limit: usize,
    pub filters: QueryFilters,
}

#[derive(Debug, Clone, Default)]
pub struct QueryFilters {
    pub source: Option<String>,
    pub memory_type: Option<crate::schema::MemoryType>,
    pub min_confidence: Option<f32>,
    pub entity_names: Vec<String>,
    pub project_path: Option<String>,
}

#[derive(Debug, Clone)]
pub struct QueryResult {
    pub memory: Memory,
    pub score: f32,
    pub path: Vec<Uuid>,
}

const MIN_SIMILARITY: f32 = 0.59;
const PROJECT_AFFINITY_BOOST: f32 = 1.15;

pub struct QueryEngine<'a, S: Store> {
    store: &'a S,
    vector_weight: f32,
    graph_weight: f32,
    recency_weight: f32,
}

impl<'a, S: Store> QueryEngine<'a, S> {
    pub fn new(store: &'a S) -> Self {
        Self {
            store,
            vector_weight: 0.5,
            graph_weight: 0.3,
            recency_weight: 0.2,
        }
    }

    pub fn with_weights(mut self, vector: f32, graph: f32, recency: f32) -> Self {
        self.vector_weight = vector;
        self.graph_weight = graph;
        self.recency_weight = recency;
        self
    }

    pub fn recall(&self, request: &QueryRequest) -> Result<Vec<QueryResult>> {
        let vector_results = self
            .store
            .vector_search(&request.embedding, request.limit * 3)?;

        let query_entity_ids = self.resolve_query_entities(&request.filters.entity_names);

        let mut scored: Vec<QueryResult> = Vec::new();

        for (memory, similarity) in vector_results {
            // Gate membership on similarity alone. Graph and recency only
            // re-rank, so folding them into the cutoff caps the similarity
            // term below any useful threshold and would silently drop aged,
            // unconnected memories that match the query perfectly.
            if similarity < MIN_SIMILARITY {
                continue;
            }
            if let Some(min_conf) = request.filters.min_confidence
                && memory.confidence < min_conf
            {
                continue;
            }
            if let Some(ref source) = request.filters.source
                && &memory.source != source
            {
                continue;
            }
            if let Some(ref mt) = request.filters.memory_type
                && &memory.memory_type != mt
            {
                continue;
            }

            let recency_score = self.compute_recency(&memory);
            let graph_score = self
                .compute_graph_relevance(&memory, &query_entity_ids)
                .unwrap_or(0.0);

            let mut rank_score = (similarity * self.vector_weight)
                + (graph_score * self.graph_weight)
                + (recency_score * self.recency_weight);

            if let Some(ref qp) = request.filters.project_path {
                if memory.project_path.as_deref() == Some(qp.as_str()) {
                    rank_score *= PROJECT_AFFINITY_BOOST;
                }
            }

            scored.push(QueryResult {
                memory,
                score: rank_score,
                path: Vec::new(),
            });
        }

        scored.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        scored.truncate(request.limit);

        Ok(scored)
    }

    fn compute_recency(&self, memory: &Memory) -> f32 {
        let hours_since_access = Utc::now()
            .signed_duration_since(memory.last_accessed)
            .num_hours() as f32;

        let decay = (-hours_since_access / (24.0 * 30.0)).exp();
        let access_boost = (memory.access_count as f32).ln_1p() / 10.0;

        (decay + access_boost).min(1.0)
    }

    /// Resolve the entity names referenced by a query to their stored entity
    /// ids. Names come from gazetteer extraction over the query text, so they
    /// match stored entity names exactly; unknown names resolve to nothing.
    fn resolve_query_entities(&self, names: &[String]) -> HashSet<Uuid> {
        let mut ids = HashSet::new();
        for name in names {
            if let Ok(Some(entity)) = self.store.find_entity_by_name(name) {
                ids.insert(entity.id);
            }
        }
        ids
    }

    fn compute_graph_relevance(
        &self,
        memory: &Memory,
        query_entity_ids: &HashSet<Uuid>,
    ) -> Result<f32> {
        let scored_types = [
            (RelationType::Reinforces, 0.3_f32),
            (RelationType::RelatesTo, 0.2),
            (RelationType::DistilledFrom, 0.15),
            (RelationType::Mentions, 0.1),
            (RelationType::DerivedFrom, 0.05),
            (RelationType::Contradicts, -0.1),
            (RelationType::Supersedes, -0.2),
        ];

        // Structural score is a query-independent prior: how reinforced and
        // non-superseded a memory is. Mention overlap with the query is
        // collected in the same pass so the Mentions edges are only fetched once.
        let mut structural = 0.0_f32;
        let mut mention_hits = 0usize;
        for (rt, boost) in &scored_types {
            if let Ok(rels) = self.store.get_relations(memory.id, Some(*rt)) {
                for rel in &rels {
                    let b = if *rt == RelationType::RelatesTo {
                        boost * rel.strength
                    } else {
                        *boost
                    };
                    structural += b;
                    if *rt == RelationType::Mentions && query_entity_ids.contains(&rel.to_id) {
                        mention_hits += 1;
                    }
                }
            }
        }
        let structural = structural.clamp(0.0, 1.0);

        if query_entity_ids.is_empty() {
            return Ok(structural);
        }

        // Query-conditioned proximity: the fraction of the query's entities this
        // memory directly mentions. A memory topically adjacent to what the
        // query is about should outrank a generically well-connected but
        // off-topic one, so proximity dominates the structural prior.
        let proximity = (mention_hits as f32 / query_entity_ids.len() as f32).min(1.0);
        Ok((0.7 * proximity + 0.3 * structural).clamp(0.0, 1.0))
    }

    pub fn store(&self) -> &S {
        self.store
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::schema::{Conversation, Entity, MemoryType, Relation};
    use chrono::Duration;

    struct StubStore {
        vector_results: Vec<(Memory, f32)>,
        relations: Vec<Relation>,
        entities: Vec<Entity>,
    }

    impl StubStore {
        fn new(vector_results: Vec<(Memory, f32)>, relations: Vec<Relation>) -> Self {
            Self {
                vector_results,
                relations,
                entities: Vec::new(),
            }
        }
    }

    impl Store for StubStore {
        fn vector_search(&self, _embedding: &[f32], _limit: usize) -> Result<Vec<(Memory, f32)>> {
            Ok(self.vector_results.clone())
        }

        fn get_relations(
            &self,
            node_id: Uuid,
            relation_type: Option<RelationType>,
        ) -> Result<Vec<Relation>> {
            Ok(self
                .relations
                .iter()
                .filter(|r| r.from_id == node_id)
                .filter(|r| relation_type.map(|rt| rt == r.relation_type).unwrap_or(true))
                .cloned()
                .collect())
        }

        fn store_memory(&self, _m: &Memory) -> Result<()> {
            unimplemented!()
        }
        fn get_memory(&self, _id: Uuid) -> Result<Option<Memory>> {
            unimplemented!()
        }
        fn delete_memory(&self, _id: Uuid) -> Result<()> {
            unimplemented!()
        }
        fn store_entity(&self, _e: &Entity) -> Result<()> {
            unimplemented!()
        }
        fn get_entity(&self, _id: Uuid) -> Result<Option<Entity>> {
            unimplemented!()
        }
        fn find_entity_by_name(&self, name: &str) -> Result<Option<Entity>> {
            Ok(self.entities.iter().find(|e| e.name == name).cloned())
        }
        fn store_conversation(&self, _c: &Conversation) -> Result<()> {
            unimplemented!()
        }
        fn store_relation(&self, _r: &Relation) -> Result<()> {
            unimplemented!()
        }
        fn traverse(&self, _id: Uuid, _depth: u32) -> Result<Vec<(Memory, Vec<Relation>)>> {
            unimplemented!()
        }
        fn memories_by_source(&self, _s: &str) -> Result<Vec<Memory>> {
            unimplemented!()
        }
        fn memories_by_type(&self, _mt: MemoryType) -> Result<Vec<Memory>> {
            unimplemented!()
        }
        fn memories_needing_decay(&self, _days: u32) -> Result<Vec<Memory>> {
            unimplemented!()
        }
        fn update_memory(&self, _m: &Memory) -> Result<()> {
            unimplemented!()
        }
        fn record_access(&self, _memory: &Memory) -> Result<()> {
            unimplemented!()
        }
        fn text_search(&self, _q: &str, _limit: usize) -> Result<Vec<Memory>> {
            unimplemented!()
        }
        fn memory_count(&self) -> Result<usize> {
            unimplemented!()
        }
        fn all_memory_ids(&self) -> Result<Vec<Uuid>> {
            unimplemented!()
        }
        fn all_relations(&self) -> Result<Vec<Relation>> {
            unimplemented!()
        }
    }

    fn memory_aged(content: &str, days_old: i64) -> Memory {
        let when = Utc::now() - Duration::days(days_old);
        let mut m = Memory::new(
            content.to_string(),
            MemoryType::Semantic,
            "test".to_string(),
            String::new(),
        );
        m.created_at = when;
        m.last_accessed = when;
        m
    }

    fn query() -> QueryRequest {
        QueryRequest {
            text: "query text".to_string(),
            embedding: vec![0.0_f32; 384],
            limit: 10,
            filters: QueryFilters::default(),
        }
    }

    #[test]
    fn aged_unconnected_memory_survives_on_strong_similarity() {
        // The memory is 180 days old with no relations, so graph and recency
        // contribute nothing; it must still be recalled because a strong
        // semantic match is what relevance means.
        let store = StubStore::new(
            vec![(
                memory_aged("kuzu was chosen as the embedded graph store", 180),
                0.92,
            )],
            Vec::new(),
        );
        let results = QueryEngine::new(&store).recall(&query()).unwrap();
        assert_eq!(results.len(), 1);
    }

    #[test]
    fn well_connected_low_similarity_memory_is_excluded() {
        // A fresh, heavily reinforced memory whose similarity is weak must be
        // dropped, because graph and recency re-rank matches but must not
        // promote an irrelevant memory past the similarity floor.
        let weak = memory_aged("an entirely unrelated topic", 0);
        let weak_id = weak.id;
        let reinforces = |to| Relation {
            from_id: weak_id,
            to_id: to,
            relation_type: RelationType::Reinforces,
            strength: 1.0,
            context: None,
        };
        let store = StubStore::new(
            vec![(weak, 0.45)],
            vec![
                reinforces(Uuid::new_v4()),
                reinforces(Uuid::new_v4()),
                reinforces(Uuid::new_v4()),
            ],
        );
        let results = QueryEngine::new(&store).recall(&query()).unwrap();
        assert!(results.is_empty());
    }

    #[test]
    fn results_are_ranked_by_blended_score() {
        // Both memories clear the similarity floor with equal similarity, so
        // recency must break the tie and rank the fresher memory first.
        let store = StubStore::new(
            vec![
                (memory_aged("older relevant note", 200), 0.80),
                (memory_aged("fresher relevant note", 0), 0.80),
            ],
            Vec::new(),
        );
        let results = QueryEngine::new(&store).recall(&query()).unwrap();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].memory.content, "fresher relevant note");
    }

    fn mentions(from: Uuid, to: Uuid) -> Relation {
        Relation {
            from_id: from,
            to_id: to,
            relation_type: RelationType::Mentions,
            strength: 1.0,
            context: None,
        }
    }

    #[test]
    fn memory_mentioning_a_query_entity_outranks_an_equal_match_that_does_not() {
        // Two memories, same similarity and same age. One mentions the entity
        // the query is about; it must rank first because graph proximity is now
        // query-conditioned rather than a raw connectedness prior.
        let on_topic = memory_aged("note about the kuzu store", 10);
        let off_topic = memory_aged("note about something else", 10);
        let kuzu = Entity::new("kuzu".to_string(), "technology".to_string());

        let mut store = StubStore::new(
            vec![(off_topic.clone(), 0.80), (on_topic.clone(), 0.80)],
            vec![mentions(on_topic.id, kuzu.id)],
        );
        store.entities = vec![kuzu];

        let mut req = query();
        req.filters.entity_names = vec!["kuzu".to_string()];

        let results = QueryEngine::new(&store).recall(&req).unwrap();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].memory.content, "note about the kuzu store");
    }

    #[test]
    fn unresolvable_query_entity_falls_back_to_structural_ranking() {
        // An entity name the store has never seen resolves to nothing, so the
        // graph term reverts to the structural prior and recency breaks the tie
        // exactly as when no entities are named.
        let store = StubStore::new(
            vec![
                (memory_aged("older relevant note", 200), 0.80),
                (memory_aged("fresher relevant note", 0), 0.80),
            ],
            Vec::new(),
        );
        let mut req = query();
        req.filters.entity_names = vec!["never-seen-entity".to_string()];

        let results = QueryEngine::new(&store).recall(&req).unwrap();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].memory.content, "fresher relevant note");
    }
}
