use std::fs::{File, OpenOptions};
use std::io::{Read, Write};
use std::net::{SocketAddr, UdpSocket};
use std::time::Duration;
use std::{env, process, time};

use dhcp::{
    Dhcp, DHCPACK, DHCPDISCOVER, DHCPNAK, DHCPOFFER, DHCPREQUEST, DHCPRELEASE,
    OPT_DNS, OPT_LEASE_TIME, OPT_MESSAGE_TYPE, OPT_REQUESTED_IP, OPT_ROUTER,
    OPT_SERVER_ID, OPT_SUBNET_MASK, OPT_END,
};

mod dhcp;

macro_rules! try_fmt {
    ($e:expr, $m:expr) => {
        match $e {
            Ok(ok) => ok,
            Err(err) => return Err(format!("{}: {}", $m, err)),
        }
    };
}

fn get_cfg_value(path: &str) -> Result<String, String> {
    let path = format!("/scheme/netcfg/{path}");
    let mut file = File::open(&path).map_err(|_| format!("Can't open {path}"))?;
    let mut result = String::new();
    file.read_to_string(&mut result)
        .map_err(|_| format!("Can't read {path}"))?;
    Ok(result)
}

fn get_iface_cfg_value(iface: &str, cfg: &str) -> Result<String, String> {
    let path = format!("ifaces/{iface}/{cfg}");
    get_cfg_value(&path)
}

fn set_cfg_value(path: &str, value: &str) -> Result<(), String> {
    let path = format!("/scheme/netcfg/{path}");
    let mut file = OpenOptions::new()
        .read(false)
        .write(true)
        .create(false)
        .open(&path)
        .map_err(|_| format!("Can't open {path}"))?;
    file.write(value.as_bytes())
        .map(|_| ())
        .map_err(|_| format!("Can't write {value} to {path}"))?;
    file.sync_data()
        .map_err(|_| format!("Can't commit {value} to {path}"))
}

fn set_iface_cfg_value(iface: &str, cfg: &str, value: &str) -> Result<(), String> {
    let path = format!("ifaces/{iface}/{cfg}");
    set_cfg_value(&path, value)
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Default)]
struct MacAddr {
    bytes: [u8; 6],
}

impl MacAddr {
    fn from_str(string: &str) -> Self {
        MacAddr::try_parse_with_delimeter(string, ':')
            .or_else(|| MacAddr::try_parse_with_delimeter(string, '-'))
            .unwrap_or_default()
    }

    fn try_parse_with_delimeter(string: &str, delimeter: char) -> Option<MacAddr> {
        let mut addr = MacAddr::default();
        let mut segments = 0;

        for part in string.split(delimeter) {
            if segments >= addr.bytes.len() {
                return None;
            }
            addr.bytes[segments] = match u8::from_str_radix(part, 16) {
                Ok(b) => b,
                _ => return None,
            };
            segments += 1;
        }

        if segments == addr.bytes.len() {
            Some(addr)
        } else {
            None
        }
    }

    fn to_string(&self) -> String {
        format!(
            "{:>02X}-{:>02X}-{:>02X}-{:>02X}-{:>02X}-{:>02X}",
            self.bytes[0],
            self.bytes[1],
            self.bytes[2],
            self.bytes[3],
            self.bytes[4],
            self.bytes[5]
        )
    }
}

fn dhcp(iface: &str, verbose: bool) -> Result<(), String> {
    let current_mac = MacAddr::from_str(get_iface_cfg_value(iface, "mac")?.trim());

    let _current_ip = get_iface_cfg_value(iface, "addr/list")?
        .lines()
        .next()
        .map(|l| l.to_owned())
        .unwrap_or("0.0.0.0".to_string());

    if verbose {
        println!("DHCP: MAC: {} Starting", current_mac.to_string());
    }

    let tid = try_fmt!(
        time::SystemTime::now().duration_since(time::UNIX_EPOCH),
        "failed to get time"
    ).subsec_nanos();

    let socket = try_fmt!(UdpSocket::bind(("0.0.0.0", 68)), "failed to bind udp");
    let broadcast = SocketAddr::from(([255, 255, 255, 255], 67));
    // 8s (was 30s): a DHCP server that is present answers a DISCOVER in well
    // under a second, so a long read timeout only serves to stall boot when no
    // server responds (e.g. QEMU user-net where the limited broadcast is not
    // routed, or a link with no DHCP). Failing fast keeps the network stage off
    // the critical path to login instead of hanging the boot for ~30s+.
    //
    // Do NOT call connect() to the broadcast address. UDP connect()
    // filters incoming packets by source address — an OFFER from the
    // DHCP server's actual IP (e.g. QEMU SLIRP at 10.0.2.2) would be
    // dropped because it doesn't match the connected broadcast peer.
    // Use send_to() per-packet and let recv() accept any source.
    try_fmt!(socket.set_read_timeout(Some(Duration::new(8, 0))), "failed to set read timeout");
    try_fmt!(socket.set_write_timeout(Some(Duration::new(8, 0))), "failed to set write timeout");

    let mut subnet_option: Option<[u8; 4]> = None;
    let mut router_option: Option<[u8; 4]> = None;
    let mut dns_option: Option<[u8; 4]> = None;
    let mut server_id_option: Option<[u8; 4]> = None;
    let mut lease_time_secs: u32 = 86400;

    // DHCPDISCOVER
    {
        let mut discover = Dhcp::default();
        init_dhcp_header(&mut discover, current_mac, tid);
        let disc_opts: &[u8] = &[OPT_MESSAGE_TYPE, 1, DHCPDISCOVER, OPT_END];
        discover.options[..disc_opts.len()].copy_from_slice(disc_opts);
        try_fmt!(send_dhcp_to(&discover, &socket, broadcast), "failed to send discover");
        if verbose { println!("DHCP: Sent Discover"); }
    }

    // Recv DHCPOFFER
    let mut offer_data = [0; 65536];
    try_fmt!(socket.recv(&mut offer_data), "failed to receive offer");
    let offer = unsafe { &*(offer_data.as_ptr() as *const Dhcp) };
    if verbose { println!("DHCP: Offer IP: {:?}", offer.yiaddr); }

    parse_options(&offer.options, &mut |code, data| {
        match code {
            OPT_SUBNET_MASK if data.len() == 4 && subnet_option.is_none() => {
                subnet_option = Some([data[0], data[1], data[2], data[3]]);
            }
            OPT_ROUTER if data.len() == 4 && router_option.is_none() => {
                router_option = Some([data[0], data[1], data[2], data[3]]);
            }
            OPT_DNS if data.len() == 4 && dns_option.is_none() => {
                dns_option = Some([data[0], data[1], data[2], data[3]]);
            }
            OPT_LEASE_TIME if data.len() == 4 => {
                lease_time_secs = u32::from_be_bytes([data[0], data[1], data[2], data[3]]);
            }
            OPT_SERVER_ID if data.len() == 4 && server_id_option.is_none() => {
                server_id_option = Some([data[0], data[1], data[2], data[3]]);
            }
            _ => {}
        }
    });

    let mask_len = compute_prefix_len(subnet_option);
    let new_ips = format!(
        "{}.{}.{}.{}/{}\n",
        offer.yiaddr[0], offer.yiaddr[1], offer.yiaddr[2], offer.yiaddr[3], mask_len
    );
    try_fmt!(set_iface_cfg_value(iface, "addr/set", &new_ips), "failed to set ip");
    apply_dhcp_config(iface, router_option, dns_option, verbose)?;

    let server_id = server_id_option.unwrap_or([0, 0, 0, 0]);

    // DHCPREQUEST
    {
        let mut request = Dhcp::default();
        init_dhcp_header(&mut request, current_mac, tid);
        let req_opts: &[u8] = &[
            OPT_MESSAGE_TYPE, 1, DHCPREQUEST,
            OPT_REQUESTED_IP, 4,
            offer.yiaddr[0], offer.yiaddr[1], offer.yiaddr[2], offer.yiaddr[3],
            OPT_SERVER_ID, 4, server_id[0], server_id[1], server_id[2], server_id[3],
            OPT_END,
        ];
        request.options[..req_opts.len()].copy_from_slice(req_opts);
        try_fmt!(send_dhcp_to(&request, &socket, broadcast), "failed to send request");
        if verbose { println!("DHCP: Sent Request"); }
    }

    // Recv DHCPACK
    let mut ack_data = [0; 65536];
    try_fmt!(socket.recv(&mut ack_data), "failed to receive ack");
    if verbose { println!("DHCP: lease acquired, {}s lease time", lease_time_secs); }

    // RFC 2131 lease lifecycle: RENEW at T1, REBIND at T2
    let t1 = Duration::from_secs(lease_time_secs as u64 / 2);
    let t2 = Duration::from_secs((lease_time_secs as u64 * 7) / 8);
    let now = time::Instant::now();
    let t1_deadline = now + t1;

    let mut remaining = t1_deadline.saturating_duration_since(time::Instant::now());
    while remaining > Duration::ZERO {
        std::thread::sleep(std::cmp::min(remaining, Duration::from_secs(60)));
        remaining = t1_deadline.saturating_duration_since(time::Instant::now());
    }

    if verbose { println!("DHCP: entering RENEW state"); }
    {
        let mut renew = Dhcp::default();
        init_dhcp_header(&mut renew, current_mac, tid.wrapping_add(1));
        renew.ciaddr = offer.yiaddr;
        let rn_opts: &[u8] = &[OPT_MESSAGE_TYPE, 1, DHCPREQUEST, OPT_END];
        renew.options[..rn_opts.len()].copy_from_slice(rn_opts);
        try_fmt!(
            send_dhcp_to(&renew, &socket, broadcast),
            "failed to send renew"
        );
    }

    socket.set_read_timeout(Some(t2.saturating_sub(t1))).ok();
    match socket.recv(&mut ack_data) {
        Ok(_) => {
            let response = unsafe { &*(ack_data.as_ptr() as *const Dhcp) };
            match get_message_type(&response.options) {
                Some(DHCPACK) => { if verbose { println!("DHCP: renewed"); } }
                Some(DHCPNAK) => {
                    if verbose { println!("DHCP: NAK, restarting"); }
                    return dhcp(iface, verbose);
                }
                _ => {}
            }
        }
        Err(_) => {
            if verbose { println!("DHCP: entering REBIND state"); }
            let bind_socket = try_fmt!(UdpSocket::bind(("0.0.0.0", 68)), "failed to bind rebind");
            let mut rebind = Dhcp::default();
            init_dhcp_header(&mut rebind, current_mac, tid.wrapping_add(2));
            rebind.ciaddr = offer.yiaddr;
            let rb_opts: &[u8] = &[OPT_MESSAGE_TYPE, 1, DHCPREQUEST, OPT_END];
            rebind.options[..rb_opts.len()].copy_from_slice(rb_opts);
            let _ = send_dhcp_to(&rebind, &bind_socket, broadcast);
            bind_socket.set_read_timeout(Some(Duration::from_secs(10))).ok();
            if let Ok(_) = bind_socket.recv(&mut ack_data) {
                if verbose { println!("DHCP: rebound"); }
            } else {
                if verbose { println!("DHCP: lease expired, restarting"); }
                return dhcp(iface, verbose);
            }
        }
    }
    Ok(())
}

fn apply_dhcp_config(
    iface: &str,
    router: Option<[u8; 4]>,
    dns: Option<[u8; 4]>,
    verbose: bool,
) -> Result<(), String> {
    if let Some(router) = router {
        let route = format!("default via {}.{}.{}.{}", router[0], router[1], router[2], router[3]);
        try_fmt!(set_cfg_value("route/add", &route), "failed to set route");
    }
    if let Some(mut dns) = dns {
        if dns[0] == 127 {
            dns = [9, 9, 9, 9];
            if verbose { println!("DHCP: replaced loopback DNS with Quad9"); }
        }
        let ns = format!("{}.{}.{}.{}", dns[0], dns[1], dns[2], dns[3]);
        try_fmt!(set_cfg_value("resolv/nameserver", &ns), "failed to set DNS");
    }
    Ok(())
}

fn compute_prefix_len(subnet: Option<[u8; 4]>) -> u32 {
    let Some(subnet) = subnet else { return 24 };
    let inverted: u32 = !u32::from_be_bytes(subnet);
    inverted.leading_zeros()
}

fn parse_options(options: &[u8], cb: &mut dyn FnMut(u8, &[u8])) {
    let mut i = 0;
    while i < options.len() {
        let code = options[i];
        if code == 0 { i += 1; continue; }
        if code == OPT_END { break; }
        i += 1;
        if i >= options.len() { break; }
        let len = options[i] as usize;
        i += 1;
        if i + len > options.len() { break; }
        cb(code, &options[i..i + len]);
        i += len;
    }
}

fn get_message_type(options: &[u8]) -> Option<u8> {
    let mut msg_type = None;
    parse_options(options, &mut |code, data| {
        if code == OPT_MESSAGE_TYPE && data.len() == 1 {
            msg_type = Some(data[0]);
        }
    });
    msg_type
}

fn init_dhcp_header(pkt: &mut Dhcp, mac: MacAddr, tid: u32) {
    *pkt = Dhcp::default();
    pkt.op = 1;
    pkt.htype = 1;
    pkt.hlen = 6;
    pkt.tid = tid;
    pkt.flags = 0x8000u16.to_be();
    pkt.chaddr[..6].copy_from_slice(&mac.bytes);
    pkt.magic = 0x63825363u32.to_be();
}

fn send_dhcp(pkt: &Dhcp, socket: &UdpSocket) -> Result<(), String> {
    let data = unsafe {
        std::slice::from_raw_parts(pkt as *const Dhcp as *const u8, std::mem::size_of::<Dhcp>())
    };
    socket.send(data).map(|_| ()).map_err(|e| format!("send: {}", e))
}

/// `send_to(addr, ...)` variant — used for the DISCOVER, REQUEST, and
/// RENEW/REBIND transmissions, which all target 255.255.255.255:67.
/// `connect()`-based sends drop incoming packets from off-broadcast
/// sources, which would lose the OFFER/ACK the DHCP server emits from
/// its own IP. Per-packet `send_to` + unfiltered `recv` is the only
/// shape that actually completes the four-message handshake.
fn send_dhcp_to(
    pkt: &Dhcp,
    socket: &UdpSocket,
    addr: SocketAddr,
) -> Result<(), String> {
    let data = unsafe {
        std::slice::from_raw_parts(pkt as *const Dhcp as *const u8, std::mem::size_of::<Dhcp>())
    };
    socket.send_to(data, addr).map(|_| ()).map_err(|e| format!("send_to: {}", e))
}

impl Default for Dhcp {
    fn default() -> Self {
        Dhcp {
            op: 0, htype: 0, hlen: 0, hops: 0, tid: 0, secs: 0, flags: 0,
            ciaddr: [0; 4], yiaddr: [0; 4], siaddr: [0; 4], giaddr: [0; 4],
            chaddr: [0; 16], sname: [0; 64], file: [0; 128], magic: 0,
            options: [0; 308],
        }
    }
}

fn main() {
    let mut verbose = false;
    let mut iface = "eth0".to_string();
    for arg in env::args().skip(1) {
        match arg.as_ref() {
            "-v" => verbose = true,
            other if !other.starts_with('-') => iface = other.to_string(),
            _ => (),
        }
    }

    // driver-manager attaches the NIC asynchronously, and smolnetd only brings up
    // /scheme/netcfg/ifaces/<iface> after it binds the adapter, so on a fresh boot
    // this path may not exist yet when dhcpd runs (10_dhcpd only requires_weak
    // that 10_smolnetd *started*). Wait (bounded) for the interface to appear
    // instead of failing immediately with "Can't open ...". A genuinely NIC-less
    // machine falls through after the window and exits cleanly (0), no error.
    let mac_path = format!("/scheme/netcfg/ifaces/{iface}/mac");
    let deadline = std::time::Instant::now() + std::time::Duration::from_secs(20);
    while std::fs::File::open(&mac_path).is_err() {
        if std::time::Instant::now() >= deadline {
            eprintln!("dhcpd: {iface} never appeared in /scheme/netcfg (no NIC?); skipping DHCP");
            return;
        }
        std::thread::sleep(std::time::Duration::from_millis(250));
    }

    if let Err(err) = dhcp(&iface, verbose) {
        eprintln!("dhcpd: {err}");
        process::exit(1);
    }
}

#[cfg(test)]
mod test {
    use super::MacAddr;

    #[test]
    fn from_str_test() {
        let mac = MacAddr { bytes: [0x01, 0x23, 0x45, 0x67, 0x89, 0xab] };
        let empty_mac = MacAddr::default();
        assert_eq!(mac, MacAddr::from_str("01:23:45:67:89:ab"));
        assert_eq!(mac, MacAddr::from_str("1:23:45:67:89:ab"));
        assert_eq!(mac, MacAddr::from_str("01:23:45:67:89:AB"));
        assert_eq!(mac, MacAddr::from_str("01-23-45-67-89-ab"));
        assert_eq!(empty_mac, MacAddr::from_str(""));
        assert_eq!(empty_mac, MacAddr::from_str("01:23:45:67:89"));
        assert_eq!(empty_mac, MacAddr::from_str("01:23:45:67:89:ab:cd"));
        assert_eq!(empty_mac, MacAddr::from_str("x1:23:45:67:89:ab"));
        assert_eq!(mac, MacAddr::from_str(&mac.to_string()));
    }
}