//! Dotted sub-verb handlers: recall.embed, recall.candidates, recall.fuse, recall.rerank, recall.score.

use serde::Deserialize;
use serde_json::{json, Value};

use khive_runtime::{NamespaceToken, RuntimeError, VerbRegistry};

use crate::config::RecallConfig;
use crate::rerank::{weighted_rerank, RerankFeatures};
use crate::MemoryPack;

use super::common::{
    compute_score, deser, fuse_candidates, make_pipeline, note_matches_tags,
    recall_candidate_count, search_source_label, to_json, validate_memory_type,
    RecallCandidateParams, RecallParams, TextSnippetPolicy, RECALL_DIAGNOSTIC_SNIPPET_CHARS,
};

impl MemoryPack {
    pub(crate) async fn handle_recall_embed(&self, params: Value) -> Result<Value, RuntimeError> {
        #[derive(Deserialize)]
        struct EmbedParams {
            query: String,
            #[serde(default)]
            include_embeddings: bool,
        }
        let p: EmbedParams = deser(params)?;

        let model_names = self.runtime.registered_embedding_model_names();
        if model_names.is_empty() {
            return to_json(&json!({
                "embedding": null,
                "model": null,
                "engines": [],
            }));
        }

        let mut engines: Vec<Value> = Vec::with_capacity(model_names.len());
        let mut primary_embedding: Option<Vec<f32>> = None;
        let primary_model = self.runtime.default_embedder_name().to_owned();

        for model_name in &model_names {
            match self.runtime.embed_with_model(model_name, &p.query).await {
                Ok(vec) => {
                    let dims = vec.len();
                    if primary_embedding.is_none() || model_name == &primary_model {
                        primary_embedding = Some(vec.clone());
                    }
                    let mut engine = json!({
                        "model": model_name,
                        "dimensions": dims,
                    });
                    if p.include_embeddings {
                        engine["embedding"] = json!(vec);
                    }
                    engines.push(engine);
                }
                Err(e) => {
                    engines.push(json!({
                        "model": model_name,
                        "error": e.to_string(),
                    }));
                }
            }
        }

        match primary_embedding {
            Some(vec) => {
                let dims = vec.len();
                let mut response = json!({
                    "dimensions": dims,
                    "engines": engines,
                });
                if p.include_embeddings {
                    response["embedding"] = json!(vec);
                }
                to_json(&response)
            }
            None => to_json(&json!({
                "embedding": null,
                "model": null,
                "engines": engines,
            })),
        }
    }

    pub(crate) async fn handle_recall_candidates(
        &self,
        token: &NamespaceToken,
        params: Value,
    ) -> Result<Value, RuntimeError> {
        let p: RecallParams = deser(params)?;
        let cfg = p.effective_config(self.active_config());
        cfg.validate()?;

        let limit = p.limit.unwrap_or(10).min(100);
        let scoring_cfg = cfg.scoring.clone().unwrap_or_default();
        let candidate_limit =
            recall_candidate_count(&cfg, limit).min(scoring_cfg.max_recall_candidates as u32);
        let effective_fts_gather_cand = crate::config::RecallFtsGatherConfig::from_env()
            .map_err(|e| RuntimeError::InvalidInput(format!("fts_gather env parse error: {e}")))?
            .unwrap_or_else(|| cfg.fts_gather.clone());

        let candidates = self
            .collect_recall_candidates(
                &p.query,
                token,
                RecallCandidateParams {
                    candidate_limit,
                    embedding_model: p.embedding_model.as_deref(),
                    is_cjk: false,
                    scoring_cfg: &scoring_cfg,
                    snippet_policy: TextSnippetPolicy::Include {
                        chars: RECALL_DIAGNOSTIC_SNIPPET_CHARS,
                    },
                    fts_gather: &effective_fts_gather_cand,
                },
            )
            .await?;

        let (memory_ids, _) = self.load_memory_candidate_notes(token, &candidates).await?;
        let text_candidates: Vec<Value> = candidates
            .text_hits
            .iter()
            .filter(|hit| memory_ids.contains(&hit.subject_id))
            .map(|hit| {
                json!({
                    "note_id": hit.subject_id.to_string(),
                    "score": hit.score.to_f64(),
                    "rank": hit.rank,
                    "title": hit.title.as_deref(),
                    "snippet": hit.snippet.as_deref(),
                })
            })
            .collect();

        let all_vector_hits = candidates.all_vector_hits();
        let vector_candidates: Vec<Value> = all_vector_hits
            .iter()
            .map(|hit| {
                json!({
                    "note_id": hit.subject_id.to_string(),
                    "score": hit.score.to_f64(),
                    "rank": hit.rank,
                })
            })
            .collect();

        let mut response = json!({
            "namespace": candidates.namespace,
            "candidate_limit": candidate_limit,
            "text_candidates": text_candidates,
            "vector_candidates": vector_candidates,
        });

        if candidates.vector_hits_per_model.len() > 1 {
            let per_model: serde_json::Map<String, Value> = candidates
                .vector_hits_per_model
                .iter()
                .map(|(model, hits)| {
                    let hits_json: Vec<Value> = hits
                        .iter()
                        .map(|h| {
                            json!({
                                "note_id": h.subject_id.to_string(),
                                "score": h.score.to_f64(),
                                "rank": h.rank,
                            })
                        })
                        .collect();
                    (model.clone(), Value::Array(hits_json))
                })
                .collect();
            response["vector_candidates_per_model"] = Value::Object(per_model);
        }

        to_json(&response)
    }

    pub(crate) async fn handle_recall_fuse(
        &self,
        token: &NamespaceToken,
        params: Value,
        _registry: &VerbRegistry,
    ) -> Result<Value, RuntimeError> {
        let p: RecallParams = deser(params)?;
        if let Some(mt) = &p.memory_type {
            validate_memory_type(mt)?;
        }

        let cfg = p.effective_config(self.active_config());
        cfg.validate()?;

        let limit = p.limit.unwrap_or(10).min(100);
        let scoring_cfg_fuse = cfg.scoring.clone().unwrap_or_default();
        let candidate_limit =
            recall_candidate_count(&cfg, limit).min(scoring_cfg_fuse.max_recall_candidates as u32);
        let effective_fts_gather_fuse = crate::config::RecallFtsGatherConfig::from_env()
            .map_err(|e| RuntimeError::InvalidInput(format!("fts_gather env parse error: {e}")))?
            .unwrap_or_else(|| cfg.fts_gather.clone());

        let candidates = self
            .collect_recall_candidates(
                &p.query,
                token,
                RecallCandidateParams {
                    candidate_limit,
                    embedding_model: p.embedding_model.as_deref(),
                    is_cjk: false,
                    scoring_cfg: &scoring_cfg_fuse,
                    snippet_policy: TextSnippetPolicy::Include {
                        chars: RECALL_DIAGNOSTIC_SNIPPET_CHARS,
                    },
                    fts_gather: &effective_fts_gather_fuse,
                },
            )
            .await?;
        let (memory_ids, notes_by_id) =
            self.load_memory_candidate_notes(token, &candidates).await?;

        let fused = fuse_candidates(&candidates, &memory_ids, &cfg, candidate_limit as usize);

        let fused_candidates: Vec<Value> = fused
            .into_iter()
            .filter_map(|hit| {
                let note = notes_by_id.get(&hit.entity_id)?;
                if let Some(mt) = &p.memory_type {
                    let stored = note
                        .properties
                        .as_ref()
                        .and_then(|props| props.get("memory_type"))
                        .and_then(|v| v.as_str());
                    if stored != Some(mt.as_str()) {
                        return None;
                    }
                }
                if let Some(filter_tags) = p.tags.as_ref().filter(|tags| !tags.is_empty()) {
                    if !note_matches_tags(note.properties.as_ref(), filter_tags, p.tag_mode) {
                        return None;
                    }
                }
                Some(json!({
                    "note_id": hit.entity_id.to_string(),
                    "fused_score": hit.score.to_f64(),
                    "source": search_source_label(hit.source),
                    "title": hit.title,
                    "snippet": hit.snippet,
                }))
            })
            .collect();

        to_json(&json!({
            "strategy": cfg.fuse_strategy,
            "candidate_limit": candidate_limit,
            "fused_candidates": fused_candidates,
        }))
    }

    pub(crate) async fn handle_recall_rerank(&self, params: Value) -> Result<Value, RuntimeError> {
        #[derive(Deserialize)]
        struct RerankParams {
            candidates: Vec<serde_json::Value>,
            config: Option<RecallConfig>,
        }
        let p: RerankParams = deser(params)?;
        let cfg = p.config.unwrap_or_else(|| self.active_config());
        cfg.validate()?;

        let active_rerankers: Vec<&String> = cfg
            .reranker_weights
            .keys()
            .filter(|k| cfg.reranker_weights[*k] > 0.0)
            .collect();

        let reranked: Vec<serde_json::Value> = p
            .candidates
            .iter()
            .map(|candidate| {
                let id = candidate
                    .get("note_id")
                    .cloned()
                    .unwrap_or(serde_json::Value::Null);

                if cfg.reranker_weights.is_empty() {
                    return json!({
                        "note_id": id,
                        "rerank_scores": {},
                        "rerank_score": 0.0_f64,
                    });
                }

                let fused_score = candidate
                    .get("fused_score")
                    .and_then(|v| v.as_f64())
                    .unwrap_or(0.0);
                let salience = candidate
                    .get("salience")
                    .and_then(|v| v.as_f64())
                    .unwrap_or(0.0);
                let decay_factor = candidate
                    .get("decay_factor")
                    .and_then(|v| v.as_f64())
                    .unwrap_or(0.01);
                let age_days = candidate
                    .get("age_days")
                    .and_then(|v| v.as_f64())
                    .unwrap_or(0.0);
                let temporal = candidate
                    .get("temporal")
                    .and_then(|v| v.as_f64())
                    .unwrap_or_else(|| {
                        let k = std::f64::consts::LN_2 / cfg.temporal_half_life_days;
                        (-k * age_days).exp()
                    });
                let effective_salience = cfg.decay_model.apply(
                    salience,
                    age_days,
                    decay_factor,
                    cfg.temporal_half_life_days,
                );
                let source_str = candidate
                    .get("source")
                    .and_then(|v| v.as_str())
                    .unwrap_or("");
                let text_match = matches!(source_str, "text" | "both");
                let vector_match = matches!(source_str, "vector" | "both");

                let features = RerankFeatures {
                    relevance: fused_score,
                    salience: effective_salience,
                    temporal,
                    text_match,
                    vector_match,
                };
                let rerank_score = weighted_rerank(&features, &cfg.reranker_weights);

                let mut rerank_scores = serde_json::Map::new();
                for (name, &weight) in &cfg.reranker_weights {
                    if weight == 0.0 {
                        continue;
                    }
                    let fv = match name.as_str() {
                        "relevance" => features.relevance,
                        "salience" => features.salience,
                        "temporal" => features.temporal,
                        "text_match" => f64::from(features.text_match),
                        "vector_match" => f64::from(features.vector_match),
                        _ => continue,
                    };
                    rerank_scores.insert(name.clone(), json!(weight * fv));
                }

                json!({
                    "note_id": id,
                    "rerank_scores": rerank_scores,
                    "rerank_score": rerank_score,
                })
            })
            .collect();

        to_json(&json!({
            "reranked": reranked,
            "active_rerankers": active_rerankers.iter().map(|n| n.as_str()).collect::<Vec<_>>(),
        }))
    }

    pub(crate) async fn handle_recall_score(&self, params: Value) -> Result<Value, RuntimeError> {
        #[derive(Deserialize)]
        struct ScoreParams {
            rrf: f64,
            salience: f64,
            decay_factor: f64,
            age_days: f64,
            config: Option<RecallConfig>,
        }
        let p: ScoreParams = deser(params)?;
        for (name, val) in [
            ("rrf", p.rrf),
            ("salience", p.salience),
            ("decay_factor", p.decay_factor),
            ("age_days", p.age_days),
        ] {
            if !val.is_finite() {
                return Err(RuntimeError::InvalidInput(format!(
                    "{name} must be a finite number, got {val}"
                )));
            }
        }
        let cfg = p.config.unwrap_or_else(|| self.active_config());
        cfg.validate()?;
        let pipeline = make_pipeline(&cfg);
        let (total, breakdown) = compute_score(
            &cfg,
            &pipeline,
            p.rrf,
            p.salience,
            p.decay_factor,
            p.age_days,
        );
        to_json(&json!({
            "total": total,
            "breakdown": breakdown,
        }))
    }
}
