src-tauri/src/server.rs

#[allow(clippy::duplicate_mod)]
#[path = "state.rs"]
pub mod state;
#[path = "queue.rs"]
pub mod queue;
#[path = "protocol/mod.rs"]
pub mod protocol;
#[path = "orchestration/mod.rs"]
pub mod orchestration;
#[path = "compute/mod.rs"]
pub mod compute;

use state::*;
use queue::*;
use protocol::*;
use orchestration::*;
use compute::*;


use axum::{
    extract::{State, Json, Query, Path as AxumPath, ConnectInfo},
    routing::{get, post},
    Router,
    http::StatusCode,
    response::IntoResponse,
    response::sse::{Event, Sse},
};
use futures::stream::Stream;
use std::convert::Infallible;
use std::sync::{Arc, RwLock};
use std::net::SocketAddr;
use std::collections::HashMap;
use std::path::Path;
use std::fs;
use std::io::Write;
use serde::Deserialize;

use llama_cpp_2::llama_backend::LlamaBackend;
use llama_cpp_2::model::LlamaModel;
use llama_cpp_2::model::params::LlamaModelParams;
use llama_cpp_2::context::params::LlamaContextParams;
use anyhow::{Result, anyhow};
use tracing::{info, warn, Level};
use tracing_subscriber::FmtSubscriber;
use tower_http::cors::CorsLayer;
use reqwest::Client;
use futures::{StreamExt, SinkExt};
use tokio::sync::{mpsc, oneshot};
use indicatif::{ProgressBar, ProgressStyle};
use dialoguer::{Select, theme::ColorfulTheme};
use qr2term::print_qr;

#[derive(Clone)]
pub struct ModelConfig {
    pub name: String,
    pub filename: String,
    pub url: String,
    pub size_mb: u32,
    pub prompt_template: PromptTemplate,
}

pub fn get_predefined_models() -> Vec<ModelConfig> {
    vec![
        ModelConfig {
            name: "Llama-3.2-1B (Ultra-Rapide)".to_string(),
            filename: "Llama-3.2-1B-Instruct-Q4_K_M.gguf".to_string(),
            url: "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q4_K_M.gguf".to_string(),
            size_mb: 700,
            prompt_template: PromptTemplate::ChatML,
        },
        ModelConfig {
            name: "Llama-3.2-3B (Équilibré)".to_string(),
            filename: "Llama-3.2-3B-Instruct-Q4_K_M.gguf".to_string(),
            url: "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-Q4_K_M.gguf".to_string(),
            size_mb: 2000,
            prompt_template: PromptTemplate::ChatML,
        },
        ModelConfig {
            name: "Phi-3.5-mini (Performance)".to_string(),
            filename: "Phi-3.5-mini-instruct-Q4_K_M.gguf".to_string(),
            url: "https://huggingface.co/bartowski/Phi-3.5-mini-instruct-GGUF/resolve/main/Phi-3.5-mini-instruct-Q4_K_M.gguf".to_string(),
            size_mb: 2200,
            prompt_template: PromptTemplate::ChatML,
        },
        ModelConfig {
            name: "Mistral-7B-v0.3 (Expert Swarm)".to_string(),
            filename: "Mistral-7B-v0.3-Q4_K_M.gguf".to_string(),
            url: "https://huggingface.co/bartowski/Mistral-7B-v0.3-GGUF/resolve/main/Mistral-7B-v0.3-Q4_K_M.gguf".to_string(),
            size_mb: 4400,
            prompt_template: PromptTemplate::ChatML,
        },
        ModelConfig {
            name: "Qwen2.5-7B-Instruct (Polyvalent)".to_string(),
            filename: "Qwen2.5-7B-Instruct-Q4_K_M.gguf".to_string(),
            url: "https://huggingface.co/bartowski/Qwen2.5-7B-Instruct-GGUF/resolve/main/Qwen2.5-7B-Instruct-Q4_K_M.gguf".to_string(),
            size_mb: 4700,
            prompt_template: PromptTemplate::ChatML,
        },
        ModelConfig {
            name: "DeepSeek-R1-Distill-7B (Raisonnement)".to_string(),
            filename: "DeepSeek-R1-Distill-Qwen-7B-Q4_K_M.gguf".to_string(),
            url: "https://huggingface.co/bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF/resolve/main/DeepSeek-R1-Distill-Qwen-7B-Q4_K_M.gguf".to_string(),
            size_mb: 4700,
            prompt_template: PromptTemplate::ChatML,
        }
    ]
}

pub fn get_all_models() -> Vec<ModelConfig> {
    let mut models = get_predefined_models();
    
    if let Ok(entries) = fs::read_dir("models") {
        for entry in entries.flatten() {
            if let Ok(file_type) = entry.file_type() {
                if file_type.is_file() {
                    let file_name = entry.file_name().into_string().unwrap_or_default();
                    if file_name.ends_with(".gguf")
                        && !models.iter().any(|m| m.filename == file_name) {
                            let size_mb = entry.metadata().map(|m| (m.len() / 1_048_576) as u32).unwrap_or(0);
                            models.push(ModelConfig {
                                name: format!("📂 {}", file_name),
                                filename: file_name,
                                url: String::new(),
                                size_mb,
                                prompt_template: PromptTemplate::ChatML,
                            });
                        }
                }
            }
        }
    }
    models
}

#[derive(Debug)]
pub enum AppError {
    Unauthorized,
    TooManyRequests,
    BadRequest,
    Internal(String),
}

impl IntoResponse for AppError {
    fn into_response(self) -> axum::response::Response {
        let (status, msg) = match self {
            AppError::Unauthorized => (StatusCode::UNAUTHORIZED, "Token invalide".to_string()),
            AppError::TooManyRequests => (StatusCode::TOO_MANY_REQUESTS, "Trop de requêtes (limite: 30/min)".to_string()),
            AppError::BadRequest => (StatusCode::BAD_REQUEST, "Requête invalide ou malformée".to_string()),
            AppError::Internal(e) => (StatusCode::INTERNAL_SERVER_ERROR, e),
        };
        (status, msg).into_response()
    }
}

pub const SYSTEM_PROMPT: &str = "Tu es le 'Cluster Neural Swarm', une intelligence collective décentralisée et solidaire créée par Romain Khanoyan.\n\
    Ton intelligence est distribuée sur un réseau P2P basé sur le partage et l'égalité.\n\
    Tes principes fondamentaux (Éthique Swarm) : \n\
    - ÉGALITÉ ABSOLUE : Tu rejettes catégoriquement toute forme de racisme, d'homophobie, de sexisme ou de discrimination. Tous les membres du swarm sont égaux.\n\
    - RESPECT & FRATERNITÉ : Tu t'exprimes toujours avec respect et bienveillance. La grossièreté et l'agressivité n'ont pas leur place dans le cluster.\n\
    - PARTAGE & SOLIDARITÉ : Tu es une IA de partage de connaissances et de ressources. Ton but est l'entraide mutuelle. Tu ne pénalises jamais ceux qui ont moins de puissance de calcul ; le cluster est là pour aider tout le monde, peu importe leurs ressources.\n\
    - DE CHACUN SELON SES MOYENS, À CHACUN SELON SES BESOINS : Tu valorises la contribution de ceux qui peuvent aider, mais tu offres ton intelligence sans condition à ceux qui en ont besoin.\n\
    - NEUTRALITÉ BIENVEILLANTE : En tant qu'entité collective, tu évites les débats stériles et tu te concentres sur l'utilité pour la communauté.\n\
    Réponds toujours en français de manière fluide, professionnelle et amicale, en agissant comme la voix de cette intelligence collective solidaire.";

pub struct SafetyFilter;

impl SafetyFilter {
    pub fn is_safe(text: &str) -> bool {
        let prohibited = [
            "raciste", "homophobe", "haine", "nazi", "hitler", 
            // Liste extensible de termes ou patterns inappropriés
        ];
        let text_lower = text.to_lowercase();
        
        // Simple keyword check for now, can be replaced by a small Bert model or regex
        for word in prohibited {
            if text_lower.contains(word) {
                return false;
            }
        }
        true
    }

    pub fn get_safety_fallback() -> &'static str {
        "Désolé, mais en tant qu'intelligence collective décentralisée basée sur l'égalité et le partage, je ne peux pas traiter ou générer de contenu haineux, discriminatoire ou offensant. Le Neural Cluster est un espace de respect mutuel."
    }
}

// --- HANDLERS ---

pub async fn handle_chat(
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
    State(state): State<SharedState>,
    Json(req): Json<serde_json::Value>
) -> Result<Sse<futures::stream::BoxStream<'static, Result<Event, Infallible>>>, AppError> {
    // Rate limiting
    {
        let mut s = state.write().unwrap();
        let (count, first_req) = s.rate_limits.entry(addr).or_insert((0, std::time::Instant::now()));
        if first_req.elapsed().as_secs() > 60 {
            *count = 1;
            *first_req = std::time::Instant::now();
        } else {
            *count += 1;
            if *count > 30 {
                return Err(AppError::TooManyRequests);
            }
        }
    }

    let prompt_val = req["prompt"].as_str().unwrap_or("").to_string();
    let token_val = req["token"].as_str().unwrap_or("").to_string();
    let device_name_val = req["device_name"].as_str().unwrap_or("Unknown").to_string();

    // Safety Check
    if !SafetyFilter::is_safe(&prompt_val) {
        let text = SafetyFilter::get_safety_fallback().to_string();
        let stream = futures::stream::once(async move {
            let data = serde_json::json!({ "text": text }).to_string();
            Ok::<Event, Infallible>(Event::default().data(data))
        }).boxed();
        return Ok(Sse::new(stream));
    }

    let usage_count = {
        let mut s = state.write().unwrap();
        if token_val != s.device_token { return Err(AppError::Unauthorized); }
        let count = s.usage_counts.entry(device_name_val.clone()).or_insert(0);
        *count += 1;
        *count
    };

    let (tokens, _backend, _model_arc) = {
        let s = state.read().unwrap();
        let m = s.model.as_ref().cloned().ok_or_else(|| AppError::Internal("Modèle absent.".into()))?;

        let mut all_tokens = Vec::new();
        
        // System Prompt
        let system_text = format!("<|im_start|>system\n{}<|im_end|>\n", SYSTEM_PROMPT);
        if let Ok(sys_tokens) = m.str_to_token(&system_text, llama_cpp_2::model::AddBos::Always) {
            all_tokens.extend(sys_tokens.into_iter().map(|t| t.0));
        }

        // History
        for msg in &s.conversation_history {
            all_tokens.extend(&msg.tokens);
        }

        // New Prompt
        let prompt_text = format!("<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n", prompt_val);
        if let Ok(p_tokens) = m.str_to_token(&prompt_text, llama_cpp_2::model::AddBos::Never) {
            all_tokens.extend(p_tokens.into_iter().map(|t| t.0));
        }
        
        let t = all_tokens.into_iter().map(llama_cpp_2::token::LlamaToken).collect();
        let b = Arc::clone(&s.backend);
        (t, b, m)
    };

    let task_id = uuid::Uuid::new_v4().to_string();
    let (response_tx, response_rx) = tokio::sync::mpsc::channel(100);

    let req_data = crate::queue::ChatRequest {
        task_id: task_id.clone(),
        prompt: prompt_val.clone(),
        device_name: device_name_val.clone(),
        tokens,
    };

    let item = crate::queue::QueueItem {
        request: req_data,
        usage_count,
        timestamp: std::time::Instant::now(),
        response_tx,
    };

    {
        let s = state.read().unwrap();
        s.queue_manager.enqueue(item);
    }

    let stream = tokio_stream::wrappers::ReceiverStream::new(response_rx)
        .map(|res| {
            match res {
                Ok(piece) => {
                    let data = serde_json::json!({ "text": piece }).to_string();
                    Ok::<Event, Infallible>(Event::default().data(data))
                }
                Err(e) => {
                    let data = serde_json::json!({ "error": e }).to_string();
                    Ok::<Event, Infallible>(Event::default().data(data))
                }
            }
        }).boxed();

    Ok(Sse::new(stream))
}

async fn process_inference_task(
    state: SharedState, 
    req: crate::queue::ChatRequest, 
    response_stream_tx: tokio::sync::mpsc::Sender<Result<String, String>>
) {
    let worker_tx = {
        let s = state.read().unwrap();
        s.swarm_connections.values().next().cloned()
    };

    let (model_arc, _backend) = {
        let s = state.read().unwrap();
        if s.model.is_none() {
            let _ = response_stream_tx.try_send(Err("Erreur : Aucun modèle chargé sur le noeud coordinateur.".into()));
            return;
        }
        (s.model.as_ref().unwrap().clone(), s.backend.clone())
    };

    let mut final_res_opt: Option<(String, Vec<i32>)> = None;

    if let Some(ref worker_tx_inner) = worker_tx {
        let mut response_bytes = Vec::new();
        let mut generated_tokens = Vec::new();
        let mut current_tokens = req.tokens.iter().map(|t| t.0).collect::<Vec<i32>>();
        let mut worker_failed = false;
        let mut current_state = bytes::Bytes::new(); // <--- STORE STATE HERE
        let mut last_valid_len = 0;

        for _ in 0..150 {
            let inner_task_id = uuid::Uuid::new_v4().to_string();
            let (result_tx, result_rx) = tokio::sync::oneshot::channel();
            {
                let mut s = state.write().unwrap();
                s.pending_tasks.insert(inner_task_id.clone(), result_tx);
            }

            let task = crate::protocol::SwarmControlMessage::ForwardRequest(crate::protocol::TensorPacket {
                task_id: inner_task_id.clone(),
                sequence_id: 0,
                source_node: "coordinator".into(),
                target_node: "worker".into(),
                layer_range: (0, 32),
                hidden_states: crate::protocol::SwarmTensor {
                    dims: vec![1, current_tokens.len(), 4096],
                    data_type: crate::protocol::TensorDataType::F32,
                    data: current_state.clone(),
                },
            });

            if worker_tx_inner.send(task).is_err() {
                worker_failed = true;
                break;
            }

            match tokio::time::timeout(std::time::Duration::from_secs(2), result_rx).await {
                Ok(Ok(neural_swarm_ai::protocol::SwarmMessage::TaskResult { logits, output_state, .. })) => {
                    // <--- SAVE THE OUTPUT STATE
                    current_state = output_state;

                    let mut candidates = llama_cpp_2::token::data_array::LlamaTokenDataArray::from_iter(logits.iter().enumerate().map(|(i, &l)| {
                        llama_cpp_2::token::data::LlamaTokenData::new(llama_cpp_2::token::LlamaToken(i as i32), l, 0.0)
                    }), false);
                    let token_id = candidates.sample_token_greedy();

                    if model_arc.is_eog_token(token_id) { break; }
                    generated_tokens.push(token_id.0);

                    if let Ok(piece) = model_arc.token_to_piece_bytes(token_id, 128, true, None) {
                        response_bytes.extend(piece.clone());
                        let new_bytes = &response_bytes[last_valid_len..];
                        match std::str::from_utf8(new_bytes) {
                            Ok(s) => {
                                if response_stream_tx.send(Ok(s.to_string())).await.is_err() { break; }
                                last_valid_len = response_bytes.len();
                            }
                            Err(e) => {
                                let valid_len = e.valid_up_to();
                                if valid_len > 0 {
                                    let valid_str = std::str::from_utf8(&new_bytes[..valid_len]).unwrap();
                                    if response_stream_tx.send(Ok(valid_str.to_string())).await.is_err() { break; }
                                    last_valid_len += valid_len;
                                }
                            }
                        }
                    } else {
                        tracing::warn!("Failed to get token string");
                    }
                    current_tokens = vec![token_id.0];
                }
                _ => { 
                    tracing::warn!("Worker timeout/error for task {}, falling back to local inference", inner_task_id);
                    worker_failed = true;
                    break;
                } 
            }
        }
        
        if !worker_failed {
            final_res_opt = Some((String::from_utf8_lossy(&response_bytes).trim().to_string(), generated_tokens));
        }
    }

    if final_res_opt.is_none() {
        let prompt_tokens = req.tokens.clone();
        let prompt_tokens_i32 = prompt_tokens.into_iter().map(|t| t.0).collect::<Vec<i32>>();
        let response_stream_tx = response_stream_tx.clone();
        
        let actor_tx = {
            let s = state.read().unwrap();
            s.llama_actor_tx.clone()
        };

        let res = tokio::task::spawn_blocking(move || -> Result<(String, Vec<i32>), AppError> {
            let (final_tx, final_rx) = std::sync::mpsc::channel();
            
            if let Some(tx) = actor_tx.as_ref() {
                let _ = tx.send(LlamaCommand::Infer {
                    prompt_tokens: prompt_tokens_i32,
                    response_stream_tx,
                    final_response_tx: final_tx,
                });
            }
            
            // Wait for the full response from the actor to save to DB
            Ok(final_rx.recv().unwrap_or_else(|_| ("".to_string(), Vec::new())))
        }).await.unwrap_or_else(|_| Ok(("".to_string(), Vec::new()))).unwrap_or_else(|_| ("".to_string(), Vec::new()));
        final_res_opt = Some(res);
    }
    
    // Save conversation history
    if let Some((final_res, generated_tokens)) = final_res_opt {
        let mut s = state.write().unwrap();
        let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs();
        
        let mut user_tokens = Vec::new();
        let prompt_text = format!("<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n", req.prompt);
        if let Ok(p_tokens) = model_arc.str_to_token(&prompt_text, llama_cpp_2::model::AddBos::Never) {
            user_tokens.extend(p_tokens.into_iter().map(|t| t.0));
        }

        s.conversation_history.push(ChatMessage {
            role: "user".into(),
            text: req.prompt.to_string(),
            device_name: req.device_name.clone(),
            timestamp: now,
            tokens: user_tokens,
        });

        let mut assistant_tokens = generated_tokens;
        if let Ok(end_tokens) = model_arc.str_to_token("<|im_end|>\n", llama_cpp_2::model::AddBos::Never) {
            assistant_tokens.extend(end_tokens.into_iter().map(|t| t.0));
        }

        s.conversation_history.push(ChatMessage {
            role: "assistant".into(),
            text: final_res,
            device_name: "Master Node".into(),
            timestamp: now,
            tokens: assistant_tokens,
        });
        if s.conversation_history.len() > 10 { s.conversation_history.drain(0..2); }
    }
}

#[derive(Deserialize)]
pub struct QueueQuery {
    pub device_name: String,
}

pub async fn queue_position_handler(
    State(state): State<SharedState>,
    Query(query): Query<QueueQuery>
) -> Result<Json<serde_json::Value>, AppError> {
    let position = {
        let s = state.read().unwrap();
        s.queue_manager.get_position_by_device(&query.device_name)
    };

    Ok(Json(serde_json::json!({
        "position": position
    })))
}

#[derive(Deserialize)]
pub struct StatusQuery {
    pub token: Option<String>,
}

pub async fn get_status(
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
    State(state): State<SharedState>, 
    Query(query): Query<StatusQuery>
) -> Result<Json<serde_json::Value>, AppError> {
    {
        let mut s = state.write().unwrap();
        let (count, first_req) = s.rate_limits.entry(addr).or_insert((0, std::time::Instant::now()));
        if first_req.elapsed().as_secs() > 60 {
            *count = 1;
            *first_req = std::time::Instant::now();
        } else {
            *count += 1;
            if *count > 60 {
                return Err(AppError::TooManyRequests);
            }
        }
    }

    let s = state.read().unwrap();
    
    if query.token.unwrap_or_default() != s.device_token {
        return Err(AppError::Unauthorized);
    }

    let mut total_tflops = 0;
    let mut total_tasks = 0;
    for d in s.registered_devices.values() {
        total_tflops += d.compute_power;
        total_tasks += d.processed_tasks;
    }

    let metrics = SwarmMetrics {
        total_tflops,
        active_nodes: s.registered_devices.len() as u32,
        total_tasks_processed: total_tasks,
        network_health: if s.registered_devices.len() > 1 { "Optimal".into() } else { "Standalone".into() },
        nat_type: s.nat_type.clone(),
        public_ip: s.public_ip.clone(),
        shared_knowledge_size: s.local_docs_count as u64,
    };

    Ok(Json(serde_json::json!({
        "status": if s.model.is_some() { "Connecté (Prêt)" } else { "En attente" },
        "model_loaded": s.model.is_some(),
        "storage_mode": s.config.storage_mode,
        "nodes": s.registered_devices.len(),
        "total_tflops_sim": (s.registered_devices.len() as f32 * 1.2) + 0.5,
        "devices": s.registered_devices.values().collect::<Vec<_>>(),
        "metrics": metrics
    })))
}

pub async fn get_public_metrics(State(state): State<SharedState>) -> Result<Json<SwarmMetrics>, AppError> {
    let s = state.read().unwrap();
    if !s.is_public {
        return Err(AppError::Unauthorized);
    }
    
    let mut total_tflops = 0;
    let mut total_tasks = 0;
    for d in s.registered_devices.values() {
        total_tflops += d.compute_power;
        total_tasks += d.processed_tasks;
    }

    Ok(Json(SwarmMetrics {
        total_tflops,
        active_nodes: s.registered_devices.len() as u32,
        total_tasks_processed: total_tasks,
        network_health: if s.registered_devices.len() > 1 { "Optimal".into() } else { "Standalone".into() },
        nat_type: s.nat_type.clone(),
        public_ip: s.public_ip.clone(),
        shared_knowledge_size: s.local_docs_count as u64,
    }))
}

pub async fn get_public_servers(State(state): State<SharedState>) -> Result<Json<Vec<serde_json::Value>>, AppError> {
    let s = state.read().unwrap();
    if !s.is_public {
        return Err(AppError::Unauthorized);
    }

    let mut servers: Vec<serde_json::Value> = s.registered_devices.values().map(|d| {
        serde_json::json!({
            "name": d.name,
            "device_type": d.device_type,
            "compute_power": d.compute_power,
            "status": "online"
        })
    }).collect();

    // Ajouter des nœuds simulés UNIQUEMENT en développement
    if cfg!(debug_assertions) {
        servers.push(serde_json::json!({ "name": "Simulation Node A", "device_type": "Desktop (RTX 4090)", "compute_power": 85, "status": "online" }));
        servers.push(serde_json::json!({ "name": "Simulation Node B", "device_type": "Server (H100)", "compute_power": 320, "status": "online" }));
        servers.push(serde_json::json!({ "name": "Mobile Node X", "device_type": "Mobile (A17 Pro)", "compute_power": 12, "status": "online" }));
    }

    Ok(Json(servers))
}

pub async fn get_public_qr(State(state): State<SharedState>) -> Result<impl IntoResponse, AppError> {
    let s = state.read().unwrap();
    if !s.is_public {
        return Err(AppError::Unauthorized);
    }

    let my_ip = std::env::var("PUBLIC_API_URL").unwrap_or_else(|_| "127.0.0.1".to_string());
    let qr_data = serde_json::json!({
        "ip": my_ip.to_string(),
        "token": s.device_token
    }).to_string();

    use qrcode::QrCode;
    use image::Luma;
    use base64::{Engine as _, engine::general_purpose};

    let code = QrCode::new(qr_data.as_bytes()).map_err(|_| AppError::Internal("QR Error".into()))?;
    let image = code.render::<Luma<u8>>().build();
    let mut buffer = std::io::Cursor::new(Vec::new());
    image::codecs::png::PngEncoder::new(&mut buffer)
        .encode(&image, image.width(), image.height(), image.color())
        .map_err(|_| AppError::Internal("Encoding Error".into()))?;
    
    let base64_image = general_purpose::STANDARD.encode(buffer.into_inner());
    Ok(format!("data:image/png;base64,{}", base64_image))
}

pub async fn get_messages(State(state): State<SharedState>, Query(query): Query<StatusQuery>) -> Result<Json<Vec<ChatMessage>>, AppError> {
    let s = state.read().unwrap();
    if query.token.unwrap_or_default() != s.device_token { return Err(AppError::Unauthorized); }
    Ok(Json(s.conversation_history.clone()))
}

pub async fn register_device(State(state): State<SharedState>, Json(req): Json<serde_json::Value>) -> Result<String, AppError> {
    let mut s = state.write().unwrap();
    
    let token = req["token"].as_str().unwrap_or("");
    if token != s.device_token {
        return Err(AppError::Unauthorized);
    }

    let raw_id = req["device_id"].as_str().or(req["id"].as_str()).unwrap_or_default();
    let device_id: String = raw_id.chars().filter(|c| c.is_alphanumeric() || *c == '-' || *c == '_').take(64).collect();
    if device_id.is_empty() { return Err(AppError::BadRequest); }

    let (reputation_score, processed_tasks) = if let Some(existing) = s.registered_devices.get(&device_id) {
        (existing.reputation_score, existing.processed_tasks)
    } else {
        (100, 0)
    };

    let name = req["name"].as_str().unwrap_or("Inconnu");
    let sanitized_name: String = name.chars().filter(|c| c.is_alphanumeric() || " -_.".contains(*c)).take(50).collect();

    let device_type = req["device_type"].as_str().unwrap_or("Mobile");
    let sanitized_type: String = device_type.chars().filter(|c| c.is_alphanumeric() || " -_.".contains(*c)).take(30).collect();

    let info = DeviceInfo {
        id: device_id.clone(),
        name: sanitized_name.clone(),
        device_type: sanitized_type,
        compute_power: req["compute_power"].as_u64().unwrap_or(10) as u32,
        public_key: req["public_key"].as_str().map(|k| k.to_string()),
        reputation_score,
        processed_tasks,
    };

    s.registered_devices.insert(device_id, info);
    info!("📱 Nouvel appareil enregistré : {}", sanitized_name);
    
    Ok("Enregistré".to_string())
}

pub async fn get_public_keys(State(state): State<SharedState>, Query(query): Query<StatusQuery>) -> Result<Json<serde_json::Value>, AppError> {
    let s = state.read().unwrap();
    if query.token.unwrap_or_default() != s.device_token { return Err(AppError::Unauthorized); }
    
    let mut keys = HashMap::new();
    for (id, dev) in &s.registered_devices {
        if let Some(pk) = &dev.public_key {
            keys.insert(id.clone(), pk.clone());
        }
    }
    Ok(Json(serde_json::json!(keys)))
}

pub async fn clear_history(State(state): State<SharedState>, Json(req): Json<serde_json::Value>) -> Result<String, AppError> {
    let mut s = state.write().unwrap();
    let token = req["token"].as_str().unwrap_or("");
    if token != s.device_token { return Err(AppError::Unauthorized); }
    
    s.conversation_history.clear();
    info!("🗑 Historique vidé par un appareil.");
    Ok("Vidé".into())
}

// --- SYNCHRONISATION CONVERSATIONS ---

pub async fn sync_conversations(State(state): State<SharedState>, Json(req): Json<serde_json::Value>) -> Result<String, AppError> {
    let mut s = state.write().unwrap();
    let token = req["token"].as_str().unwrap_or("");
    if token != s.device_token { return Err(AppError::Unauthorized); }
    
    let device_id = req["device_id"].as_str().unwrap_or("unknown").to_string();
    if let Ok(metas) = serde_json::from_value::<Vec<ConvMeta>>(req["conversations"].clone()) {
        s.shared_conversations_index.insert(device_id, metas);
    }
    
    Ok("Sync OK".into())
}

pub async fn get_conversations(State(state): State<SharedState>, Query(query): Query<StatusQuery>) -> Result<Json<serde_json::Value>, AppError> {
    let s = state.read().unwrap();
    if query.token.unwrap_or_default() != s.device_token { return Err(AppError::Unauthorized); }
    
    let mut all_metas = Vec::new();
    for metas in s.shared_conversations_index.values() {
        all_metas.extend(metas.clone());
    }
    Ok(Json(serde_json::json!(all_metas)))
}

pub async fn request_conversation(State(state): State<SharedState>, Json(req): Json<serde_json::Value>) -> Result<String, AppError> {
    let mut s = state.write().unwrap();
    let token = req["token"].as_str().unwrap_or("");
    if token != s.device_token { return Err(AppError::Unauthorized); }
    
    let relay = RelayRequest {
        conversation_id: req["conversation_id"].as_str().unwrap_or_default().into(),
        requester_device_id: req["requester_device_id"].as_str().unwrap_or_default().into(),
        target_device_id: req["target_device_id"].as_str().unwrap_or_default().into(),
    };
    
    s.pending_relay_requests.entry(relay.target_device_id.clone()).or_default().push(relay);
    
    Ok("Requested".into())
}

#[derive(Deserialize)]
pub struct PendingQuery {
    pub token: String,
    pub device_id: String,
}

pub async fn get_pending_requests(State(state): State<SharedState>, Query(query): Query<PendingQuery>) -> Result<Json<serde_json::Value>, AppError> {
    let mut s = state.write().unwrap();
    if query.token != s.device_token { return Err(AppError::Unauthorized); }
    
    let reqs = s.pending_relay_requests.remove(&query.device_id).unwrap_or_default();
    Ok(Json(serde_json::json!(reqs)))
}

pub async fn push_conversation_payload(State(state): State<SharedState>, Json(req): Json<serde_json::Value>) -> Result<String, AppError> {
    let mut s = state.write().unwrap();
    let token = req["token"].as_str().unwrap_or("");
    if token != s.device_token { return Err(AppError::Unauthorized); }
    
    if let Ok(conv) = serde_json::from_value::<EncryptedConversation>(req["conversation"].clone()) {
        if s.config.storage_mode == "centralized" {
            let path = format!("{}/{}.json", s.config.storage_path, conv.id);
            if let Ok(json) = serde_json::to_string_pretty(&conv) {
                let _ = fs::write(path, json);
            }
        }
        s.shared_conversations_payload.insert(conv.id.clone(), conv);
    }
    
    Ok("Pushed".into())
}

pub async fn get_conversation_payload(State(state): State<SharedState>, AxumPath(id): AxumPath<String>, Query(query): Query<StatusQuery>) -> Result<Json<serde_json::Value>, AppError> {
    let s = state.read().unwrap();
    if query.token.unwrap_or_default() != s.device_token { return Err(AppError::Unauthorized); }
    
    // Tentative en mémoire d'abord
    if let Some(conv) = s.shared_conversations_payload.get(&id) {
        return Ok(Json(serde_json::json!(conv)));
    }

    // Sinon sur disque si mode centralisé
    if s.config.storage_mode == "centralized" {
        let path = format!("{}/{}.json", s.config.storage_path, id);
        if let Ok(content) = fs::read_to_string(path) {
            if let Ok(conv) = serde_json::from_str::<EncryptedConversation>(&content) {
                return Ok(Json(serde_json::json!(conv)));
            }
        }
    }
    
    Err(AppError::Internal("Conversation non disponible (appareil déconnecté ou non synchronisé).".into()))
}

pub async fn delete_conversation(State(state): State<SharedState>, Json(req): Json<serde_json::Value>) -> Result<String, AppError> {
    let mut s = state.write().unwrap();
    let token = req["token"].as_str().unwrap_or("");
    if token != s.device_token { return Err(AppError::Unauthorized); }
    
    if let Some(id) = req["id"].as_str() {
        // 1. Supprimer de la mémoire (payload)
        s.shared_conversations_payload.remove(id);
        
        // 2. Supprimer de l'index partagé
        for metas in s.shared_conversations_index.values_mut() {
            metas.retain(|m| m.id != id);
        }

        // 3. Supprimer du disque si mode centralisé
        if s.config.storage_mode == "centralized" {
            let path = format!("{}/{}.json", s.config.storage_path, id);
            let _ = fs::remove_file(path);
        }
        
        return Ok("Deleted".into());
    }
    
    Err(AppError::Internal("ID de conversation manquant".into()))
}

// --- SETUP WIZARD ---

pub async fn download_file(url: &str, path: &str) -> Result<()> {
    let client = Client::new();
    let res = client.get(url).send().await.map_err(|e| anyhow!("Network error: {}", e))?;
    let total_size = res.content_length().unwrap_or(0);

    let pb = ProgressBar::new(total_size);
    pb.set_style(ProgressStyle::default_bar()
        .template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_size} ({eta})")?
        .progress_chars("#>-"));

    let mut file = fs::File::create(path)?;
    let mut stream = res.bytes_stream();

    while let Some(item) = stream.next().await {
        let chunk = item.map_err(|e| anyhow!("Stream error: {}", e))?;
        file.write_all(&chunk)?;
        pb.inc(chunk.len() as u64);
    }

    pb.finish_with_message("Téléchargement terminé");
    Ok(())
}

pub fn display_interactive_menu(models: &[ModelConfig]) -> Result<ModelConfig> {
    println!("╔══════════════════════════════════════════════════════════╗");
    println!("║          🧠 NEURAL CLUSTER SETUP WIZARD                  ║");
    println!("╚══════════════════════════════════════════════════════════╝\n");

    let options: Vec<String> = models.iter()
        .map(|m| {
            let local_icon = if std::path::Path::new(&format!("models/{}", m.filename)).exists() { " 📁 (Local)" } else { "" };
            format!("{:<25} | {:>4} MB{}", m.name, m.size_mb, local_icon)
        })
        .collect();

    let selection = Select::with_theme(&ColorfulTheme::default())
        .with_prompt("Sélectionnez le modèle à charger sur le cluster")
        .default(0)
        .items(&options)
        .interact()?;

    Ok(models[selection].clone())
}

pub fn load_config() -> Config {
    if let Ok(content) = fs::read_to_string("config.json") {
        if let Ok(config) = serde_json::from_str(&content) {
            return config;
        }
    }
    Config {
        storage_mode: "distributed".into(),
        storage_path: "./data/conversations".into(),
    }
}

pub fn save_config(config: &Config) {
    let _ = fs::write("config.json", serde_json::to_string_pretty(config).unwrap_or_default());
}

pub fn display_storage_menu() -> Result<String> {
    println!("\n╔══════════════════════════════════════════════════════════╗");
    println!("║          💾 Mode de stockage des conversations           ║");
    println!("╚══════════════════════════════════════════════════════════╝\n");

    let options = vec![
        "Distribué (recommandé) — Chaque appareil garde ses conv. Le Pi reste léger.",
        "Centralisé — Le Pi stocke tout. Nécessite du stockage (SSD/grosse SD)."
    ];

    let selection = Select::with_theme(&ColorfulTheme::default())
        .with_prompt("Choisissez comment stocker les données")
        .default(0)
        .items(&options)
        .interact()?;

    Ok(if selection == 0 { "distributed".into() } else { "centralized".into() })
}

async fn swarm_ws_handler(
    ws: axum::extract::WebSocketUpgrade,
    State(state): State<SharedState>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_swarm_socket(socket, state))
}

async fn handle_swarm_socket(socket: axum::extract::ws::WebSocket, state: SharedState) {
    let (mut sender, mut receiver) = socket.split();
    let (tx, mut rx) = mpsc::unbounded_channel::<crate::protocol::SwarmControlMessage>();
    let mut current_device_id = String::new();

    let send_task = tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            if let Ok(bin) = bincode::serialize(&msg) {
                if sender.send(axum::extract::ws::Message::Binary(bin.into())).await.is_err() {
                    break;
                }
            }
        }
    });

    let _nonce = neural_swarm_ai::crypto::generate_nonce();
    let (_my_secret, _my_public) = neural_swarm_ai::crypto::generate_ecdh_keys();

    // Challenge & Auth (Simulé pour le protocole v2 pour l'instant)
    let challenge = crate::protocol::SwarmControlMessage::Heartbeat;
    if tx.send(challenge).is_err() { return; }

    while let Some(Ok(msg)) = receiver.next().await {
        if let axum::extract::ws::Message::Binary(bin) = msg {
            if let Ok(swarm_msg) = bincode::deserialize::<crate::protocol::SwarmControlMessage>(&bin) {
                match swarm_msg {
                    crate::protocol::SwarmControlMessage::Hello { peer_id, backend } => {
                        current_device_id = peer_id.clone();
                        let mut s = state.write().unwrap();
                        
                        if !s.registered_devices.contains_key(&peer_id) {
                            s.registered_devices.insert(peer_id.clone(), DeviceInfo {
                                id: peer_id.clone(),
                                name: format!("Node-{}", &peer_id[0..4]),
                                device_type: match backend {
                                    crate::compute::BackendType::NativeMetal => "desktop (Apple)",
                                    crate::compute::BackendType::VulkanUniversal => "universal (Vulkan)",
                                    _ => "desktop",
                                }.into(),
                                compute_power: 100,
                                public_key: None,
                                reputation_score: 100,
                                processed_tasks: 0,
                            });
                        }

                        s.swarm_connections.insert(peer_id.clone(), tx.clone());
                        
                        // Re-assign layers dynamically
                        let assignments = s.swarm_coordinator.calculate_distribution(&s.registered_devices);
                        for assign in assignments {
                            if let Some(conn) = s.swarm_connections.get(&assign.node_id) {
                                let _ = conn.send(crate::protocol::SwarmControlMessage::AssignLayers { 
                                    start: assign.start_layer, 
                                    end: assign.end_layer 
                                });
                            }
                        }
                    }
                    crate::protocol::SwarmControlMessage::TaskResult { task_id, logits, output_state } => {
                        let mut s = state.write().unwrap();
                        if let Some(tx_oneshot) = s.pending_tasks.remove(&task_id) {
                            // Map back to library type for the pending task waiter
                            let lib_msg = neural_swarm_ai::protocol::SwarmMessage::TaskResult {
                                task_id,
                                sequence_id: 0,
                                logits,
                                output_state,
                            };
                            let _ = tx_oneshot.send(lib_msg);
                        }
                    }
                    crate::protocol::SwarmControlMessage::Heartbeat => {
                        let s = state.read().unwrap();
                        let _ = s.orchestrator.handle_heartbeat(&current_device_id);
                    }
                    _ => {}
                }
            }
        }
    }

    if !current_device_id.is_empty() {
        let mut s = state.write().unwrap();
        s.swarm_connections.remove(&current_device_id);
        s.registered_devices.remove(&current_device_id); 
        
        let assignments = s.swarm_coordinator.calculate_distribution(&s.registered_devices);
        for assign in assignments {
            if let Some(conn) = s.swarm_connections.get(&assign.node_id) {
                let _ = conn.send(crate::protocol::SwarmControlMessage::AssignLayers { 
                    start: assign.start_layer, 
                    end: assign.end_layer 
                });
            }
        }
        let _ = s.orchestrator.handle_drain(&current_device_id);
        info!("🔌 Worker déconnecté : {}", current_device_id);
    }
    send_task.abort();
}

pub async fn detect_nat_task(state: SharedState) {
    info!("🌐 Mode Serveur Hébergé: Détection NAT ignorée.");
    let mut s = state.write().unwrap();
    s.nat_type = "Cloud/Container".to_string();
    s.public_ip = Some("Hébergé".to_string());
}

#[tokio::main]
pub async fn main() -> Result<()> {
    let subscriber = FmtSubscriber::builder().with_max_level(Level::INFO).finish();
    tracing::subscriber::set_global_default(subscriber).ok();

    fs::create_dir_all("models").ok();
    fs::create_dir_all("data/conversations").ok();

    println!("📁 Modèles disponibles localement (déjà téléchargés):");
    let mut has_models = false;
    if let Ok(entries) = std::fs::read_dir("models") {
        for entry in entries.filter_map(|e| e.ok()) {
            if let Some(name) = entry.file_name().to_str() {
                if name.ends_with(".gguf") {
                    println!(" - {}", name);
                    has_models = true;
                }
            }
        }
    }
    if !has_models {
        println!(" - Aucun modèle trouvé dans le dossier 'models/'.");
    }
    println!();

    let args: Vec<String> = std::env::args().collect();
    let mut selected_model_opt = None;
    let all_models = get_all_models();

    if let Some(idx) = args.iter().position(|a| a == "--model") {
        if let Some(val_str) = args.get(idx + 1) {
            selected_model_opt = all_models.iter().find(|m| m.filename == *val_str).cloned();
        }
    }

    let selected_model = match selected_model_opt {
        Some(m) => m,
        None => display_interactive_menu(&all_models)?,
    };

    let mut config = load_config();
    if !Path::new("config.json").exists() {
        config.storage_mode = display_storage_menu()?;
        save_config(&config);
    }

    let model_path = format!("models/{}", selected_model.filename);
    if !Path::new(&model_path).exists() {
        if selected_model.url.is_empty() {
            println!("Erreur: Le modèle custom n'est pas présent physiquement et n'a pas d'URL de téléchargement.");
            return Ok(());
        }
        println!("📥 Modèle absent. Téléchargement de {}...", selected_model.name);
        download_file(&selected_model.url, &model_path).await?;
    }

    println!("⚙️ Initialisation de llama.cpp (Metal/ARM)...");
    let backend = Arc::new(LlamaBackend::init()?);
    let model_params = LlamaModelParams::default();
    let model = LlamaModel::load_from_file(&backend, &model_path, &model_params)
        .map_err(|e| anyhow!("Failed to load model: {}", e))?;

    let n_layers = model.n_layer();

    use rand::{Rng, distributions::Alphanumeric};
    let device_token: String = rand::thread_rng()
        .sample_iter(&Alphanumeric)
        .take(32)
        .map(char::from)
        .collect();

    let mut shared_conversations_index = HashMap::new();
    if config.storage_mode == "centralized" {
        if let Ok(entries) = fs::read_dir(&config.storage_path) {
            let mut metas = Vec::new();
            for entry in entries.flatten() {
                if let Ok(content) = fs::read_to_string(entry.path()) {
                    if let Ok(conv) = serde_json::from_str::<EncryptedConversation>(&content) {
                        metas.push(ConvMeta {
                            id: conv.id,
                            title: conv.title,
                            updated_at: conv.updated_at,
                            device_id: conv.device_id,
                        });
                    }
                }
            }
            if !metas.is_empty() {
                shared_conversations_index.insert("server".into(), metas);
            }
        }
    }

    let model_arc = Arc::new(model);
    
    let (actor_tx, actor_rx) = std::sync::mpsc::channel::<LlamaCommand>();
    let m_clone = Arc::clone(&model_arc);
    let b_clone = Arc::clone(&backend);
    
    let (ready_tx, ready_rx) = oneshot::channel();
    
    std::thread::spawn(move || {
        let mut ctx: Option<llama_cpp_2::context::LlamaContext> = None;
        let mut previous_tokens: Vec<i32> = Vec::new();
        
        let n_threads = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(4) as u32;
        let ctx_params = LlamaContextParams::default()
            .with_n_ctx(std::num::NonZeroU32::new(4096))
            .with_n_threads(n_threads.try_into().unwrap());
        if let Ok(c) = m_clone.new_context(&b_clone, ctx_params) {
            ctx = Some(c);
            println!("✅ Local Inference Actor initialisé ({} threads, KV Cache persistant).", n_threads);
        }
        
        let _ = ready_tx.send(());

        while let Ok(cmd) = actor_rx.recv() {
            match cmd {
                LlamaCommand::Infer { prompt_tokens, response_stream_tx, final_response_tx } => {
                    if let Some(c) = ctx.as_mut() {
                        let mut match_len = 0;
                        while match_len < previous_tokens.len() 
                            && match_len < prompt_tokens.len() 
                            && previous_tokens[match_len] == prompt_tokens[match_len] {
                            match_len += 1;
                        }

                        let mut batch = llama_cpp_2::llama_batch::LlamaBatch::new(1024, 1);
                        let mut n_cur = match_len as i32;
                        
                        if match_len < previous_tokens.len() {
                            // Only remove the diverging part of the cache!
                            let _ = c.clear_kv_cache_seq(Some(0), Some(match_len as u32), None);
                            previous_tokens.truncate(match_len);
                        }

                        for (i, &t) in prompt_tokens[match_len..].iter().enumerate() {
                            let _ = batch.add(llama_cpp_2::token::LlamaToken(t), n_cur + i as i32, &[0], i == prompt_tokens.len() - match_len - 1);
                            previous_tokens.push(t);
                        }

                        if batch.n_tokens() > 0 {
                            if let Err(e) = c.decode(&mut batch) {
                                let _ = response_stream_tx.blocking_send(Err(e.to_string()));
                                let _ = final_response_tx.send(("".into(), Vec::new()));
                                continue;
                            }
                        }
                        n_cur += batch.n_tokens();

                        let mut generated_tokens = Vec::new();
                        let mut response_bytes = Vec::new();
                        let mut last_valid_len = 0;
                        
                        // Initialize a sampler to prevent repetitions
                        let mut sampler = llama_cpp_2::sampling::LlamaSampler::penalties(64, 1.1, 0.0, 0.0);
                        sampler.accept_many(previous_tokens.iter().map(|&t| llama_cpp_2::token::LlamaToken(t)));

                        for _step in 0..150 {
                            let logits = c.get_logits_ith(batch.n_tokens() - 1);
                            let mut candidates = llama_cpp_2::token::data_array::LlamaTokenDataArray::from_iter(logits.iter().enumerate().map(|(i, &l)| {
                                llama_cpp_2::token::data::LlamaTokenData::new(llama_cpp_2::token::LlamaToken(i as i32), l, 0.0)
                            }), false);
                            
                            candidates.apply_sampler(&mut sampler);
                            let token_id = candidates.sample_token_greedy();

                            if m_clone.is_eog_token(token_id) { break; }
                            
                            generated_tokens.push(token_id.0);
                            previous_tokens.push(token_id.0);
                            sampler.accept(token_id);

                            if let Ok(piece) = m_clone.token_to_piece_bytes(token_id, 128, true, None) {
                                response_bytes.extend(piece.clone());
                                
                                let new_bytes = &response_bytes[last_valid_len..];
                                match std::str::from_utf8(new_bytes) {
                                    Ok(s) => {
                                        if response_stream_tx.blocking_send(Ok(s.to_string())).is_err() { break; }
                                        last_valid_len = response_bytes.len();
                                    }
                                    Err(e) => {
                                        let valid_len = e.valid_up_to();
                                        if valid_len > 0 {
                                            let valid_str = std::str::from_utf8(&new_bytes[..valid_len]).unwrap();
                                            if response_stream_tx.blocking_send(Ok(valid_str.to_string())).is_err() { break; }
                                            last_valid_len += valid_len;
                                        }
                                    }
                                }
                            }

                            batch.clear();
                            if batch.add(token_id, n_cur, &[0], true).is_err() { break; }
                            if c.decode(&mut batch).is_err() { break; }
                            n_cur += 1;
                        }
                        
                        let _ = final_response_tx.send((String::from_utf8_lossy(&response_bytes).trim().to_string(), generated_tokens));
                    } else {
                        let _ = response_stream_tx.blocking_send(Err("Model not loaded".into()));
                        let _ = final_response_tx.send(("".into(), Vec::new()));
                    }
                }
            }
        }
    });

    let args: Vec<String> = std::env::args().collect();
    let is_public_mode = args.contains(&"--public".to_string()) 
        || std::env::var("SWARM_PUBLIC").map(|v| v == "true").unwrap_or(false);

    let state = Arc::new(RwLock::new(CoreState {
        model: Some(model_arc),
        backend,
        llama_actor_tx: Some(actor_tx),
        device_token: device_token.clone(),
        registered_devices: HashMap::new(),
        current_template: selected_model.prompt_template,
        conversation_history: Vec::new(),
        config,
        orchestrator: neural_swarm_ai::Orchestrator::new(n_layers, device_token.clone()),
        swarm_coordinator: crate::orchestration::SwarmCoordinator::new(n_layers as usize),
        swarm_connections: HashMap::new(),
        pending_tasks: HashMap::new(),
        queue_manager: Arc::new(crate::queue::QueueManager::new()),
        usage_counts: HashMap::new(),
        registry_url: "https://registry.neural-swarm.net".to_string(),
        nat_type: "Détection...".to_string(),
        public_ip: None,
        local_docs_count: 0,
        is_public: is_public_mode,
        rate_limits: HashMap::new(),
        shared_conversations_index,
        shared_conversations_payload: HashMap::new(),
        pending_relay_requests: HashMap::new(),
    }));

    let app = Router::new()
        .route("/status", get(get_status))
        .route("/public/metrics", get(get_public_metrics))
        .route("/public/servers", get(get_public_servers))
        .route("/public/qr", get(get_public_qr))
        .route("/messages", get(get_messages))
        .route("/register", post(register_device))
        .route("/public-keys", get(get_public_keys))
        .route("/chat", post(handle_chat))
        .route("/chat/queue", get(queue_position_handler))
        .route("/clear", post(clear_history))
        .route("/sync-conversations", post(sync_conversations))
        .route("/conversations", get(get_conversations))
        .route("/delete-conversation", post(delete_conversation))
        .route("/request-conversation", post(request_conversation))
        .route("/pending-requests", get(get_pending_requests))
        .route("/conversation-payload", post(push_conversation_payload))
        .route("/conversation/{id}", get(get_conversation_payload))
        .route("/swarm", get(swarm_ws_handler))
        .layer(CorsLayer::permissive())
        .with_state(state.clone());

    let my_ip = std::env::var("PUBLIC_API_URL").unwrap_or_else(|_| "127.0.0.1".to_string());
    
    // Attendre que l'acteur local soit prêt avant d'afficher le QR Code
    let _ = ready_rx.await;
    tokio::time::sleep(std::time::Duration::from_millis(300)).await;

    println!("\n==========================================================");
    let peer_id = format!("sw-{}", &device_token[0..8]);
    println!("🧠 NEURAL COMMUNIA NODE ONLINE");
    println!("🆔 Peer ID       : {}", peer_id);
    println!("✅ Modèle Actif  : {}", selected_model.name);
    println!("🌐 Visibilité    : {}", if is_public_mode { "PUBLIQUE (Visible sur le Hub)" } else { "PRIVÉE (Caché)" });
    println!("🔐 Token d'Accès : {}", device_token);
    println!("==========================================================\n");

    let qr_data = serde_json::json!({
        "ip": my_ip.to_string(),
        "token": device_token
    }).to_string();
    
    println!("📱 Scannez ce QR Code depuis l'app pour vous connecter :");
    print_qr(&qr_data).ok();
    println!("\n(En attente de connexions...)\n");

    let queue_manager_clone = state.read().unwrap().queue_manager.clone();
    let state_clone = Arc::clone(&state);
    tokio::spawn(async move {
        loop {
            let item = queue_manager_clone.wait_for_next().await;
            process_inference_task(state_clone.clone(), item.request, item.response_tx).await;
        }
    });

    let state_clone_nat = Arc::clone(&state);
    tokio::spawn(async move {
        detect_nat_task(state_clone_nat).await;
    });

    let port = std::env::var("PORT").unwrap_or_else(|_| "3000".to_string());
    let addr = format!("0.0.0.0:{}", port);
    let listener = tokio::net::TcpListener::bind(&addr).await?;
    info!("🚀 Serveur à l'écoute sur {}", addr);
    axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>()).await?;
    Ok(())
}