use pnet::transport::{transport_channel, TransportChannelType};
use pnet::packet::ip::IpNextHeaderProtocols;
use pnet::packet::tcp::MutableTcpPacket;
use pnet::packet::ipv4::{MutableIpv4Packet, Ipv4Packet};
use pnet::packet::Packet;
use rand::Rng;
use std::net::{Ipv4Addr, IpAddr};

use super::console::fatal;

pub struct Reflector {
    pub ip: String,
    pub port: u16,
}

impl Reflector {
    pub fn from_str(refstr: &str) -> Result<Self, String> {
        let parts: Vec<&str> = refstr.split(':').collect();
        let ip = parts[0].to_string();
        let port: u16 = parts[1].parse().unwrap_or(0);

        if port == 0 {
            return Err(format!("invalid port \"{}\" in {}", parts[1], refstr));
        }

        Ok(Self {ip, port})
    }
}

pub fn random_field() -> u32 {
    let mut rng = rand::thread_rng();
    rng.gen()
}

pub fn mkpacket4(saddr: &str, sport: u16, daddr: &str, dport: u16, seed: u32) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
    let mut ipbuf = vec![0u8; 40];
    let mut iph = MutableIpv4Packet::new(&mut ipbuf).ok_or_else(|| "unable to create ipv4 packet")?;

    let saddr: Ipv4Addr = saddr.parse()?;
    let daddr: Ipv4Addr = daddr.parse()?;
    
    iph.set_version(4);
    iph.set_header_length(5);
    iph.set_total_length(40);
    iph.set_identification(seed as u16);
    iph.set_flags(0);
    iph.set_ttl(255);
    iph.set_next_level_protocol(IpNextHeaderProtocols::Tcp);
    iph.set_source(saddr);
    iph.set_destination(daddr);
    
    let mut tcpbuf = vec![0u8; 20];
    let mut tcph = MutableTcpPacket::new(&mut tcpbuf).unwrap();
    
    tcph.set_source(sport);
    tcph.set_destination(dport);
    tcph.set_sequence(seed);
    tcph.set_acknowledgement(0);
    tcph.set_data_offset(5);
    tcph.set_flags(10); // SYN + PSH
    tcph.set_window(65535);
    tcph.set_urgent_ptr(0);
    
    let checksum = pnet::packet::tcp::ipv4_checksum(&tcph.to_immutable(), &saddr, &daddr);
    tcph.set_checksum(checksum);
    
    let mut packet = vec![0u8; 40];
    packet[..20].copy_from_slice(&iph.packet()[..20]);
    packet[20..].copy_from_slice(tcph.packet());
    
    Ok(packet)
}

pub fn send(packet_bytes: Vec<u8>) -> Result<(), Box<dyn std::error::Error>> {    
    let packet = Ipv4Packet::owned(packet_bytes).ok_or_else(|| "error reading packet bytes back into type")?;
    
    let (mut sender, _) = transport_channel(packet.packet().len(), TransportChannelType::Layer3(IpNextHeaderProtocols::Tcp))?;

    sender.send_to(&packet, IpAddr::V4(packet.get_destination()))?; // could check if # bytes sent matches packet len for integrity
    Ok(())
}

pub fn transmit(source: &str, port: u16, r: Reflector, seed: u32) {
    if let Ok(pkt) = mkpacket4(source, port, &r.ip, r.port, seed) {
        match send(pkt) {
            Ok(_) => {
                // println!("transmitted to reflector {}:{}", r.ip, r.port);
            }
            Err(e) => {
                if e.to_string().contains("permitted") {
                    fatal("unable to transmit over raw socket, are you root?".to_string());
                }
                println!("error transmitting to reflector {}:{}: {}", r.ip, r.port, e);
            },
        }
    } else {
        println!("error creating packet for reflector {}:{}", r.ip, r.port);
    }
}