use anyhow::{
    Context,
    Result,
};
use clap::Parser;
use dhcproto::{
    Decodable,
    Decoder,
    Encodable,
    Encoder,
    v6::{
        DhcpOption,
        Message,
        MessageType,
        OptionCode,
        ORO,
    },
};
use ipnetwork;
use log::{
    debug,
    log_enabled,
    info,
    Level,
};
use pnet;
use serde::Serialize;
use std::net::SocketAddrV6;
use tokio::net::UdpSocket;

const ALL_DHCP_RELAY_AGENTS_AND_SERVERS: &str = "ff02::1:2";
const DHCP_CLIENT_PORT: u16 = 546;
const DHCP_RELAY_AGENT_AND_SERVER_PORT: u16 = 547;

/// Requests some simple parameters using DHCPv6 from specified network
#[derive(Parser)]
#[command(version, about, long_about = None)]
struct Cli {
    /// Network interface name used for DHCPv6
    interface: String,
    /// Format output as json
    #[arg(long)]
    json: bool,
}

#[derive(Serialize)]
struct MessageMap {
    #[serde(skip_serializing_if = "Option::is_none")]
    server_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    domain_name_servers: Option<Vec<String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    domain_search_list: Option<Vec<String>>,
}

impl MessageMap {
    fn new() -> Self {
        Self {
            server_id: None,
            domain_name_servers: None,
            domain_search_list: None,
        }
    }
}

fn message_to_map(msg: &Message) -> MessageMap {
    let mut map = MessageMap::new();

    if let Some(DhcpOption::ServerId(server_id)) = &msg.opts().get(OptionCode::ServerId) {
        map.server_id = Some(hex::encode(&server_id));
    }

    if let Some(DhcpOption::DomainNameServers(domain_name_servers)) = &msg.opts().get(OptionCode::DomainNameServers) {
        map.domain_name_servers = Some(domain_name_servers.iter().map(|a| a.to_string()).collect());
    }
    if let Some(DhcpOption::DomainSearchList(domain_search_list)) = &msg.opts().get(OptionCode::DomainSearchList) {
        map.domain_search_list = Some(domain_search_list.iter().map(|n| n.to_string()).collect());
    }

    return map;
}

#[tokio::main]
async fn main() -> Result<()>{
    env_logger::init();

    let cli = Cli::parse();

    let all_interfaces = pnet::datalink::interfaces();
    let selected_interface = all_interfaces
        .iter()
        .find(|i| i.name == cli.interface)
        .with_context(|| format!("No interface found with name: {}", cli.interface))?;

    let ipv6_addresses = selected_interface.ips
        .iter()
        .filter_map(|a| match a {
            ipnetwork::IpNetwork::V6(a) => Some(a),
            _ => None,
        });
    let mut ipv6_link_local_addresses = ipv6_addresses
        .filter(|a| a.is_subnet_of("fe80::/10".parse().unwrap()));

    // Just take the first address found on the interface
    let selected_address = ipv6_link_local_addresses.next()
        .context("No IPv6 link local address assigned to this interface")?;

    let socket_addr = SocketAddrV6::new(selected_address.ip(), DHCP_CLIENT_PORT, 0, selected_interface.index.clone());
    let sock = UdpSocket::bind(socket_addr).await
        .context("Unable to bind UDP socket")?;

    info!("Listening on {}", sock.local_addr()
        .context("Failed to fetch address socket bound to")?);

    let remote_addr = SocketAddrV6::new(ALL_DHCP_RELAY_AGENTS_AND_SERVERS.parse().unwrap(), DHCP_RELAY_AGENT_AND_SERVER_PORT, 0, selected_interface.index.clone());

    let mut msg = Message::new(MessageType::InformationRequest);
    let msg_opts = msg.opts_mut();
    //msg_opts.insert(DhcpOption::ClientId(Vec::from(duid)));
    msg_opts.insert(DhcpOption::ElapsedTime(0));
    msg_opts.insert(DhcpOption::ORO(ORO { opts: vec![
        OptionCode::InfMaxRt,
        OptionCode::InformationRefreshTime,
        OptionCode::DomainNameServers,
        OptionCode::DomainSearchList,
    ],})); // request options

    let mut msg_buf = Vec::new();
    let mut msg_encoder = Encoder::new(&mut msg_buf);
    msg.encode(&mut msg_encoder)
        .context("Unable to serialize DHCP options for message to send")?; // Serializes msg to msg_buf

    sock.send_to(&msg_buf, remote_addr).await
        .context("Failed to send DHCPv6 INFORMATION-REQUEST message")?;

    let response_msg = loop {
        let mut recv_buf = [0; 1024];
        let (recv_len, recv_addr) = sock.recv_from(&mut recv_buf).await
            .context("Failed to receive DHCPv6 response")?;

        if log_enabled!(Level::Debug) {
            debug!("{:?} bytes received from {:?}", recv_len, recv_addr);
            debug!("Received packet: {}", hex::encode(&recv_buf[..recv_len]));
        }

        let recv_msg = Message::decode(&mut Decoder::new(&recv_buf))
            .context("Unable to parse DHCPv6 response message")?;

        if log_enabled!(Level::Debug) {
            debug!("Packet contains transaction id: 0x{}", hex::encode(&msg.xid()));
        }

        if recv_msg.xid() == msg.xid() {
            info!("Found reponse with matching transaction id");
            break recv_msg;
        }
    };

    let message_map = message_to_map(&response_msg);

    if cli.json {
        let message_map_json = serde_json::to_string(&message_map)
            .context("Unable to format message to json")?;
        println!("{}", message_map_json);
    }
    else {
        let message_map_yaml = serde_yaml::to_string(&message_map)
            .context("Unable to format message to yaml")?;
        println!("{}", message_map_yaml);
    }

    Ok(())
}