mod ctx;

use std::{net::Ipv4Addr, process::ExitCode};

use ctx::Context;
use dhcproto::{
    v4::{DhcpOption, Message, MessageType, Opcode, OptionCode},
    Decodable, Decoder, Encodable, Encoder,
};
use nzr_api::{config::Config, net::mac::MacAddr};
use std::net::SocketAddr;
use tracing::instrument;

const EMPTY_V4: Ipv4Addr = Ipv4Addr::new(0, 0, 0, 0);
const DEFAULT_LEASE: u32 = 86400;

fn make_reply(msg: &Message, msg_type: MessageType, lease_addr: Option<Ipv4Addr>) -> Message {
    let mut resp = Message::new(
        EMPTY_V4,
        lease_addr.unwrap_or(EMPTY_V4),
        EMPTY_V4,
        msg.giaddr(),
        msg.chaddr(),
    );
    resp.set_opcode(Opcode::BootReply)
        .set_xid(msg.xid())
        .set_htype(msg.htype())
        .set_flags(msg.flags());
    resp.opts_mut().insert(DhcpOption::MessageType(msg_type));
    resp
}

#[instrument(skip(ctx, msg))]
async fn handle_message(ctx: &Context, from: SocketAddr, msg: &Message) {
    if msg.opcode() != Opcode::BootRequest {
        tracing::warn!("Invalid incoming opcode {:?}", msg.opcode());
        return;
    }

    let Some(DhcpOption::MessageType(msg_type)) = msg.opts().get(OptionCode::MessageType) else {
        tracing::warn!("Missing DHCP message type");
        return;
    };

    let Ok(client_mac) = MacAddr::from_bytes(msg.chaddr()) else {
        tracing::info!("Received DHCP payload with invalid addr (different media type?)");
        return;
    };
    tracing::debug!("Client MAC is {client_mac}!!!");

    let instance = match ctx.instance_by_mac(client_mac).await {
        Ok(Some(i)) => i,
        Ok(None) => {
            tracing::info!("{msg_type:?} from unknown host {client_mac}, ignoring");
            return;
        }
        Err(err) => {
            tracing::error!("Error getting instance for {client_mac}: {err}");
            return;
        }
    };

    tracing::info!(
        "Recieved {msg_type:?} from {client_mac} (assuming {})",
        &instance.name
    );

    let mut lease_time = None;
    let mut nak = false;

    let mut response = match msg_type {
        MessageType::Discover => {
            lease_time = Some(DEFAULT_LEASE);
            make_reply(msg, MessageType::Offer, Some(instance.lease.addr.addr))
        }
        MessageType::Request => {
            if let Some(DhcpOption::RequestedIpAddress(addr)) =
                msg.opts().get(OptionCode::RequestedIpAddress)
            {
                if *addr == instance.lease.addr.addr {
                    make_reply(msg, MessageType::Ack, Some(instance.lease.addr.addr))
                } else {
                    nak = true;
                    make_reply(msg, MessageType::Nak, None)
                }
            } else {
                nak = true;
                make_reply(msg, MessageType::Nak, None)
            }
        }
        MessageType::Decline => {
            tracing::warn!(
                "Client (assumed to be {}) informed us that {} is in use by another server",
                &instance.name,
                instance.lease.addr.addr
            );
            return;
        }
        MessageType::Release => {
            // We only provide static leases
            tracing::debug!("Ignoring DHCPRELEASE");
            return;
        }
        MessageType::Inform => make_reply(msg, MessageType::Ack, None),
        other => {
            tracing::info!("Received unhandled message {other:?}");
            return;
        }
    };

    {
        let opts = response.opts_mut();
        let giaddr = if msg.giaddr().is_unspecified() {
            todo!("no relay??")
        } else {
            msg.giaddr()
        };

        opts.insert(DhcpOption::ServerIdentifier(giaddr));
        if let Some(time) = lease_time {
            opts.insert(DhcpOption::AddressLeaseTime(time));
        }

        if !nak {
            // Get general networking info
            let subnet = match ctx.get_subnet(&instance.lease.subnet).await {
                Ok(Some(net)) => net,
                Ok(None) => {
                    tracing::error!("nzrd says '{}' isn't a subnet", &instance.lease.subnet);
                    return;
                }
                Err(err) => {
                    tracing::error!("Error getting subnet: {err}");
                    return;
                }
            };

            opts.insert(DhcpOption::Hostname(instance.name.clone()));

            if !subnet.dns.is_empty() {
                opts.insert(DhcpOption::DomainNameServer(subnet.dns));
            }

            if let Some(name) = subnet.domain_name {
                opts.insert(DhcpOption::DomainName(name.to_utf8()));
            }

            if let Some(gw) = subnet.gateway4 {
                opts.insert(DhcpOption::Router(Vec::from(&[gw])));
            }

            opts.insert(DhcpOption::SubnetMask(instance.lease.addr.netmask()));
        }
    }

    tracing::info!(
        "Sending message {:?} with yiaddr {}",
        response
            .opts()
            .get(OptionCode::MessageType)
            .unwrap_or(&DhcpOption::End),
        response.yiaddr()
    );

    // unicast it back
    let mut resp_buf = Vec::new();
    let mut enc = Encoder::new(&mut resp_buf);
    if let Err(err) = response.encode(&mut enc) {
        tracing::error!("Couldn't encode response: {err}");
        return;
    }

    if let Err(err) = ctx.sock().send_to(&resp_buf, from).await {
        tracing::error!("Couldn't send response: {err}");
    }
}

#[tokio::main]
async fn main() -> ExitCode {
    tracing_subscriber::fmt().init();
    let cfg: Config = match Config::figment().extract() {
        Ok(cfg) => cfg,
        Err(err) => {
            tracing::error!("Unable to get configuration: {err}");
            return ExitCode::FAILURE;
        }
    };

    let ctx = match Context::new(&cfg).await {
        Ok(ctx) => ctx,
        Err(err) => {
            tracing::error!("{err}");
            return ExitCode::FAILURE;
        }
    };

    tracing::info!("nzrdhcp ready! Listening on {}", ctx.addr());

    loop {
        let mut buf = [0u8; 576];
        let (_, src) = match ctx.sock().recv_from(&mut buf).await {
            Ok(x) => x,
            Err(err) => {
                tracing::error!("recv_from error: {err}");
                return ExitCode::FAILURE;
            }
        };

        let msg = match Message::decode(&mut Decoder::new(&buf)) {
            Ok(msg) => msg,
            Err(err) => {
                tracing::error!("Couldn't process message from {}: {}", src, err);
                continue;
            }
        };

        handle_message(&ctx, src, &msg).await;
    }
}