use crate::{
    channel::{AwgnChannel, Channel},
    decoder::Decoder,
    encoder::Encoder,
    Gf2, Result,
};
use image::{ImageBuffer, Rgb};
use indicatif::{ProgressBar, ProgressStyle};
use rayon::prelude::*;
use std::time::Instant;

// Convertit un tableau d'octets en un flux de bits
pub fn bytes_to_bits(bytes: &[u8]) -> Vec<Gf2> {
    let mut bits = Vec::with_capacity(bytes.len() * 8);
    for &byte in bytes {
        for i in (0..8).rev() {
            bits.push((byte >> i) & 1);
        }
    }
    bits
}

// Convertit un flux de bits en tableau d'octets
pub fn bits_to_bytes(bits: &[Gf2]) -> Vec<u8> {
    let mut bytes = Vec::with_capacity(bits.len() / 8);
    for chunk in bits.chunks(8) {
        let mut byte = 0u8;
        for (i, &bit) in chunk.iter().enumerate() {
            byte |= bit << (7 - i);
        }
        bytes.push(byte);
    }
    bytes
}

// Transmet une image à travers le canal avec codage LDPC
// pub fn transmit_image(
//     input_path: &str,
//     noisy_out_path: &str,
//     decoded_out_path: &str,
//     encoder: &dyn Encoder,
//     decoder: &dyn Decoder,
//     channel: &AwgnChannel,
// ) -> Result<()> {
//     println!("[*] Chargement de l'image : {}", input_path);
//     let img = image::open(input_path)
//         .expect("Erreur de chargement de l'image")
//         .to_rgb8();
//     let (width, height) = img.dimensions();
//     let raw_bytes = img.into_raw();
//
//     let mut bits = bytes_to_bits(&raw_bytes);
//     let original_bit_len = bits.len();
//
//     // Padding
//     let k = encoder.message_len();
//     let remainder = bits.len() % k;
//     if remainder != 0 {
//         bits.resize(bits.len() + (k - remainder), 0);
//     }
//
//     let num_blocks = bits.len() / k;
//     println!("    - Taille: {}x{} pixels", width, height);
//     println!("    - Blocs à transmettre (k={}): {}", k, num_blocks);
//
//     let mut rng = rand::rngs::StdRng::seed_from_u64(42);
//
//     let mut noisy_bits = Vec::with_capacity(num_blocks * k);
//     let mut decoded_bits = Vec::with_capacity(num_blocks * k);
//
//     let mut frame_errors = 0;
//
//     println!("[*] Transmission et Décodage en cours...");
//
//     for (i, block) in bits.chunks(k).enumerate() {
//         if i % 100 == 0 && i > 0 {
//             println!("    - Progession: {} / {} blocs...", i, num_blocks);
//         }
//
//         let codeword = encoder.encode(block)?;
//
//         let rx_llr = channel.transmit(&codeword, &mut rng);
//
//         // Sans correction LDPC
//         let hard_codeword: Vec<Gf2> = rx_llr
//             .iter()
//             .map(|&l| if l < 0.0 { 1 } else { 0 })
//             .collect();
//         let noisy_block = encoder.extract_message(&hard_codeword);
//         noisy_bits.extend_from_slice(&noisy_block);
//
//         // Décodage LDPC
//         let res = decoder.decode(&rx_llr);
//         if let Some(decoded_codeword) = res.codeword() {
//             let decoded_msg = encoder.extract_message(decoded_codeword);
//             decoded_bits.extend_from_slice(&decoded_msg);
//
//             if decoded_msg != block {
//                 frame_errors += 1;
//             }
//         } else {
//             decoded_bits.extend_from_slice(&noisy_block);
//             frame_errors += 1;
//         }
//     }
//
//     println!(
//         "[*] Transmission terminée. FER : {:.2}%",
//         (frame_errors as f64 / num_blocks as f64) * 100.0
//     );
//
//     // Suppression du padding
//     noisy_bits.truncate(original_bit_len);
//     decoded_bits.truncate(original_bit_len);
//
//     // Reconstitution des images
//     let noisy_bytes = bits_to_bytes(&noisy_bits);
//     let decoded_bytes = bits_to_bytes(&decoded_bits);
//
//     let noisy_img = ImageBuffer::<Rgb<u8>, _>::from_raw(width, height, noisy_bytes)
//         .expect("Erreur de reconstruction de l'image bruitée");
//     noisy_img.save(noisy_out_path).unwrap();
//
//     let decoded_img = ImageBuffer::<Rgb<u8>, _>::from_raw(width, height, decoded_bytes)
//         .expect("Erreur de reconstruction de l'image décodée");
//     decoded_img.save(decoded_out_path).unwrap();
//
//     println!(
//         "[*] Images sauvegardées : {} et {}",
//         noisy_out_path, decoded_out_path
//     );
//     Ok(())
// }
pub fn transmit_image(
    input_path: &str,
    noisy_out_path: &str,
    decoded_out_path: &str,
    encoder: &dyn Encoder,
    decoder: &dyn Decoder,
    channel: &AwgnChannel,
) -> Result<()> {
    println!("\n[*] Chargement de l'image : {}", input_path);
    let img = image::open(input_path)
        .expect("Erreur : Impossible de trouver ou lire l'image")
        .to_rgb8();
    let (width, height) = img.dimensions();
    let raw_bytes = img.into_raw();

    let mut bits = bytes_to_bits(&raw_bytes);
    let original_bit_len = bits.len();

    let k = encoder.message_len();
    let remainder = bits.len() % k;
    if remainder != 0 {
        bits.resize(bits.len() + (k - remainder), 0);
    }

    let num_blocks = bits.len() / k;
    println!("    - Résolution : {}x{} pixels", width, height);
    println!(
        "    - Poids total : {:.2} Mo",
        raw_bytes.len() as f64 / 1_048_576.0
    );
    println!("    - Découpage en : {} blocs de {} bits\n", num_blocks, k);

    let start_time = Instant::now();

    let pb = ProgressBar::new(num_blocks as u64);
    pb.set_style(
        ProgressStyle::default_bar()
            .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} blocs ({eta})")
            .unwrap()
            .progress_chars("=>-")
    );

    let results: Vec<_> = bits
        .par_chunks(k)
        .map(|block| {
            let mut rng = rand::thread_rng();
            let codeword = encoder.encode(block).unwrap();
            let rx_llr = channel.transmit(&codeword, &mut rng);
            let hard_codeword: Vec<Gf2> = rx_llr
                .iter()
                .map(|&l| if l < 0.0 { 1 } else { 0 })
                .collect();
            let noisy_block = encoder.extract_message(&hard_codeword);
            let res = decoder.decode(&rx_llr);

            let (decoded_block, has_error) = if let Some(decoded_codeword) = res.codeword() {
                let decoded_msg = encoder.extract_message(decoded_codeword);
                let err = if decoded_msg != block { 1 } else { 0 };
                (decoded_msg, err)
            } else {
                (noisy_block.clone(), 1)
            };
            pb.inc(1);
            (noisy_block, decoded_block, has_error)
        })
        .collect();

    pb.finish_with_message("Décodage terminé !");

    let mut noisy_bits = Vec::with_capacity(num_blocks * k);
    let mut decoded_bits = Vec::with_capacity(num_blocks * k);
    let mut frame_errors = 0;

    for (n_block, d_block, err) in results {
        noisy_bits.extend_from_slice(&n_block);
        decoded_bits.extend_from_slice(&d_block);
        frame_errors += err;
    }

    let duration = start_time.elapsed();
    let fer = (frame_errors as f64 / num_blocks as f64) * 100.0;
    println!("[*] Transmission terminée en {:.2?}", duration);
    println!("    - Taux d'Erreur Trame (FER) : {:.2}%", fer);

    noisy_bits.truncate(original_bit_len);
    decoded_bits.truncate(original_bit_len);

    let noisy_bytes = bits_to_bytes(&noisy_bits);
    let decoded_bytes = bits_to_bytes(&decoded_bits);

    let noisy_img = ImageBuffer::<Rgb<u8>, _>::from_raw(width, height, noisy_bytes)
        .expect("Erreur de reconstruction de l'image bruitée");
    noisy_img.save(noisy_out_path).unwrap();

    let decoded_img = ImageBuffer::<Rgb<u8>, _>::from_raw(width, height, decoded_bytes)
        .expect("Erreur de reconstruction de l'image décodée");
    decoded_img.save(decoded_out_path).unwrap();

    println!(
        "[*] Succès ! Images sauvegardées : {} et {}",
        noisy_out_path, decoded_out_path
    );
    Ok(())
}
