use std::collections::HashMap;
use sqlx::Row;
use crate::mcp::types::AppState;

/// 検索クエリ文字列をベクトルに変換する。エディションに応じて LSA または埋め込みモデルを使用（1 箇所で切り替え）。
#[cfg(feature = "community")]
async fn get_query_vector(state: &AppState, text: &str) -> Option<Vec<f32>> {
    let lsa_guard = state.lsa_model.read().await;
    let model = lsa_guard.as_ref()?;
    let mut query_counts = HashMap::new();
    let tokens = state.tokenizer.tokenize_to_vec(text).unwrap_or_default();
    for token in tokens {
        if let Some(&tid) = model.vocabulary.get(&token) {
            *query_counts.entry(tid).or_insert(0.0) += 1.0;
        }
    }
    if query_counts.is_empty() {
        return None;
    }
    let mut query_vec = ndarray::Array1::zeros(model.vocabulary.len());
    for (tid, count) in query_counts {
        query_vec[tid] = count;
    }
    let query_lsa = model.project_query(&query_vec).ok()?;
    let mut out: Vec<f32> = query_lsa.iter().map(|&x| x as f32).collect();
    if out.len() < 50 {
        out.resize(50, 0.0);
    } else {
        out.truncate(50);
    }
    Some(out)
}

#[cfg(feature = "pro")]
async fn get_query_vector(state: &AppState, text: &str) -> Option<Vec<f32>> {
    let guard = state.embedding_model.read().await;
    let model = guard.as_ref()?;
    model.encode(text).ok()
}

pub async fn handle_search_text(
    state: &AppState,
    actual_method: &str,
    args: &serde_json::Map<String, serde_json::Value>,
) -> Option<serde_json::Value> {
    let search_content = if actual_method == "lsa_search" {
        args.get("query").and_then(|v| v.as_str()).unwrap_or("")
    } else {
        args.get("content").and_then(|v| v.as_str()).unwrap_or("")
    };
    let search_limit = args.get("limit").and_then(|v| v.as_i64()).unwrap_or(5);
    // true のとき検索結果を文書単位にまとめ、同一文書のチャンクを結合して返す。未指定時は true。
    let group_by_document = args.get("group_by_document")
        .or_else(|| args.get("groupByDocument"))
        .and_then(|v| v.as_bool())
        .unwrap_or(true);
    // 足切り: この値未満の similarity は返さない。未指定時は 0.3（ノイズを落としつつ関連しそうなものを残す）
    let min_score = args.get("min_score")
        .or_else(|| args.get("minScore"))
        .and_then(|v| v.as_f64())
        .unwrap_or(0.3)
        .clamp(0.0, 1.0) as f32;
    let filter_category = args.get("category")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .to_string();

    if search_content.is_empty() {
        log::info!("[search] query=empty -> skipped");
        let empty_payload = serde_json::json!({
            "items": [],
            "vector_search_used": false
        });
        let empty_text = serde_json::to_string_pretty(&empty_payload).unwrap_or_else(|_| r#"{"items":[],"vector_search_used":false}"#.to_string());
        return Some(serde_json::json!({
            "content": [{ "type": "text", "text": empty_text }]
        }));
    }

    let query_preview = if search_content.len() > 60 {
        format!("{}...", &search_content[..60])
    } else {
        search_content.to_string()
    };
    log::info!("[search] query={:?} limit={} min_score={}", query_preview, search_limit, min_score);

    // 1. FTS5 (BM25) search - Elasticsearch-like statistical ranking
    let mut fts_results = HashMap::new();
    if let Ok(rows) = sqlx::query(
        "SELECT rowid, bm25(items_fts) as score 
         FROM items_fts 
         WHERE items_fts MATCH ? 
         ORDER BY score LIMIT ?"
    ).bind(search_content).bind(search_limit).fetch_all(&state.db_pool).await {
        for row in rows {
            let id: i64 = row.get(0);
            let bm25_score: f64 = row.get(1);
            let sim = (1.0 - (bm25_score / 10.0).tanh()).clamp(0.0, 1.0) as f32;
            fts_results.insert(id, sim);
        }
    }
    // FTS が 0 件のとき（1文字クエリなど trigram がヒットしにくい）: LIKE でフォールバック（スコアはキーワード一致として 0.4 に統一）
    if fts_results.is_empty() && !search_content.trim().is_empty() {
        let like_pattern = format!("%{}%", search_content.trim());
        if let Ok(rows) = sqlx::query(
            "SELECT i.id FROM items i JOIN documents d ON i.document_id = d.id WHERE i.content LIKE ? LIMIT ?"
        ).bind(&like_pattern).bind(search_limit).fetch_all(&state.db_pool).await {
            let like_count = rows.len();
            for row in rows {
                let id: i64 = row.get(0);
                fts_results.insert(id, 0.4f32);
            }
            log::info!("[search] FTS=0, LIKE fallback hits={} (scores fixed 0.4)", like_count);
        }
    }
    log::info!("[search] FTS hits={}", fts_results.len());

    // 2. Vector Search（クエリベクトルは get_query_vector で 1 箇所に集約）
    let mut final_results: HashMap<i64, serde_json::Value> = HashMap::new();
    let mut vector_search_used = false;
    if let Some(query_embedding) = get_query_vector(state, search_content).await {
        let mut vector_hits = Vec::new();
        let _hnsw_is_some = {
            let hnsw_idx_guard = state.hnsw_index.read().await;
            if let Some(h_ptr) = hnsw_idx_guard.as_ref() {
                let neighbors = h_ptr.search(&query_embedding, (search_limit * 2) as usize, 100);
                for n in neighbors {
                    vector_hits.push((n.d_id as i64, 1.0f32 - n.distance));
                }
                true
            } else {
                false
            }
        };
        if vector_hits.is_empty() {
            if let Ok(rows) = sqlx::query(
                "SELECT id, distance FROM vec_items WHERE embedding MATCH ? AND k = ?",
            )
            .bind(serde_json::to_string(&query_embedding).unwrap_or_else(|_| "[]".to_string()))
            .bind(search_limit * 2)
            .fetch_all(&state.db_pool)
            .await
            {
                for r in rows {
                    let id: i64 = r.get(0);
                    let dist: f64 = r.get(1);
                    vector_hits.push((id, (1.0 - (dist / 2.0)) as f32));
                }
            }
        }
        if vector_hits.is_empty() {
            #[cfg(feature = "pro")]
            {
                let vec_count: i64 = sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM vec_items")
                    .fetch_one(&state.db_pool)
                    .await
                    .unwrap_or(0);
                log::warn!(
                    "[search] Pro: 近似近傍が 0 件 (HNSW={} vec_items={})。RE-INDEX を実行するか、バックグラウンドで再構築を試行します。",
                    if _hnsw_is_some { "あり" } else { "なし" },
                    vec_count
                );
                use std::sync::atomic::Ordering;
                if !_hnsw_is_some
                    && state
                        .pro_hnsw_rebuild_requested
                        .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
                        .is_ok()
                {
                    let state_clone = state.clone();
                    tokio::spawn(async move {
                        log::info!("[search] Pro: HNSW 遅延構築を開始します（vec_items 再読込・HNSW 構築）");
                        crate::mcp::system::train_lsa_and_sync_hnsw(state_clone).await;
                        log::info!("[search] Pro: HNSW 遅延構築完了。次回検索から近似近傍が使われます。");
                    });
                }
            }
        } else {
            vector_search_used = true;
            log::info!("[search] vector hits={}", vector_hits.len());
        }
        for (id, v_sim) in vector_hits {
            let f_sim = fts_results.get(&id).cloned().unwrap_or(0.0);
            let final_sim = v_sim.max(f_sim);
            if let Ok(Some((content, path, mime, document_id, category))) =
                crate::db::get_item_content_with_doc(&state.db_pool, id).await
            {
                if !filter_category.is_empty() && category != filter_category {
                    continue;
                }
                final_results.insert(id, serde_json::json!({
                    "id": id,
                    "document_id": document_id,
                    "content": content,
                    "path": path,
                    "mime": mime,
                    "category": category,
                    "similarity": final_sim.clamp(0.0, 1.0)
                }));
            }
        }
    }

    // 3. Add FTS results (and for Community, remaining FTS not in vector results)
    for (id, f_sim) in fts_results {
        if !final_results.contains_key(&id) {
            if let Ok(Some((content, path, mime, document_id, category))) =
                crate::db::get_item_content_with_doc(&state.db_pool, id).await
            {
                if !filter_category.is_empty() && category != filter_category {
                    continue;
                }
                final_results.insert(id, serde_json::json!({
                    "id": id,
                    "document_id": document_id,
                    "content": content,
                    "path": path,
                    "mime": mime,
                    "category": category,
                    "similarity": f_sim.clamp(0.0, 1.0)
                }));
            }
        }
    }

    let mut sorted: Vec<_> = final_results.into_values().collect();
    sorted.sort_by(|a, b| {
        b.get("similarity").and_then(|v| v.as_f64()).unwrap_or(0.0)
            .partial_cmp(&a.get("similarity").and_then(|v| v.as_f64()).unwrap_or(0.0))
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    let filtered: Vec<_> = sorted
        .into_iter()
        .filter(|v| v.get("similarity").and_then(|x| x.as_f64()).unwrap_or(0.0) >= min_score as f64)
        .take((search_limit * 3) as usize) // 文書結合時はチャンク多めに取得
        .collect();

    let final_items: Vec<serde_json::Value> = if group_by_document && !filtered.is_empty() {
        // 文書単位にまとめる: document_id ごとにチャンクを結合
        use std::collections::BTreeMap;
        let mut by_doc: BTreeMap<i64, (f64, Option<String>, Option<String>, String)> = BTreeMap::new();
        for it in &filtered {
            let doc_id = it.get("document_id").and_then(|v| v.as_i64()).unwrap_or(0);
            let sim = it.get("similarity").and_then(|v| v.as_f64()).unwrap_or(0.0);
            let path = it.get("path").and_then(|v| v.as_str()).map(String::from);
            let mime = it.get("mime").and_then(|v| v.as_str()).map(String::from);
            let cat = it.get("category").and_then(|v| v.as_str()).unwrap_or("").to_string();
            let entry = by_doc.entry(doc_id).or_insert((0.0f64, path, mime, cat));
            if sim > entry.0 {
                entry.0 = sim;
            }
        }
        let mut doc_list: Vec<(i64, f64, Option<String>, Option<String>, String)> = by_doc
            .into_iter()
            .map(|(doc_id, (sim, path, mime, cat))| (doc_id, sim, path, mime, cat))
            .collect();
        doc_list.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        let mut out = Vec::with_capacity(doc_list.len().min(search_limit as usize));
        for (doc_id, best_sim, path, mime, category) in doc_list.into_iter().take(search_limit as usize) {
            let chunks: Vec<String> = sqlx::query_scalar::<_, String>(
                "SELECT content FROM items WHERE document_id = ? ORDER BY chunk_index"
            )
            .bind(doc_id)
            .fetch_all(&state.db_pool)
            .await
            .unwrap_or_default();
            let content = chunks.join("\n\n");
            out.push(serde_json::json!({
                "id": doc_id,
                "document_id": doc_id,
                "path": path,
                "mime": mime,
                "category": category,
                "content": content,
                "similarity": best_sim
            }));
        }
        out
    } else {
        filtered.into_iter().take(search_limit as usize).collect()
    };
    let result_count = final_items.len();
    log::info!("[search] result_count={} ok={} vector_search_used={}", result_count, result_count > 0, vector_search_used);

    let response_payload = serde_json::json!({
        "items": final_items,
        "vector_search_used": vector_search_used
    });
    let result_text = serde_json::to_string_pretty(&response_payload).unwrap_or_else(|_| "{}".to_string());

    Some(serde_json::json!({ 
        "content": [{ "type": "text", "text": result_text }] 
    }))
}
