use std::{io, net::UdpSocket};

use image::{Rgb, RgbImage};
use thiserror::Error;

use crate::{Layout, Strandifier, StrandifierError};

#[derive(Debug, Error)]
#[allow(clippy::module_name_repetitions)]
pub enum SenderError {
    #[error("IO error")]
    Io(#[from] io::Error),
    #[error("Strandifier failed")]
    Strandifier(#[from] StrandifierError),
    #[error("Controller failed to confirm frame")]
    ConfirmationTimeout,
    #[error("Invalid confirmation packet: {0:?}")]
    ConfirmationInvalid(Vec<u8>),

    #[error("Confirmation packet for wrong frame number: expected {expected}, got {got}")]
    ConfirmationMismatched { expected: u32, got: u32 },
}

/// Sends a strand's data to the controller.
///
/// # Errors
///
/// Returns an error if network communication fails.
pub fn send_strand<I: IntoIterator<Item = Rgb<u8>>>(
    socket: &UdpSocket,
    strand_num: u32,
    frame_num: u32,
    pixels: I,
) -> Result<(), SenderError> {
    const MAGIC: u32 = u32::from_be_bytes(*b"PIXL");

    let pixels = pixels.into_iter().map(|c| {
        let mut bytes = [0; 4];
        bytes[1..].copy_from_slice(&c.0);

        u32::from_be_bytes(bytes)
    });

    let buf: Vec<_> = [MAGIC, strand_num, frame_num]
        .into_iter()
        .chain(pixels)
        .flat_map(u32::to_be_bytes)
        .collect();

    socket.send(&buf)?;

    Ok(())
}

/// Sends an entire image frame to the controller.
///
/// # Errors
///
/// Returns an error if the image does not have the correct dimensions, if network communication
/// fails or if the controller does not confirm reception of the frame.
pub fn send_frame(
    socket: &UdpSocket,
    layout: Layout,
    frame_num: u32,
    image: &RgbImage,
) -> Result<(), SenderError> {
    for strand_num in 0..layout.total_strands {
        if strand_num >= layout.first_strand_index
            && strand_num < layout.num_strands() + layout.first_strand_index
        {
            let data =
                Strandifier::make_strand(layout, image, strand_num - layout.first_strand_index)?;
            send_strand(socket, strand_num as u32, frame_num, data)?;
        } else {
            send_strand(
                socket,
                strand_num as u32,
                frame_num,
                vec![Rgb([0, 0, 0]); layout.strand_len() as usize],
            )?;
        }
    }

    let mut buf = vec![0; 100];
    let len = socket
        .recv(&mut buf)
        .map_err(|_| SenderError::ConfirmationTimeout)?;
    let buf = &buf[..len];
    let response_frame_num = u32::from_be_bytes(
        buf.try_into()
            .map_err(|_| SenderError::ConfirmationInvalid(buf.to_vec()))?,
    );

    if response_frame_num != frame_num {
        return Err(SenderError::ConfirmationMismatched {
            expected: frame_num,
            got: response_frame_num,
        });
    }

    Ok(())
}