use pyo3::prelude::*;
use pyo3::exceptions::PyRuntimeError;
use candle_core::{Tensor, Device};
use std::collections::HashMap;
use std::fs::File;
use std::path::Path;

/// Helper function to perform raw tensor manipulation/quantization
fn perform_quantization(input_path: &str, output_path: &str, precision: &str) -> Result<(), String> {
    let device = Device::Cpu; // Can switch to Device::new_cuda(0) for GPU
    
    // 1. Read raw Llama 3.2 safetensors from the input directory
    let content = std::fs::read(input_path)
        .map_err(|e| format!("Failed to read input safetensor file: {}", e))?;
    
    let safetensor = safetensors::SafeTensors::deserialize(&content)
        .map_err(|e| format!("Invalid Safetensor format: {}", e))?;

    let mut new_tensors = HashMap::new();

    // 2. Iterate through layers (Weights, Attention blocks, Tensors)
    for (name, tensor_view) in safetensor.tensors() {
        let candle_tensor = candle_core::safetensors::load_buffer(tensor_view, &device)
            .map_err(|e| format!("Candle failed to parse tensor {}: {}", name, e))?;

        // 3. Low-level optimization based on UI input
        let optimized_tensor = match precision {
            "F16" => candle_tensor.to_dtype(candle_core::DType::F16)
                .map_err(|e| e.to_string())?,
            "BF16" => candle_tensor.to_dtype(candle_core::DType::BF16)
                .map_err(|e| e.to_string())?,
            _ => candle_tensor // Keep original if layout matches
        };

        new_tensors.insert(name, optimized_tensor);
    }

    // 4. Safely serialize and stream the weights back to disk
    let output_file = Path::new(output_path);
    candle_core::safetensors::save(&new_tensors, output_file)
        .map_err(|e| format!("Failed to write optimized Llama 3.2 model: {}", e))?;

    Ok(())
}

/// This is the actual function exposed directly to Python
#[pyfunction]
fn compress_safetensors(repo_id: String, precision: String) -> PyResult<bool> {
    // Fake directory path routing for local example layout
    let input_file = format!("./models/{}/model.safetensors", repo_id);
    let output_file = format!("./models/{}/model_{}.safetensors", repo_id, precision.to_lowercase());

    println!("[Rust Engine] Quantizing Llama 3.2 weights for: {}", repo_id);

    match perform_quantization(&input_file, &output_file, &precision) {
        Ok(_) => {
            println!("[Rust Engine] Success. Optimized file written to {}", output_file);
            Ok(true)
        },
        Err(err_msg) => {
            // Converts native Rust system errors into explicit Python exceptions
            Err(PyRuntimeError::new_err(err_msg))
        }
    }
}

/// Defines the Python library entry point structures
#[pymodule]
fn llama32_rust_core(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(compress_safetensors, m)?)?;
    Ok(())
}