use std::net::{SocketAddr, UdpSocket};
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Barrier, Mutex, RwLock};
use std::thread;
use std::time::{Duration, Instant};

use clap::Parser;
use socket2::{Domain, Socket, Type};

mod bouncing_image;
mod circle;
mod color;
mod display;
mod drawable;
mod pixel_buf;
mod png_data;
mod raw_socket;
mod pixel;
mod gif_image;

use bouncing_image::BouncingImage;
use circle::Circle;
use display::Display;
use drawable::Drawable;
use gif_image::GifImage;
use pixel_buf::PixelBuf;

pub const DISPLAY_WIDTH: i32 = 1920;
pub const DISPLAY_HEIGHT: i32 = 1080;

#[derive(Parser)]
#[command(name = "flood-rs", about = "Pixel flooding tool")]
struct Args {
    /// Target IP address or hostname
    target: String,

    /// Target port
    #[arg(short, long, default_value_t = 5005)]
    port: u16,

    /// Source network interface (e.g. eth0). Auto-discovered if not set.
    #[arg(short = 'i', long)]
    interface: Option<String>,
}

/// Unpacks a 4-byte slice into two u16 values (little-endian).
fn unpack_coordinates(buffer: &[u8]) -> Option<(u16, u16)> {
    if buffer.len() != 4 {
        return None;
    }
    let x_bytes: [u8; 2] = buffer[0..2].try_into().ok()?;
    let y_bytes: [u8; 2] = buffer[2..4].try_into().ok()?;
    let x = u16::from_le_bytes(x_bytes);
    let y = u16::from_le_bytes(y_bytes);
    Some((x, y))
}

fn main() {
    let args = Args::parse();

    let x: Arc<Mutex<u32>> = Arc::new(Mutex::new(0));
    let x_thread = x.clone();

    let y: Arc<Mutex<u32>> = Arc::new(Mutex::new(0));
    let y_thread = y.clone();

    let set: Arc<Mutex<bool>> = Arc::new(Mutex::new(false));
    let set_thread = set.clone();

    let circle = Box::new(Circle::new(x, y, set));
    let mut images: Vec<Box<dyn Drawable>> = vec![
        // Velocities are in pixels per second.
//        Box::new(BouncingImage::new("images/unicorn_cc.png", 30.0, -30.0, -1, -1)),
//        Box::new(BouncingImage::new("images/windows_logo.png", -20.0, 20.0, -1, -1)),
//        Box::new(BouncingImage::new("images/spade.png", 90.0, -60.0, 0, 0)),
//        Box::new(BouncingImage::new("images/dvdvideo.png", 60.0, 18.0, 1000, 800)),
//        Box::new(BouncingImage::new("images/hackaday.png", 40.0, 50.0, 500, 800)),
        circle,
        Box::new(GifImage::new("images/netto.gif", 200, 200)),
    ];

    let mut display = Display::new(&args.target, args.port, args.interface.as_deref());

    // Spawn a UDP listener thread for receiving coordinates
    thread::spawn(move || {
        let bind_address = "0.0.0.0:12345";
        let socket = Socket::new(Domain::IPV4, Type::DGRAM, None).unwrap();
        socket.set_reuse_address(true).unwrap();
        socket
            .bind(&"0.0.0.0:1234".parse::<SocketAddr>().unwrap().into())
            .unwrap();
        let socket: UdpSocket = socket.into();

        println!("Listening for UDP packets on {}", bind_address);

        let mut buf = [0u8; 4];

        loop {
            match socket.recv_from(&mut buf) {
                Ok((number_of_bytes, src_addr)) => {
                    println!("\nReceived {} bytes from {}", number_of_bytes, src_addr);

                    if number_of_bytes == 4 {
                        if let Some((x_rev, y_rev)) = unpack_coordinates(&buf) {
                            println!("Received Coordinates: X = {}, Y = {}", x_rev, y_rev);
                            *x_thread.lock().unwrap() = x_rev.into();
                            *y_thread.lock().unwrap() = y_rev.into();
                            *set_thread.lock().unwrap() = true;
                        } else {
                            eprintln!("Error: Failed to unpack coordinate data.");
                        }
                    } else {
                        eprintln!(
                            "Warning: Received packet with incorrect size ({} bytes). Expected 4.",
                            number_of_bytes
                        );
                    }
                }
                Err(e) => {
                    eprintln!("Error receiving data: {}", e);
                }
            }
        }
    });

    let num_drawables = images.len();

    // Wrap drawables in Arc<RwLock<>> for shared access:
    // - Workers take read locks for draw(&self) -- parallel, no contention
    // - Main takes write locks for update(&mut self) -- exclusive, 30fps only
    let images: Vec<Arc<RwLock<Box<dyn Drawable>>>> = images
        .into_iter()
        .map(|b| Arc::new(RwLock::new(b)))
        .collect();

    // Shared elapsed time for draw threads (as u32 bits of f32).
    let elapsed_bits = Arc::new(AtomicU32::new(0));

    // Barriers: main signals "go", workers draw, then signal "done".
    let go_barrier = Arc::new(Barrier::new(num_drawables + 1));
    let done_barrier = Arc::new(Barrier::new(num_drawables + 1));

    // Each worker gets its own PixelBuf wrapped in Mutex so the main
    // thread can read the results after the done barrier.
    let pixel_bufs: Vec<Arc<Mutex<PixelBuf>>> = (0..num_drawables)
        .map(|_| Arc::new(Mutex::new(PixelBuf::new(60_000))))
        .collect();

    // Spawn persistent worker threads.
    for i in 0..num_drawables {
        let img = Arc::clone(&images[i]);
        let buf = Arc::clone(&pixel_bufs[i]);
        let elapsed = Arc::clone(&elapsed_bits);
        let go = Arc::clone(&go_barrier);
        let done = Arc::clone(&done_barrier);

        thread::spawn(move || {
            loop {
                go.wait();

                let e = f32::from_bits(elapsed.load(Ordering::Relaxed));
                let drawable = img.read().unwrap();
                let mut pb = buf.lock().unwrap();
                pb.clear();
                drawable.draw(&mut pb, e);
                drop(pb);
                drop(drawable);

                done.wait();
            }
        });
    }

    let start_time = Instant::now();
    let mut last_update = Instant::now();
    let update_interval = Duration::from_secs_f32(1.0 / 30.0);

    loop {
        let now = Instant::now();
        let elapsed = (now - start_time).as_secs_f32();

        // Update positions at 30 fps (main thread only, workers are idle).
        if now - last_update >= update_interval {
            let dt = (now - last_update).as_secs_f32();
            for img in &images {
                img.write().unwrap().update(dt, elapsed);
            }
            last_update = now;
        }

        // Publish elapsed time and signal workers to draw.
        elapsed_bits.store(elapsed.to_bits(), Ordering::Relaxed);
        go_barrier.wait();

        // Wait for all workers to finish drawing.
        done_barrier.wait();

        // Merge all pixel bufs into the single Display for efficient packing.
        for pb in &pixel_bufs {
            let buf = pb.lock().unwrap();
            display.write_raw_pixels(&buf.pixels);
        }
        display.flush_frame();
    }
}
