use std::path::Path;
use anyhow::{Result, anyhow};
use super::backend::{ComputeBackend, BackendType};
use llama_cpp_2::model::LlamaModel;
use llama_cpp_2::llama_backend::LlamaBackend;
use llama_cpp_2::model::params::LlamaModelParams;
use std::sync::Arc;
use tracing::info;

pub struct LlamaCppBackend {
    backend: Arc<LlamaBackend>,
    model: Option<LlamaModel>,
    start_layer: usize,
    end_layer: usize,
}

impl LlamaCppBackend {
    pub fn new() -> Result<Self> {
        let backend = LlamaBackend::init().map_err(|e| anyhow!("Failed to init llama backend: {}", e))?;
        Ok(Self {
            backend: Arc::new(backend),
            model: None,
            start_layer: 0,
            end_layer: 0,
        })
    }
}

impl ComputeBackend for LlamaCppBackend {
    fn load_layer_slice(&mut self, model_path: &Path, start_layer: usize, end_layer: usize) -> Result<()> {
        let model_params = LlamaModelParams::default();
        // NOTE: Dans une vraie implémentation de slicing custom, on utiliserait 
        // les paramètres de llama.cpp pour ne charger que certaines couches en VRAM.
        let model = LlamaModel::load_from_file(&self.backend, model_path, &model_params)
            .map_err(|e| anyhow!("Failed to load model: {}", e))?;
        
        self.model = Some(model);
        self.start_layer = start_layer;
        self.end_layer = end_layer;
        Ok(())
    }

    fn forward_slice(&mut self, hidden_states: &[f32], seq_len: usize) -> Result<Vec<f32>> {
        // Dans une architecture Swarm, cette fonction :
        // 1. Reçoit l'hidden_state du noeud précédent
        // 2. L'injecte dans le contexte local (start_layer)
        // 3. Calcule jusqu'à end_layer
        // 4. Extrait le nouvel hidden_state pour le suivant
        info!("Calcul Swarm sur couches {} à {} (seq_len: {})", self.start_layer, self.end_layer, seq_len);
        
        // Simulation pour le protocole : on retourne un vecteur modifié
        Ok(hidden_states.iter().map(|x| x + 0.01).collect())
    }

    fn get_backend_type(&self) -> BackendType {
        // Logique de priorité : Native > Vulkan > CPU
        if cfg!(feature = "metal") {
            BackendType::NativeMetal
        } else if cfg!(feature = "cuda") {
            BackendType::NativeCuda
        } else if cfg!(feature = "rocm") {
            BackendType::NativeRocm
        } else if cfg!(feature = "vulkan") {
            // Vulkan est notre pont universel pour AMD Windows et Android (Snapdragon/Adreno)
            BackendType::VulkanUniversal
        } else if cfg!(feature = "opencl") {
            BackendType::OpenCL
        } else {
            BackendType::Cpu
        }
    }
}