use std::fs::File;
use std::io::BufReader;
use std::net::{ToSocketAddrs, UdpSocket};
use std::time::Duration;

// Constants from the C code
const QUEUE_LEN: usize = 1000;
const MSG_PAYLOAD_SIZE: usize = 7 * 160;
const MSGSIZE: usize = 2 + MSG_PAYLOAD_SIZE;

const DISPLAY_HOST: &str = "100.65.0.2";
const DISPLAY_PORT: u16 = 5005;
const DISPLAY_WIDTH: i32 = 1920;
const DISPLAY_HEIGHT: i32 = 1080;

/// Represents the data decoded from a PNG file.
struct PngData {
    width: u32,
    height: u32,
    pixels: Vec<u8>,
}

impl PngData {
    /// Loads and decodes a PNG image from the given path.
    fn open(path: &str) -> Result<Self, png::DecodingError> {
        let file = File::open(path).expect("Failed to open PNG file");
        let decoder = png::Decoder::new(BufReader::new(file));
        let mut reader = decoder.read_info()?;
        let mut buf = vec![0; reader.output_buffer_size()];
        let info = reader.next_frame(&mut buf)?;

        Ok(PngData {
            width: info.width,
            height: info.height,
            pixels: buf,
        })
    }
}

/// Represents a bouncing image on the screen.
struct BouncingImage {
    img: PngData,
    x: i32,
    y: i32,
    x1: i32,
    y1: i32,
    x2: i32,
    y2: i32,
    move_x: i32,
    move_y: i32,
    rate: u32,
}

impl BouncingImage {
    /// Initializes a new BouncingImage.
    fn new(img_file: &str, move_x: i32, move_y: i32, rate: u32, start_x: i32, start_y: i32) -> Self {
        let img = PngData::open(img_file).expect("Could not load image");
        let mut bb = BouncingImage {
            x2: DISPLAY_WIDTH - img.width as i32,
            y2: DISPLAY_HEIGHT - img.height as i32,
            img,
            x: start_x,
            y: start_y,
            x1: 0,
            y1: 0,
            move_x,
            move_y,
            rate,
        };
        if bb.x == -1 {
            bb.x = (bb.x1 + bb.x2) / 2;
        }
        if bb.y == -1 {
            bb.y = (bb.y1 + bb.y2) / 2;
        }
        bb
    }

    /// Draws the image and updates its position.
    fn draw_and_move(&mut self, display: &mut Display) {
        display.draw_png(&self.img, self.x, self.y);

        self.x += self.move_x;
        self.y += self.move_y;

        if self.x < self.x1 || self.x > self.x2 {
            self.move_x *= -1;
        }
        if self.y < self.y1 || self.y > self.y2 {
            self.move_y *= -1;
        }
    }
}

/// Manages the connection and data sent to the display.
struct Display {
    socket: UdpSocket,
    bufs: Vec<[u8; MSGSIZE]>,
    next_buf: usize,
    pos_in_buf: usize,
}

impl Display {
    /// Creates a new Display and connects to the specified host and port.
    fn new(host: &str, port: u16) -> Self {
        let remote_addr = (host, port)
            .to_socket_addrs()
            .expect("Invalid remote address")
            .next()
            .expect("Could not resolve host");

        let socket = UdpSocket::bind("0.0.0.0:0").expect("Could not bind to local port");
        socket.connect(remote_addr).expect("Could not connect to remote");

        let mut bufs = vec![[0; MSGSIZE]; QUEUE_LEN];
        for buf in bufs.iter_mut() {
            buf[0] = 0x00;
            buf[1] = 0x01;
        }

        Display {
            socket,
            bufs,
            next_buf: 0,
            pos_in_buf: 0,
        }
    }

    /// Flushes the current buffer if it contains pixel data.
    fn flush_frame(&mut self) {
        if self.pos_in_buf > 0 {
            let len = 2 + self.pos_in_buf * 7;
            let buf_to_send = &self.bufs[self.next_buf][..len];
            self.socket.send(buf_to_send).expect("Failed to send data");
            self.next_buf = (self.next_buf + 1) % QUEUE_LEN;
            self.pos_in_buf = 0;
        }
    }

    /// Sets a pixel color at a specific coordinate.
    fn set_pixel(&mut self, x: u16, y: u16, r: u8, g: u8, b: u8) {
        let offset = 2 + self.pos_in_buf * 7;
        let buf = &mut self.bufs[self.next_buf][offset..offset + 7];
        buf[0] = x as u8;
        buf[1] = (x >> 8) as u8;
        buf[2] = y as u8;
        buf[3] = (y >> 8) as u8;
        buf[4] = r;
        buf[5] = g;
        buf[6] = b;

        self.pos_in_buf += 1;
        if self.pos_in_buf == 160 {
            self.flush_frame();
        }
    }
    
    /// Draws a PNG image at the given coordinates.
    fn draw_png(&mut self, png: &PngData, x: i32, y: i32) {
        for sy in 0..png.height {
            for sx in 0..png.width {
                let index = (sy * png.width + sx) as usize * 4;
                let rgba = &png.pixels[index..index + 4];
                if rgba[3] > 0 { // Check alpha channel
                    self.set_pixel((x + sx as i32) as u16, (y + sy as i32) as u16, rgba[0], rgba[1], rgba[2]);
                }
            }
        }
    }

    /// Clears the entire screen to black.
    #[allow(dead_code)]
    fn blank_screen(&mut self) {
        for x in 0..DISPLAY_WIDTH {
            for y in 0..DISPLAY_HEIGHT {
                self.set_pixel(x as u16, y as u16, 0, 0, 0);
            }
        }
        self.flush_frame();
    }
}

fn main() {
    let mut images = vec![
        BouncingImage::new("images/unicorn_cc.png", 13, -10, 1, -1, -1),
        BouncingImage::new("images/windows_logo.png", -8, 3, 2, -1, -1),
        BouncingImage::new("images/spade.png", 32, -12, 1, 0, 0),
        BouncingImage::new("images/dvdvideo.png", 20, 6, 5, 1000, 800),
        BouncingImage::new("images/hackaday.png", 40, 18, 3, 500, 800),
    ];

    let mut display = Display::new(DISPLAY_HOST, DISPLAY_PORT);
    let mut frame_counter: u32 = 0;

    // display.blank_screen();

    loop {
        for (i, bb) in images.iter_mut().enumerate() {
            if bb.rate > 0 && frame_counter % bb.rate != 0 {
                continue;
            }
            bb.draw_and_move(&mut display);
        }
        display.flush_frame();
        frame_counter += 1;
        
        // A small delay to control the frame rate
        std::thread::sleep(Duration::from_millis(16));
    }
}