use std::io;
use std::net::{SocketAddr, ToSocketAddrs};

use std::time::{Duration, SystemTime, SystemTimeError};

use clap::{arg, Parser};
use tokio::net::UdpSocket;
use tracing::{debug, error};


use crate::Destination::{Address, SockAddr};

const TTL: Duration = Duration::from_secs(300);

/// Simple program to greet a person
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// The address and port to listen and accept cleartext UDP datagrams from
    #[arg(short, long)]
    bind: String,

    /// The local endpoint to read cleartext datagrams from & write to
    #[arg(short, long)]
    local: String,

    /// The destination to send obfuscated UDP datagrams to
    #[arg(short, long)]
    sendto: Option<String>,

    /// The hex representation of XOR key to use
    #[arg(short, long)]
    key: String,
}

#[derive(Debug, Eq, PartialEq)]
enum Destination<'a> {
    None,
    Address(&'a String),
    SockAddr(SocketAddr),
}

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt::init();
    let args = Args::parse();
    let key = hex::decode(args.key)
        .expect("failed to decode key hex string");
    let sock = UdpSocket::bind(&args.bind).await
        .expect("failed to bind UDP endpoint");
    let mut buf = [0; 65536];
    let local = args.local.parse().expect("failed to parse local address");
    let mut sendto = Destination::None;
    let args_sendto: String;
    if let Some(s) = args.sendto {
        args_sendto = s;
        sendto = Address(&args_sendto);
    }
    let mut addr_sendto: Option<SocketAddr> = None;
    let mut addr_sendto_t = SystemTime::UNIX_EPOCH;
    loop {
        match sock.recv_from(&mut buf).await {
            Ok((len, addr)) => {
                debug!("recv_from OK, len {}, addr {}", len, addr);
                let mut i = 0;
                let n = key.len();
                for j in 0..len {
                    buf[j] ^= key[i];
                    i += 1;
                    if i == n {
                        i = 0;
                    }
                }
                if sendto == Destination::None && addr != local {
                    // assume the packet is coming from target
                    // update missing sendto parameter
                    sendto = SockAddr(addr);
                }
                let sendto = match sendto {
                    Destination::None => Destination::None,
                    Address(addr) => {
                        let now = SystemTime::now();
                        if now.duration_since(addr_sendto_t)
                            .or::<SystemTimeError>(Ok(Duration::from_secs(0))).unwrap() > TTL {
                            // update
                            debug!("resolving address: {}", &addr);
                            match resolve_addr(&addr) {
                                Ok(addr) => {
                                    addr_sendto = Some(addr);
                                    addr_sendto_t = SystemTime::now();
                                    SockAddr(addr)
                                }
                                Err(why) => {
                                    error!("Failed to resolve address {}: {}", addr, why);
                                    Destination::None
                                }
                            }
                        } else {
                            // cache is still valid
                            match addr_sendto {
                                None => Destination::None,
                                Some(addr) => SockAddr(addr),
                            }
                        }
                    }
                    SockAddr(addr) => SockAddr(addr),
                };
                let target = match sendto {
                    Destination::None => None,
                    Address(_) => panic!("unresolved address"),
                    SockAddr(sendto_addr) => {
                        if addr == sendto_addr {
                            // remote -> local
                            Some(local)
                        } else {
                            Some(sendto_addr)
                        }
                    }
                };
                match target {
                    None => {
                        error!("cannot send_to: destination not determined");
                        continue;
                    }
                    Some(target) => {
                        if let Err(why) = sock.send_to(&buf[..len], target).await {
                            debug!("send_to failed, dest: {}, error: {}", target, why);
                        } else {
                            debug!("send_to OK, len: {}, dest: {}", len, target);
                        }
                    }
                }
            }
            Err(why) => {
                debug!("recv_from failed: {}", why);
            }
        }
    }
}

fn resolve_addr(addr: &String) -> io::Result<SocketAddr> {
    match addr.to_socket_addrs() {
        Ok(mut addrs) => match addrs.next() {
            None => Err(io::Error::new(io::ErrorKind::Other, "no available address")),
            Some(addr) => Ok(addr),
        },
        Err(why) => Err(why),
    }
}