use anyhow::{Context, Result};
use clap::Parser;
use log::{error, info};
use socket2::{Domain, Protocol, Socket, Type};
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::str::FromStr;
use std::time::{Duration, Instant};
use tokio::net::UdpSocket;

#[derive(Parser, Debug)]
#[command(author, version, about = "Multicast packet generator for testing CastRepeat")]
struct Args {
    #[arg(short, long, default_value = "239.192.55.1")]
    multicast_addr: String,

    #[arg(short, long, default_value = "1681")]
    port: u16,

    #[arg(short, long, default_value = "1000")]
    interval_ms: u64,

    #[arg(short, long, default_value = "60")]
    duration_sec: u64,

    #[arg(short, long, default_value = "Test packet")]
    message: String,
    
    #[arg(short, long)]
    interface: Option<String>,
    
    #[arg(short = 'b', long, action)]
    binary_mode: bool,
}

#[tokio::main]
async fn main() -> Result<()> {
    env_logger::init();
    let args = Args::parse();

    info!("CastRepeat Multicast Packet Generator");
    info!("--------------------------------");
    info!("Multicast Address: {}", args.multicast_addr);
    info!("Port:              {}", args.port);
    info!("Interval:          {} ms", args.interval_ms);
    info!("Duration:          {} sec", args.duration_sec);
    if let Some(interface) = &args.interface {
        info!("Interface:         {}", interface);
    }
    info!("Mode:              {}", if args.binary_mode { "Binary" } else { "Text" });
    info!("--------------------------------");

    // Verify multicast address
    let mcast_addr = match IpAddr::from_str(&args.multicast_addr) {
        Ok(IpAddr::V4(addr)) if addr.is_multicast() => addr,
        _ => {
            error!("Invalid multicast address: {}", args.multicast_addr);
            return Ok(());
        }
    };

    // Create socket
    let socket = Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::UDP))
        .context("Failed to create socket")?;

    socket.set_multicast_ttl_v4(4)?;
    socket.set_nonblocking(true)?;
    
    // Set the multicast interface if specified
    if let Some(if_str) = &args.interface {
        if let Ok(if_addr) = Ipv4Addr::from_str(if_str) {
            socket.set_multicast_if_v4(&if_addr)?;
            info!("Using interface: {}", if_addr);
        } else {
            error!("Invalid interface address: {}", if_str);
            return Ok(());
        }
    }

    let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0);
    socket.bind(&addr.into())?;

    let socket = UdpSocket::from_std(socket.into())?;
    let dest_addr = SocketAddr::new(IpAddr::V4(mcast_addr), args.port);

    // Start sending packets
    info!("Sending packets to {}...", dest_addr);

    let start = Instant::now();
    let end = start + Duration::from_secs(args.duration_sec);
    let mut counter: u64 = 0;  // Specify counter type as u64

    while Instant::now() < end {
        counter += 1;
        
        // Create either text or binary test data
        let data = if args.binary_mode {
            // Create binary test data (similar to what we might see in the field)
            let mut packet = Vec::with_capacity(16);
            packet.extend_from_slice(b"REL\0");          // 4 bytes header
            packet.extend_from_slice(&counter.to_be_bytes()[4..]); // 4 bytes counter 
            packet.extend_from_slice(&[0x00, 0x10, 0x02, 0x00]); // 4 bytes
            packet.extend_from_slice(&[0x00, 0x00, 0x00, 0x90]); // 4 bytes
            packet
        } else {
            // Create text test data
            format!("{} #{}", args.message, counter).into_bytes()
        };

        match socket.send_to(&data, &dest_addr).await {
            Ok(bytes) => {
                info!("Sent packet #{}: {} bytes", counter, bytes);
            }
            Err(e) => {
                error!("Error sending packet: {}", e);
            }
        }

        tokio::time::sleep(Duration::from_millis(args.interval_ms)).await;
    }

    info!("Done. Sent {} packets in {} seconds.", counter, args.duration_sec);
    Ok(())
}