use std::sync::atomic::Ordering;
use crate::mcp::types::AppState;
#[cfg(feature = "community")]
use sqlx::Row;
#[cfg(feature = "community")]
use std::sync::Arc;

// Pro 版: 既存 items を vec_items にバックフィルしてから HNSW を構築する
#[cfg(feature = "pro")]
pub async fn train_lsa_and_sync_hnsw(state: AppState) {
    use hnsw_rs::prelude::*;
    use sqlx::Row;
    let dim = crate::utils::embedding_pro::ProEmbeddingModel::embed_dim();

    {
        let mut st = state.indexing_status.write().await;
        *st = "syncing".to_string();
    }
    let _ = state.tx.send("indexing:syncing".to_string());

    // vec_items に無い items だけを埋め込みでバックフィルする（全件ではない）
    let guard = state.embedding_model.read().await;
    let mut data_to_insert: Vec<(Vec<f32>, usize)> = Vec::new();
    if let Some(ref model) = *guard {
        let items_missing: Vec<(i64, String)> = match sqlx::query(
            "SELECT i.id, i.content FROM items i
             LEFT JOIN vec_items v ON i.id = v.id
             WHERE v.id IS NULL AND i.content IS NOT NULL AND trim(i.content) != ''"
        )
        .fetch_all(&state.db_pool)
        .await
        {
            Ok(rows) => rows.iter().map(|r| (r.get::<i64, _>(0), r.get::<String, _>(1))).collect(),
            Err(e) => {
                log::warn!("[BOOT] Pro backfill: failed to list items: {}", e);
                Vec::new()
            }
        };
        if !items_missing.is_empty() {
            log::info!("[BOOT] Pro backfill: embedding {} items (missing from vec_items)...", items_missing.len());
            let mut encode_fail_count = 0usize;
            let mut first_encode_err: Option<String> = None;
            for (id, content) in items_missing {
                match model.encode(&content) {
                    Ok(embedding) if embedding.len() == dim => {
                        data_to_insert.push((embedding, id as usize));
                    }
                    Ok(_) => {
                        encode_fail_count += 1;
                        if first_encode_err.is_none() {
                            first_encode_err = Some(format!("id={} dimension mismatch", id));
                        }
                    }
                    Err(e) => {
                        encode_fail_count += 1;
                        if first_encode_err.is_none() {
                            first_encode_err = Some(format!("id={} {}", id, e));
                        }
                    }
                }
            }
            if encode_fail_count > 0 {
                log::warn!("[BOOT] Pro backfill: {} encode failures (first: {:?})", encode_fail_count, first_encode_err);
            }
            log::info!("[BOOT] Pro backfill: encoded {} vectors.", data_to_insert.len());
            // 不足分だけ vec_items に INSERT（既存は触らない）
            if let Ok(mut tx) = state.db_pool.begin().await {
                let mut insert_ok = 0usize;
                let mut insert_err = 0usize;
                for (vec, id) in &data_to_insert {
                    let emb_json = serde_json::to_string(vec).unwrap_or_else(|_| "[]".to_string());
                    match sqlx::query("INSERT INTO vec_items (id, embedding) VALUES (?, ?)")
                        .bind(*id as i64)
                        .bind(&emb_json)
                        .execute(&mut *tx)
                        .await
                    {
                        Ok(r) => { insert_ok += r.rows_affected() as usize; }
                        Err(e) => {
                            insert_err += 1;
                            log::warn!("[BOOT] Pro backfill: insert vec_items id={} failed: {}", id, e);
                        }
                    }
                }
                // 診断: トランザクション内で INSERT 直後の件数
                let count_in_tx: Option<i64> = sqlx::query_scalar("SELECT COUNT(*) FROM vec_items")
                    .fetch_one(&mut *tx)
                    .await
                    .ok();
                log::info!("[BOOT] Pro backfill vec_items: insert_ok={} insert_err={} COUNT(*) within tx={:?}", insert_ok, insert_err, count_in_tx);
                if let Err(e) = tx.commit().await {
                    log::warn!("[BOOT] Pro backfill: vec_items commit failed: {}", e);
                }
            }
        }
    }
    drop(guard);

    // 診断: COMMIT 後に別コネクションで vec_items 件数を確認（永続化の有無）
    if let Ok((count_after,)) = sqlx::query_as::<_, (i64,)>("SELECT COUNT(*) FROM vec_items")
        .fetch_one(&state.db_pool)
        .await
    {
        log::info!("[BOOT] Pro vec_items COUNT(*) after commit (pool): {}", count_after);
    }

    // HNSW は vec_items 全体から構築（既存＋今回バックフィルした分）
    let from_db: Vec<(Vec<f32>, usize)> = if let Ok(rows) = sqlx::query("SELECT id, vec_to_json(embedding) FROM vec_items")
        .fetch_all(&state.db_pool)
        .await
    {
        let mut out = Vec::new();
        for row in rows {
            let id: i64 = row.get(0);
            let embedding_str: String = row.get(1);
            if let Ok(vec) = serde_json::from_str::<Vec<f32>>(&embedding_str) {
                if vec.len() == dim {
                    out.push((vec, id as usize));
                }
            }
        }
        out
    } else {
        Vec::new()
    };

    if !from_db.is_empty() {
        log::info!("[BOOT] Pro HNSW: building index for {} vectors (from vec_items)...", from_db.len());
        let hnsw = tokio::task::spawn_blocking(move || {
            let hnsw: Hnsw<'static, f32, DistCosine> = Hnsw::new(16, from_db.len().max(100), 16, 200, DistCosine {});
            let refs: Vec<(&Vec<f32>, usize)> = from_db.iter().map(|(v, id)| (v, *id)).collect();
            hnsw.parallel_insert(&refs);
            log::info!("[BOOT] Pro HNSW: index built.");
            hnsw
        }).await.ok();
        if let Some(hnsw) = hnsw {
            *state.hnsw_index.write().await = Some(hnsw);
        }
    } else if !data_to_insert.is_empty() {
        // vec_items が空（永続化されていない等）だが今回エンコードした分はある → メモリ上のデータで HNSW 構築
        log::info!("[BOOT] Pro HNSW: vec_items empty, building from {} in-memory vectors...", data_to_insert.len());
        let hnsw = tokio::task::spawn_blocking(move || {
            let hnsw: Hnsw<'static, f32, DistCosine> = Hnsw::new(16, data_to_insert.len().max(100), 16, 200, DistCosine {});
            let refs: Vec<(&Vec<f32>, usize)> = data_to_insert.iter().map(|(v, id)| (v, *id)).collect();
            hnsw.parallel_insert(&refs);
            hnsw
        }).await.ok();
        if let Some(hnsw) = hnsw {
            *state.hnsw_index.write().await = Some(hnsw);
        }
    } else {
        log::info!("[BOOT] Pro HNSW: no vectors (add documents and run RE-INDEX).");
    }
    *state.indexing_status.write().await = "idle".to_string();
    let _ = state.tx.send("indexing:idle".to_string());
}

#[cfg(feature = "community")]
use std::collections::HashMap;
#[cfg(feature = "community")]
use hnsw_rs::prelude::*;
#[cfg(feature = "community")]
use crate::utils::lsa::LsaModel;

#[cfg(feature = "community")]
pub async fn train_lsa_and_sync_hnsw(state: AppState) {
    let boot_start = std::time::Instant::now();
    log::info!("[BOOT] LSA: starting (fetching documents)");
    {
        let mut st = state.indexing_status.write().await;
        *st = "training".to_string();
    }
    let _ = state.tx.send("indexing:training".to_string());

    let rows = match sqlx::query("SELECT content FROM items").fetch_all(&state.db_pool).await {
        Ok(rows) if !rows.is_empty() => {
            log::info!("[BOOT] LSA: fetched {} documents ({}ms)", rows.len(), boot_start.elapsed().as_millis());
            rows
        }
        _ => {
            log::info!("[BOOT] LSA: no documents, skip ({}ms)", boot_start.elapsed().as_millis());
            *state.indexing_status.write().await = "idle".to_string();
            let _ = state.tx.send("indexing:idle".to_string());
            return;
        }
    };

    log::info!("[BOOT] LSA: training in blocking thread...");
    let state_inner = state.clone();
    let join_result = tokio::task::spawn_blocking(move || {
        let mut builder = crate::utils::lsa::TermDocumentMatrixBuilder::new();
        for row in rows {
            let content: String = row.get(0);
            let tokens = state_inner.tokenizer.tokenize_to_vec(&content).unwrap_or_default();
            builder.add_document(tokens);
        }
        let (matrix, idfs) = builder.build_matrix();
        LsaModel::train(&matrix, builder.vocabulary, idfs, 50).map(|m| (m, builder.counts.len()))
    }).await;

    let result = match join_result {
        Ok(r) => r,
        Err(e) => {
            log::error!("[BOOT] LSA: training task panicked: {} ({}ms)", e, boot_start.elapsed().as_millis());
            *state.indexing_status.write().await = "idle".to_string();
            let _ = state.tx.send("indexing:idle".to_string());
            return;
        }
    };

    match result {
        Ok((model, doc_count)) => {
            let model_arc = Arc::new(model);
            {
                let mut lsa = state.lsa_model.write().await;
                *lsa = Some((*model_arc).clone());
            }
            log::info!("[BOOT] LSA: trained {} documents ({}ms)", doc_count, boot_start.elapsed().as_millis());

            log::info!("[BOOT] HNSW: building index...");
            let hnsw: Hnsw<'static, f32, DistCosine> = Hnsw::new(16, doc_count.max(100), 16, 200, DistCosine {});

            {
                let mut st = state.indexing_status.write().await;
                *st = "syncing".to_string();
            }
            let _ = state.tx.send("indexing:syncing".to_string());
            sync_all_vectors(state.clone(), Some(hnsw)).await;
            *state.indexing_status.write().await = "idle".to_string();
            let _ = state.tx.send("indexing:idle".to_string());
            log::info!("[BOOT] LSA+HNSW: all done ({}ms)", boot_start.elapsed().as_millis());
        }
        Err(e) => {
            log::error!("[BOOT] LSA: training failed: {} ({}ms)", e, boot_start.elapsed().as_millis());
            *state.indexing_status.write().await = "idle".to_string();
            let _ = state.tx.send("indexing:idle".to_string());
        }
    }
}

/// DB 内の全アイテムをチェックし、ベクトルが欠落または異常（全て0）なものを補完する（Community 版）
#[cfg(feature = "community")]
pub async fn sync_all_vectors(state: AppState, startup_hnsw: Option<Hnsw<'static, f32, DistCosine>>) {
    log::info!("[BOOT] HNSW: checking vectors in vec_items...");
    
    let rows = match sqlx::query(
        "SELECT i.id, i.content, 
                CASE WHEN v.embedding IS NOT NULL THEN vec_to_json(v.embedding) ELSE NULL END
         FROM items i 
         LEFT JOIN vec_items v ON i.id = v.id"
    )
    .fetch_all(&state.db_pool)
    .await {
        Ok(rows) => rows,
        Err(e) => {
            log::error!("Failed to fetch items for sync: {}", e);
            return;
        }
    };

    let mut to_sync = Vec::new();
    for row in rows {
        let id: i64 = row.get(0);
        let content: String = row.get(1);
        let embedding_str: Option<String> = row.get(2);

        let needs_sync = if let Some(s) = embedding_str {
            if let Ok(vec) = serde_json::from_str::<Vec<f32>>(&s) {
                // すべて 0.0 なら異常（ダミー）とみなす
                vec.iter().all(|&x| x == 0.0)
            } else {
                true // パース失敗も異常
            }
        } else {
            true // 不在
        };

        if needs_sync {
            to_sync.push((id, content));
        }
    }

    if !to_sync.is_empty() {
        log::info!("[BOOT] HNSW: {} items need vector update, processing...", to_sync.len());
        
        let lsa_guard = state.lsa_model.read().await;
        if let Some(model) = lsa_guard.as_ref() {
            let model_inner = model.clone();
            let state_inner = state.clone();
            
            let synced_data = match tokio::task::spawn_blocking(move || {
                let mut results = Vec::new();
                for (id, content) in to_sync {
                    let mut query_counts = HashMap::new();
                    let tokens = state_inner.tokenizer.tokenize_to_vec(&content).unwrap_or_default();
                    for token in tokens {
                        if let Some(&tid) = model_inner.vocabulary.get(&token) {
                            *query_counts.entry(tid).or_insert(0.0) += 1.0;
                        }
                    }
                    let mut query_vec = ndarray::Array1::zeros(model_inner.vocabulary.len());
                    for (tid, count) in query_counts {
                        query_vec[tid] = count;
                    }

                    if let Ok(projected) = model_inner.project_query(&query_vec) {
                        let mut proj_f32: Vec<f32> = projected.iter().map(|&x| x as f32).collect();
                        if proj_f32.len() < 50 { proj_f32.resize(50, 0.0); } else { proj_f32.truncate(50); }
                        results.push((id, proj_f32));
                    }
                }
                results
            }).await {
                Ok(data) => data,
                Err(e) => {
                    log::error!("[BOOT] HNSW: sync vector computation task panicked: {}", e);
                    return;
                }
            };

            let mut count = 0;
            for (id, proj_f32) in synced_data {
                let mut tx = match state.db_pool.begin().await {
                    Ok(t) => t,
                    Err(_) => continue,
                };

                let _ = sqlx::query("DELETE FROM vec_items WHERE id = ?").bind(id).execute(&mut *tx).await;
                let _ = sqlx::query("INSERT INTO vec_items (id, embedding) VALUES (?, ?)")
                    .bind(id)
                    .bind(serde_json::to_string(&proj_f32).unwrap_or("[]".to_string()))
                    .execute(&mut *tx)
                    .await;

                let vector_blob = bincode::serialize(&proj_f32).unwrap_or_default();
                let _ = sqlx::query("INSERT OR REPLACE INTO items_lsa (id, vector) VALUES (?, ?)")
                    .bind(id)
                    .bind(vector_blob)
                    .execute(&mut *tx)
                    .await;

                if tx.commit().await.is_ok() {
                    count += 1;
                }
            }
            log::info!("[BOOT] HNSW: synchronized {} vectors.", count);
        } else {
            log::warn!("[BOOT] HNSW: LSA model not available for sync.");
        }
    } else {
        log::info!("[BOOT] HNSW: all vectors healthy.");
    }

    if let Some(hnsw) = startup_hnsw {
        log::info!("[BOOT] HNSW: populating index from database...");
        let db = state.db_pool.clone();
        
        tokio::task::spawn_blocking(move || {
            // I/O は blocking thread 内で行なう (sqlx は本来 async 全体で使いたいが、
            // parallel_insert が CPU を占有する間、他が動けないため、このスレッド内で同期的に回すか、
            // 先に async で取ってきたデータを move する)
            // ここでは I/O を async で済ませてデータを move し、計算だけを blocking thread でやるのが美しい。
            (hnsw, db)
        }).await.map(|(hnsw, _db)| {
            let state_inner = state.clone();
            tokio::spawn(async move {
                if let Ok(rows) = sqlx::query("SELECT id, vec_to_json(embedding) FROM vec_items").fetch_all(&state_inner.db_pool).await {
                    let mut data_to_insert = Vec::new();
                    for row in rows {
                        let id: i64 = row.get(0);
                        let embedding_str: String = row.get(1);
                        if let Ok(vec) = serde_json::from_str::<Vec<f32>>(&embedding_str) {
                            if vec.len() == 50 {
                                data_to_insert.push((vec, id as usize));
                            }
                        }
                    }
                    
                    if !data_to_insert.is_empty() {
                        log::info!("[BOOT] HNSW: inserting {} items...", data_to_insert.len());
                        tokio::task::spawn_blocking(move || {
                             let refs: Vec<(&Vec<f32>, usize)> = data_to_insert.iter().map(|(v, id)| (v, *id)).collect();
                             hnsw.parallel_insert(&refs);
                             hnsw
                        }).await.map(|hnsw_ready| {
                            tokio::spawn(async move {
                                let mut idx = state_inner.hnsw_index.write().await;
                                *idx = Some(hnsw_ready);
                                log::info!("[BOOT] HNSW: index ready.");
                            });
                        }).unwrap_or_else(|e| log::error!("HNSW parallel insert panicked: {}", e));
                    } else {
                        let mut idx = state_inner.hnsw_index.write().await;
                        *idx = Some(hnsw);
                        log::info!("[BOOT] HNSW: index ready (no items).");
                    }
                }
            });
        }).unwrap();
    }
}

/// 追加・更新・削除が「登録件数に対する割合」で閾値を超えたら、デバウンス後にバックグラウンドで LSA 再学習を実行する。（Community では LSA、Pro では no-op）
/// 二重実行防止のため、1本だけスケジュールされる。
pub async fn schedule_retrain_if_needed(state: &AppState) {
    const CHANGE_RATIO: f64 = 0.2; // 登録ドキュメント数の 20%
    const MAX_THRESHOLD: u64 = 5;  // 多くても5件で再学習を検討（数が多いときに学習しなくならないように）
    const MIN_THRESHOLD: u64 = 1;
    const DEBOUNCE_SECS: u64 = 90;

    let doc_count: i64 = crate::db::get_document_count(&state.db_pool)
        .await
        .unwrap_or(0);
    let by_ratio = (doc_count as f64 * CHANGE_RATIO).ceil() as u64;
    let threshold = by_ratio.min(MAX_THRESHOLD).max(MIN_THRESHOLD);

    let prev = state.changes_since_train.fetch_add(1, Ordering::Relaxed);
    let count = prev + 1;
    if count < threshold {
        return;
    }
    if state
        .retrain_scheduled
        .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
        .is_err()
    {
        return;
    }
    let state_clone = state.clone();
    tokio::spawn(async move {
        tokio::time::sleep(std::time::Duration::from_secs(DEBOUNCE_SECS)).await;
        log::info!(
            "Auto LSA retrain starting (debounced {}s after {} changes, threshold was {}).",
            DEBOUNCE_SECS, count, threshold
        );
        train_lsa_and_sync_hnsw(state_clone.clone()).await;
        state_clone.changes_since_train.store(0, Ordering::Relaxed);
        state_clone.retrain_scheduled.store(false, Ordering::Relaxed);
        log::info!("Auto LSA retrain completed.");
    });
}
