use crate::mcp::types::AppState;

#[cfg(feature = "pro")]
pub async fn handle_lsa_retrain(state: &AppState) -> Option<serde_json::Value> {
    use hnsw_rs::prelude::*;
    use sqlx::Row;

    {
        let guard = state.embedding_model.read().await;
        if guard.is_none() {
            return Some(serde_json::json!({
                "content": [{ "type": "text", "text": "Pro: 埋め込みモデルが読み込まれていません。RE-INDEX を実行できません。" }],
                "isError": true
            }));
        }
    }

    let state_clone = state.clone();
    tokio::spawn(async move {
        let dim = crate::utils::embedding_pro::ProEmbeddingModel::embed_dim();

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

        let rows = match sqlx::query("SELECT id, content FROM items WHERE content IS NOT NULL AND trim(content) != ''")
            .fetch_all(&state_clone.db_pool)
            .await
        {
            Ok(r) if !r.is_empty() => r,
            Ok(_) => {
                log::info!("Pro RE-INDEX: no documents, skip.");
                *state_clone.indexing_status.write().await = "idle".to_string();
                let _ = state_clone.tx.send("indexing:idle".to_string());
                return;
            }
            Err(e) => {
                log::error!("Pro RE-INDEX: failed to fetch items: {}", e);
                *state_clone.indexing_status.write().await = "idle".to_string();
                let _ = state_clone.tx.send("indexing:idle".to_string());
                return;
            }
        };

        log::info!("Pro RE-INDEX: embedding {} items...", rows.len());

        let mut data_to_insert: Vec<(Vec<f32>, usize)> = Vec::new();
        let mut tx = match state_clone.db_pool.begin().await {
            Ok(t) => t,
            Err(e) => {
                log::error!("Pro RE-INDEX: begin tx failed: {}", e);
                *state_clone.indexing_status.write().await = "idle".to_string();
                let _ = state_clone.tx.send("indexing:idle".to_string());
                return;
            }
        };

        sqlx::query("DELETE FROM vec_items").execute(&mut *tx).await.ok();
        sqlx::query("DELETE FROM items_fts").execute(&mut *tx).await.ok();

        let mut vec_insert_ok = 0usize;
        let mut vec_insert_err = 0usize;
        {
            let model_guard = state_clone.embedding_model.read().await;
            let model = match model_guard.as_ref() {
                Some(m) => m,
                None => {
                    log::error!("Pro RE-INDEX: embedding model not loaded.");
                    *state_clone.indexing_status.write().await = "idle".to_string();
                    let _ = state_clone.tx.send("indexing:idle".to_string());
                    return;
                }
            };
            for row in &rows {
                let id: i64 = row.get(0);
                let content: String = row.get(1);
                if let Ok(embedding) = model.encode(&content) {
                    if embedding.len() == dim {
                        data_to_insert.push((embedding.clone(), id as usize));
                    }
                    let emb_json = serde_json::to_string(&embedding).unwrap_or_else(|_| "[]".to_string());
                    match sqlx::query("INSERT INTO vec_items (id, embedding) VALUES (?, ?)")
                        .bind(id)
                        .bind(&emb_json)
                        .execute(&mut *tx)
                        .await
                    {
                        Ok(r) => { vec_insert_ok += r.rows_affected() as usize; }
                        Err(e) => {
                            vec_insert_err += 1;
                            log::warn!("Pro RE-INDEX: insert vec_items id={} failed: {}", id, e);
                        }
                    }
                    let _ = sqlx::query("INSERT INTO items_fts (rowid, content) VALUES (?, ?)")
                        .bind(id)
                        .bind(content)
                        .execute(&mut *tx)
                        .await;
                }
            }
        }
        // 診断: トランザクション内での vec_items 件数
        let count_in_tx: Result<i64, _> = sqlx::query_scalar("SELECT COUNT(*) FROM vec_items")
            .fetch_one(&mut *tx)
            .await;
        log::info!("Pro RE-INDEX vec_items: insert_ok={} insert_err={} COUNT(*) within tx={:?}", vec_insert_ok, vec_insert_err, count_in_tx);

        if let Err(e) = tx.commit().await {
            log::error!("Pro RE-INDEX: commit failed: {}", e);
            *state_clone.indexing_status.write().await = "idle".to_string();
            let _ = state_clone.tx.send("indexing:idle".to_string());
            return;
        }
        // 診断: COMMIT 後の vec_items 件数（永続化の有無）
        if let Ok((c,)) = sqlx::query_as::<_, (i64,)>("SELECT COUNT(*) FROM vec_items")
            .fetch_one(&state_clone.db_pool)
            .await
        {
            log::info!("Pro RE-INDEX vec_items COUNT(*) after commit (pool): {}", c);
        }

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

        if data_to_insert.is_empty() {
            log::info!("Pro RE-INDEX: no vectors encoded.");
            *state_clone.hnsw_index.write().await = None;
        } else {
            log::info!("Pro RE-INDEX: building HNSW for {} vectors (in-memory)...", 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_clone.hnsw_index.write().await = Some(hnsw);
                log::info!("Pro RE-INDEX: HNSW built.");
            }
        }

        *state_clone.indexing_status.write().await = "idle".to_string();
        let _ = state_clone.tx.send("indexing:idle".to_string());
        log::info!("Pro RE-INDEX: completed.");
    });

    Some(serde_json::json!({
        "content": [{ "type": "text", "text": "Pro: 埋め込みモデルで全件再計算を開始しました。完了までしばらくお待ちください。" }]
    }))
}

#[cfg(feature = "community")]
use sqlx::Row;
#[cfg(feature = "community")]
use hnsw_rs::prelude::*;
#[cfg(feature = "community")]
use crate::mcp::system::sync_all_vectors;

#[cfg(feature = "community")]
pub async fn handle_lsa_retrain(
    state: &AppState,
) -> Option<serde_json::Value> {
    log::info!("Manual LSA retrain triggered.");
    let state_clone = state.clone();
    tokio::spawn(async move {
        if let Ok(rows) = sqlx::query("SELECT id, content FROM items").fetch_all(&state_clone.db_pool).await {
            if !rows.is_empty() {
                 let mut builder = crate::utils::lsa::TermDocumentMatrixBuilder::new();
                 let mut doc_records = Vec::new();
                 for row in rows {
                     let id: i64 = row.get(0);
                     let content: String = row.get(1);
                     let tokens = state_clone.tokenizer.tokenize_to_vec(&content).unwrap_or_default();
                     builder.add_document(tokens);
                     doc_records.push((id, content));
                 }
                 let (matrix, idfs) = builder.build_matrix();
                 match crate::utils::lsa::LsaModel::train(&matrix, builder.vocabulary, idfs, 50) {
                     Ok(model) => {
                         let mut tx = state_clone.db_pool.begin().await.unwrap();
                         sqlx::query("DELETE FROM items_lsa").execute(&mut *tx).await.unwrap();
                         sqlx::query("DELETE FROM vec_items").execute(&mut *tx).await.unwrap();
                         sqlx::query("DELETE FROM items_fts").execute(&mut *tx).await.unwrap();
                         
                         for (i, (id, content)) in doc_records.iter().enumerate() {
                             let mut doc_tf = ndarray::Array1::zeros(model.vocabulary.len());
                             for (&tid, &count) in &builder.counts[i] {
                                 doc_tf[tid] = count;
                             }
                             if let Ok(projected) = model.project_query(&doc_tf) {
                                 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); }
                                 
                                 let vector_blob = bincode::serialize(&proj_f32).unwrap_or_default();
                                 sqlx::query("INSERT INTO items_lsa (id, vector) VALUES (?, ?)")
                                     .bind(*id)
                                     .bind(vector_blob)
                                     .execute(&mut *tx)
                                     .await
                                     .unwrap();

                                 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
                                     .unwrap();

                                 sqlx::query("INSERT INTO items_fts (rowid, content) VALUES (?, ?)")
                                     .bind(*id)
                                     .bind(content)
                                     .execute(&mut *tx)
                                     .await
                                     .unwrap();
                             }
                         }
                         tx.commit().await.unwrap();

                         let mut lsa = state_clone.lsa_model.write().await;
                         *lsa = Some(model);

                         // HNSW インデックスの再構築
                         let ids: Vec<i64> = doc_records.iter().map(|(id, _)| *id).collect();
                         let hnsw: Hnsw<f32, DistCosine> = Hnsw::new(16, ids.len().max(100), 16, 200, DistCosine {});
                         log::info!("Manual LSA retrain completed successfully. Rebuilding HNSW...");
                         sync_all_vectors(state_clone.clone(), Some(hnsw)).await;
                     }
                    Err(e) => log::error!("Manual LSA training failed: {}", e),
                }
            }
        }
    });
    Some(serde_json::json!({ "content": [{ "type": "text", "text": "LSA retrain started in background." }] }))
}
