use clap::ValueEnum;
use log::{debug, info};
use std::net::{IpAddr, ToSocketAddrs, UdpSocket};
use tailstun::{AddrPort, StunMessage, TxId};

#[derive(Debug, Clone, ValueEnum)]
enum OutputFormat {
    Text,
    Json,
    Yaml,
}

impl OutputFormat {
    fn format_stun(&self, msg: &StunMessage) -> String {
        match self {
            OutputFormat::Text => format!("{}", msg),
            OutputFormat::Json => serde_json::to_string_pretty(msg).unwrap(),
            OutputFormat::Yaml => serde_yaml::to_string(msg).unwrap(),
        }
    }
    fn format_address(&self, a: &AddrPort) -> String {
        let a = match a.address {
            IpAddr::V4(_) => a.address,
            IpAddr::V6(v6) => {
                if let Some(v4) = v6.to_ipv4_mapped() {
                    IpAddr::V4(v4)
                } else {
                    a.address
                }
            }
        };
        match self {
            OutputFormat::Text => format!("{}", a),
            OutputFormat::Json => serde_json::to_string_pretty(&a).unwrap(),
            OutputFormat::Yaml => serde_yaml::to_string(&a).unwrap(),
        }
    }
}

#[derive(clap::Parser)]
#[command(version, about, long_about = None)]
#[command(about = "Test a Tailscale derp node's stun service")]
struct Cli {
    host: String,
    #[clap(default_value = "3478")]
    port: u16,
    #[clap(short = '4', conflicts_with = "v6_only", default_value = "false")]
    v4_only: bool,
    #[clap(short = '6', conflicts_with = "v4_only", default_value = "false")]
    v6_only: bool,
    #[clap(short, long, default_value = "text")]
    format: OutputFormat,
    #[clap(
        short,
        long,
        default_value = "false",
        help = "Only output the first mapped address & convert IPv6-mapped to IPv4"
    )]
    address_only: bool,
    #[command(flatten)]
    verbose: clap_verbosity_flag::Verbosity,
}

fn main() {
    let cli = <Cli as clap::Parser>::parse();
    env_logger::Builder::new()
        .filter_level(cli.verbose.log_level_filter())
        .init();

    let dest = (cli.host.as_str(), cli.port)
        .to_socket_addrs()
        .expect("Unable to resolve host")
        .find(|a| {
            if cli.v4_only {
                a.is_ipv4()
            } else if cli.v6_only {
                a.is_ipv6()
            } else {
                true
            }
        })
        .expect("No address found for host");

    let socket = UdpSocket::bind("[::]:0").expect("Unable to bind a UDP socket");
    socket
        .connect(dest)
        .expect("Unable to connect to the destination");

    info!(
        "Connected to {:?} from {:?}",
        socket.peer_addr(),
        socket.local_addr()
    );

    let req = TxId::new().make_request();

    debug!("Sending request {:?}", &req);

    socket.send(&req).expect("Unable to send request");

    let mut buf = [0u8; 1500];

    match socket.recv(&mut buf) {
        Ok(received) => {
            let buf = &buf[..received];
            debug!("Received response: {:?}", buf);

            let msg = StunMessage::parse(buf).unwrap();
            info!("Parsed message from {}:", socket.peer_addr().unwrap());
            if cli.address_only {
                match msg.attributes.mapped_address() {
                    Some(addr) => println!("{}", cli.format.format_address(addr)),
                    None => {
                        // No mapped address
                        std::process::exit(1);
                    }
                }
            } else {
                println!("{}", cli.format.format_stun(&msg));
            }
        }
        Err(e) => {
            println!("recv function failed: {e:?}");
        }
    }
}