pub mod types;
pub mod handlers;
pub mod system;
pub mod tools;
pub mod watch;

pub use types::AppState;
use axum::{
    routing::{get, post},
    Router,
};
use std::sync::atomic::{AtomicBool, AtomicU64};
use std::sync::Arc;
use tokio::sync::{broadcast, RwLock};
use tower_http::cors::{Any, CorsLayer};
use std::collections::HashMap;

pub fn create_mcp_app(state: AppState) -> Router {
    let cors = CorsLayer::new()
        .allow_origin(Any)
        .allow_methods(Any)
        .allow_headers(Any);

    Router::new()
        .route("/sse", get(handlers::sse_handler))
        .route("/messages", post(mcp_messages_handler))
        .route("/edition", get(handlers::edition_handler))
        .route("/doc_count", get(handlers::doc_count_handler))
        .route("/model_name", get(handlers::model_name_handler))
        .route("/indexing_status", get(handlers::indexing_status_handler))
        .route("/heal", get(handlers::heal_handler))
        .route("/version", get(handlers::version_handler))
        .route("/settings", get(handlers::settings_get_handler).post(handlers::settings_post_handler))
        .layer(cors)
        .with_state(state)
}

/// MCP 用の AppState を構築する。ルーティング・listen とは分離し、テストや起動オーケストレーションで再利用する。
pub fn build_app_state(
    app_data_dir: std::path::PathBuf,
    db_pool: sqlx::SqlitePool,
    model_name: String,
    edition: String,
    #[cfg(feature = "pro")] embedding_model_dir: Option<std::path::PathBuf>,
    watcher_restart_tx: Option<crate::mcp::types::WatcherRestartSender>,
) -> AppState {
    let (tx, _rx) = broadcast::channel(100);
    let tokenizer = Arc::new(crate::utils::tokenizer::JapaneseTokenizer::new().unwrap());

    #[cfg(feature = "pro")]
    let embedding_model = {
        let had_dir = embedding_model_dir.is_some();
        let model = embedding_model_dir.and_then(|dir| {
            log::info!("[BOOT] Pro: loading embedding model from {:?}...", dir);
            match crate::utils::embedding_pro::ProEmbeddingModel::load(&dir) {
                Ok(m) => {
                    log::info!("[BOOT] Pro: embedding model loaded.");
                    Some(m)
                }
                Err(e) => {
                    log::warn!("[BOOT] Pro: embedding model load failed: {}", e);
                    None
                }
            }
        });
        if model.is_none() && had_dir {
            log::warn!("[BOOT] Pro: running without embedding (vec search disabled until model is available).");
        }
        Arc::new(RwLock::new(model))
    };

    AppState {
        app_data_dir,
        db_pool,
        tx,
        model_name,
        edition,
        sessions: Arc::new(RwLock::new(HashMap::new())),
        tokenizer,
        #[cfg(feature = "community")]
        lsa_model: Arc::new(RwLock::new(None)),
        #[cfg(feature = "pro")]
        embedding_model,
        #[cfg(feature = "pro")]
        pro_hnsw_rebuild_requested: Arc::new(AtomicBool::new(false)),
        hnsw_index: Arc::new(RwLock::new(None)),
        changes_since_train: Arc::new(AtomicU64::new(0)),
        retrain_scheduled: Arc::new(AtomicBool::new(false)),
        indexing_status: Arc::new(RwLock::new("idle".to_string())),
        watch_ingestion_status: Arc::new(std::sync::RwLock::new(String::new())),
        watcher_restart_tx,
    }
}

/// 起動オーケストレーション: 状態構築 → 即 listen → バックグラウンドで LSA/HNSW 構築。
/// GUI はすぐ 3001 に接続でき「LSA学習中」等を表示。検索は準備完了まで FTS のみ等で応答。
pub async fn run_server(
    port: u16,
    app_data_dir: std::path::PathBuf,
    db_pool: sqlx::SqlitePool,
    model_name: String,
    edition: String,
    #[cfg(feature = "pro")] embedding_model_dir: Option<std::path::PathBuf>,
) {
    let mcp_start = std::time::Instant::now();
    log::info!("[BOOT] MCP: run_server started (port={})", port);

    let (watcher_tx, watcher_rx) = std::sync::mpsc::channel();
    #[cfg(feature = "community")]
    let state = build_app_state(app_data_dir.clone(), db_pool, model_name, edition, Some(watcher_tx));
    #[cfg(feature = "pro")]
    let state = build_app_state(app_data_dir.clone(), db_pool, model_name, edition, embedding_model_dir, Some(watcher_tx));

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

    let watcher_config: crate::mcp::types::WatcherConfig = {
        let path = app_data_dir.join("settings.json");
        let (monitor_paths, extensions, category_map) = if path.exists() {
            if let Ok(s) = std::fs::read_to_string(&path) {
                if let Ok(v) = serde_json::from_str::<serde_json::Value>(&s) {
                    let (paths, cat_map) = v
                        .get("monitor_paths")
                        .and_then(|a| a.as_array())
                        .map(|a| crate::mcp::types::parse_monitor_paths(a))
                        .unwrap_or_default();
                    let exts: Vec<String> = v
                        .get("watch_extensions")
                        .and_then(|a| a.as_array())
                        .map(|a| a.iter().filter_map(|v| v.as_str().map(String::from)).collect())
                        .unwrap_or_else(|| {
                            watch::DEFAULT_WATCH_EXTENSIONS
                                .iter()
                                .map(|s| s.to_string())
                                .collect()
                        });
                    (paths, exts, cat_map)
                } else {
                    (vec![], watch::DEFAULT_WATCH_EXTENSIONS.iter().map(|s| s.to_string()).collect(), Default::default())
                }
            } else {
                (vec![], watch::DEFAULT_WATCH_EXTENSIONS.iter().map(|s| s.to_string()).collect(), Default::default())
            }
        } else {
            (vec![], watch::DEFAULT_WATCH_EXTENSIONS.iter().map(|s| s.to_string()).collect(), Default::default())
        };
        crate::mcp::types::WatcherConfig { paths: monitor_paths, extensions, category_map }
    };
    if let Some(ref tx) = state.watcher_restart_tx {
        let _ = tx.send(watcher_config);
    }
    let state_arc = std::sync::Arc::new(state.clone());
    let runtime_handle = tokio::runtime::Handle::current();
    watch::spawn_folder_watcher(runtime_handle, state_arc, watcher_rx);

    let skip_boot_train = std::env::var_os("TELOS_SKIP_BOOT_TRAIN")
        .map(|v| v == "1" || v == "true" || v == "yes")
        .unwrap_or(false);
    if skip_boot_train {
        log::info!("[BOOT] TELOS_SKIP_BOOT_TRAIN=1: skipping LSA/HNSW on startup (use RE-INDEX when needed)");
        *state.indexing_status.write().await = "idle".to_string();
        let _ = state.tx.send("indexing:idle".to_string());
    } else {
        let state_for_train = state.clone();
        tokio::task::spawn(async move {
            system::train_lsa_and_sync_hnsw(state_for_train).await;
            log::info!("[BOOT] MCP: LSA/HNSW background task finished");
        });
    }

    let app = create_mcp_app(state);
    let bind_addr = format!("127.0.0.1:{}", port);
    let listener = tokio::net::TcpListener::bind(&bind_addr)
        .await
        .unwrap();
    log::info!("[BOOT] MCP: Listening on {} ({}ms) — ready for connections", listener.local_addr().unwrap(), mcp_start.elapsed().as_millis());
    axum::serve(listener, app).await.unwrap();
}

// ----------------------------------------------------------------------------
// Main Message Handler (Dispatching to Tools)
// ----------------------------------------------------------------------------
use axum::{
    extract::{Query, State},
    response::{IntoResponse, Response},
    Json,
    body::Bytes,
};
use crate::mcp::types::{JsonRpcRequest, JsonRpcResponse, MessageQuery};

pub async fn mcp_messages_handler(
    State(state): State<AppState>,
    Query(query): Query<MessageQuery>,
    body: Bytes,
) -> Response {
    // 1. Raw body logging for diagnostics (サーバーコンソールに操作時の受信メッセージを出力)
    let body_str = String::from_utf8_lossy(&body);
    log::info!("[server] MCP receive: {}", body_str);

    // 2. Manual JSON parsing to avoid silent extraction errors
    let req: JsonRpcRequest = match serde_json::from_slice(&body) {
        Ok(r) => r,
        Err(e) => {
            log::error!("MCP JSON Deserialization Error: {}. Raw: {}", e, body_str);
            return (
                axum::http::StatusCode::UNPROCESSABLE_ENTITY,
                Json(serde_json::json!({
                    "jsonrpc": "2.0",
                    "error": { "code": -32700, "message": "Parse error" },
                    "id": null
                }))
            ).into_response();
        }
    };

    let method = req.method.as_str();
    let actual_method = if method == "tools/call" {
        req.params
            .as_ref()
            .and_then(|p| p.get("name"))
            .and_then(|n| n.as_str())
            .unwrap_or("unknown")
    } else {
        method
    };

    log::info!("MCP Request: {} (Actual: {}, Session: {:?})", method, actual_method, query.session_id);

    // tools/call をブロードキャストして UI の MCP ACTIVITY に通知
    if method == "tools/call" || matches!(actual_method, "get_item_by_id" | "add_item_text" | "search_text" | "lsa_search" | "update_item" | "delete_item" | "list_documents" | "list_categories" | "get_document_count" | "get_document" | "delete_document" | "lsa_retrain") {
        let _ = state.tx.send(format!("mcp:call:{}", actual_method));
    }

    let result = match method {
        "initialize" => {
            let client_version = req.params.as_ref()
                .and_then(|p| p.as_object())
                .and_then(|p| p.get("protocolVersion"))
                .and_then(|v| v.as_str())
                .unwrap_or("2024-11-05");
            
            log::info!("MCP Handshake: Client requested protocol version {}", client_version);
            
            Some(serde_json::json!({
                "protocolVersion": client_version,
                "capabilities": { 
                    "tools": { "listChanged": false },
                    "resources": { "listChanged": false, "subscribe": false },
                    "prompts": { "listChanged": false },
                    "logging": {}
                },
                "serverInfo": { "name": "TelosDB", "version": "0.3.0" }
            }))
        },
        "notifications/initialized" => {
            log::info!("MCP Client Initialized");
            None
        }
        "resources/list" => Some(serde_json::json!({ "resources": [] })),
        "prompts/list" => Some(serde_json::json!({ "prompts": [] })),
        "tools/list" => Some(serde_json::json!({ "tools": tools::registry::tool_list() })),
        "tools/call" | "get_item_by_id" | "add_item_text" | "search_text" | "lsa_search" | "update_item" | "delete_item" | "list_documents" | "list_categories" | "get_document_count" | "get_document" | "delete_document" | "lsa_retrain" => {
            let empty_map = serde_json::Map::new();
            let mut args = req.params.as_ref().and_then(|p| p.as_object()).unwrap_or(&empty_map);
            
            // tools/call の場合は、実際の引数は "arguments" フィールドにある
            if method == "tools/call" {
                if let Some(inner_args) = args.get("arguments").and_then(|a| a.as_object()) {
                    args = inner_args;
                }
            }
            
            tools::dispatch_tool(&state, method, actual_method, args).await
        }
        _ => Some(serde_json::json!({
            "content": [{ "type": "text", "text": format!("Method not implemented: {}", method) }],
            "isError": true
        })),
    };

    // Notifications (id == null) MUST NOT receive a response
    if req.id.is_none() || req.id.as_ref().is_some_and(|v| v.is_null()) {
        log::info!("MCP Notification received: {} (No response sent)", method);
        return axum::http::StatusCode::NO_CONTENT.into_response();
    }

    if let Some(id_val) = req.id {
        let resp = JsonRpcResponse {
            jsonrpc: "2.0",
            result,
            error: None,
            id: Some(id_val),
        };

        if let Some(sid) = query.session_id {
            // MCP SPEC over SSE: "The server SHOULD send responses via the SSE stream... if a session ID is provided"
            // オリジナルの動作に戻し、SSE経由で配送する。
            let resp_str = serde_json::to_string(&resp).unwrap();
            log::info!("Sending MCP Response via SSE (Session: {}, ID: {:?})", sid, resp.id);
            let sessions = state.sessions.read().await;
            if let Some(tx) = sessions.get(&sid) {
                let _ = tx.send(resp_str);
            }
            axum::http::StatusCode::ACCEPTED.into_response()
        } else {
            // Direct request (not via SSE)
            log::info!("Sending Direct MCP Response (ID: {:?})", resp.id);
            Json(resp).into_response()
        }
    } else {
        axum::http::StatusCode::NO_CONTENT.into_response()
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn test_text_chunking_logic() {
        // 800文字ずつの分割を確認する
        let chunk_size = 800;
        
        // 1. ちょうど 800 文字
        let text_800 = "a".repeat(800);
        let chunks_800: Vec<String> = text_800.chars()
            .collect::<Vec<char>>()
            .chunks(chunk_size)
            .map(|c| c.iter().collect())
            .collect();
        assert_eq!(chunks_800.len(), 1);
        assert_eq!(chunks_800[0].len(), 800);

        // 2. 801 文字 (2 チャンク)
        let text_801 = "a".repeat(801);
        let chunks_801: Vec<String> = text_801.chars()
            .collect::<Vec<char>>()
            .chunks(chunk_size)
            .map(|c| c.iter().collect())
            .collect();
        assert_eq!(chunks_801.len(), 2);
        assert_eq!(chunks_801[0].len(), 800);
        assert_eq!(chunks_801[1].len(), 1);

        // 3. 1600 文字 (2 チャンク)
        let text_1600 = "a".repeat(1600);
        let chunks_1600: Vec<String> = text_1600.chars()
            .collect::<Vec<char>>()
            .chunks(chunk_size)
            .map(|c| c.iter().collect())
            .collect();
        assert_eq!(chunks_1600.len(), 2);
        
        // 4. 空文字列
        let text_empty = "";
        let chunks_empty: Vec<String> = text_empty.chars()
            .collect::<Vec<char>>()
            .chunks(chunk_size)
            .map(|c| c.iter().collect())
            .collect();
        assert_eq!(chunks_empty.len(), 0);
    }
}
