use tokio::sync::mpsc;
use tokio::{net::UdpSocket, sync::Mutex, time};
use x25519_dalek::{PublicKey, StaticSecret};
use base64::prelude::*;
use log::{error, info};
use core::str;
use std::sync::Arc;
use std::net::{ SocketAddr, Ipv4Addr, IpAddr };
use std::collections::HashMap;
use aes_gcm::{ aead::{Aead, AeadCore, KeyInit, OsRng},
Aes256Gcm, Nonce };

#[cfg(target_os = "linux")]
use network_interface::{NetworkInterface, NetworkInterfaceConfig};

use frida_core::config::{ ServerConfiguration, ServerPeer};
use frida_core::udp::{UDPKeepAlive, UDPSerializable, UDPVpnHandshake, UDPVpnPacket};
use frida_core::{DeviceReader, DeviceWriter, create, device::AbstractDevice};

#[cfg(target_os = "linux")]
fn configure_routes(s_interface: Option<&str>) {
    let interfaces = NetworkInterface::show().unwrap();

    let net_inter = interfaces.iter()
        .filter(|i| !i.addr.iter().any(|b| b.ip().to_string() == "127.0.0.1" || b.ip().to_string() == "::1") )
        .min_by(|x, y| x.index.cmp(&y.index))
        .unwrap();

    info!("Main network interface: {:?}", net_inter.name);

    let inter_name = if s_interface.is_some() { s_interface.unwrap() } else { &net_inter.name };

    let mut ip_output = std::process::Command::new("iptables")
        .arg("-A")
        .arg("FORWARD")
        .arg("-i")
        .arg("tun0")
        .arg("-o")
        .arg(inter_name)
        .arg("-j")
        .arg("ACCEPT")
        .output()
        .expect("Failed to execute iptables command.");

    if !ip_output.status.success() {
        error!("Failed to forward packets: {:?}", String::from_utf8_lossy(&ip_output.stderr));
    }

    ip_output = std::process::Command::new("iptables")
        .arg("-A")
        .arg("FORWARD")
        .arg("-i")
        .arg(inter_name)
        .arg("-o")
        .arg("tun0")
        .arg("-m")
        .arg("state")
        .arg("--state")
        .arg("ESTABLISHED,RELATED")
        .arg("-j")
        .arg("ACCEPT")
        .output()
        .expect("Failed to execute iptables command.");

    if !ip_output.status.success() {
        error!("Failed to forward packets: {:?}", String::from_utf8_lossy(&ip_output.stderr));
    }
    
    ip_output = std::process::Command::new("iptables")
        .arg("-t")
        .arg("nat")
        .arg("-A")
        .arg("POSTROUTING")
        .arg("-o")
        .arg(inter_name)
        .arg("-j")
        .arg("MASQUERADE")
        .output()
        .expect("Failed to execute iptables command.");

    if !ip_output.status.success() {
        error!("Failed to forward packets: {:?}", String::from_utf8_lossy(&ip_output.stderr));
    }
}

#[cfg(not(target_os = "android"))]
pub async fn server_mode(server_config: ServerConfiguration, s_interface: Option<&str>) {
    info!("Starting server...");

    let mut config = AbstractDevice::default();
    config.address(server_config.interface.internal_address.parse().unwrap())
        .netmask(Ipv4Addr::new(255, 255, 255, 0)) // 255.255.255.0
        .tun_name("tun0");
    
    let (mut dev_reader, mut dev_writer) = create(config);

    let sock = UdpSocket::bind(&server_config.interface.bind_address).await.unwrap();
    let sock_rec = Arc::new(sock);
    let sock_hnd = sock_rec.clone();
    let addresses = Arc::new(Mutex::new(HashMap::<IpAddr, UDPeer>::new()));
    let peers = Arc::new(Mutex::new(Vec::<ServerPeer>::new()));

    let (send2tun, mut recv2tun) = mpsc::unbounded_channel::<Vec<u8>>(); // unbounded::<Vec<u8>>();

    let (send2hnd, mut recv2hnd) = mpsc::unbounded_channel::<(Vec<u8>, SocketAddr)>(); // unbounded::<(Vec<u8>, SocketAddr)>();

    let (send2loop, mut recv_from_tun) = mpsc::unbounded_channel::<Vec<u8>>();

    #[cfg(target_os = "linux")]
    configure_routes(s_interface);

    let tun_rdr_task = tokio::spawn(async move {
        let mut buf = vec![0; 4096];
        loop {
            if let Ok(n) = dev_reader.read(&mut buf).await {
                let _ = send2loop.send(buf[..n].to_vec());
            }
        }
    });

    let tun_writer_task = tokio::spawn(async move {
        loop {
            if let Some(bytes) = recv2tun.recv().await {
                //info!("Sent to tun!");
                let _ = dev_writer.write(&bytes).await;
            }
        }
    });

    let keepalive_sec = server_config.interface.keepalive.clone();
    let send2hnd_cl = send2hnd.clone();
    let addrs_lcl = addresses.clone();

    let alive_task = tokio::spawn(async move {
        let kp_sc = keepalive_sec.clone();
        if kp_sc <= 0 { return; }
        loop {
            time::sleep(time::Duration::from_secs(kp_sc.into())).await;
            let mmp = addrs_lcl.lock().await;
            mmp.values().for_each(|p| {
                let _ = send2hnd_cl.send((UDPKeepAlive{}.serialize(), p.addr));
            });
            drop(mmp);
        }
    });

    let sock_writer_task = tokio::spawn(async move {
        loop {
            if let Some((handshake, addr)) = recv2hnd.recv().await {
                //info!("I SENT THAT STUFF");
                let _ = sock_hnd.send_to(&handshake, addr).await;
            }
        }
    });

    let addrs_cl = addresses.clone();
    let send2hnd_sr = send2hnd.clone();
    let tun_reader_task = tokio::spawn(async move {
        loop {
            if let Some(buf) = recv_from_tun.recv().await {
                if buf.len() <= 19 { continue; }
            
                let ip = IpAddr::V4(Ipv4Addr::new(buf[16], buf[17], buf[18], buf[19]));
                let mp = addrs_cl.lock().await;
                if let Some(peer) = mp.get(&ip) {
                    let aes = Aes256Gcm::new(&peer.shared_secret.into());
                    let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
    
                    let ciphered_data = aes.encrypt(&nonce, &buf[..]);
    
                    if let Ok(ciphered_d) = ciphered_data {
                        let vpn_packet = UDPVpnPacket{ data: ciphered_d, nonce: nonce.to_vec()};
                        let _ = send2hnd_sr.send((vpn_packet.serialize(), peer.addr));
                    } else {
                        error!("Traffic encryption failed.");
                    }
                } else {
                    // TODO: check in config is broadcast mode enabled (if not, do not send this to everyone)
                    //mp.values().for_each(| peer | { sock_snd.send_to(&buf[..n], peer.addr); });
                }
                drop(mp);
            }
        }
    });
    
    let addrs_lp = addresses.clone();
    let peers_lp = peers.clone();

    let mut f_plp = peers_lp.lock().await;
    server_config.peers.iter().for_each(|c| f_plp.push(c.clone()));
    drop(f_plp);

    let send2hnd_ssr = send2hnd.clone();

    let sock_reader_task = tokio::spawn(async move {
        let mut buf = vec![0; 2048];
        loop {
            if let Ok((len, addr)) = sock_rec.recv_from(&mut buf).await {
                //info!("There is packet!");
                let mut mp = addrs_lp.lock().await;
                let plp = peers_lp.lock().await;
                match buf.first() {
                    Some(h) => {
                        match h {
                            0 => {
                                let handshake = UDPVpnHandshake::deserialize(&buf);
                                //info!("Got handshake from {:?}", handshake.request_ip);
                                let skey = BASE64_STANDARD.encode(&handshake.public_key);
                                if plp.iter().any(|c| c.ip == handshake.request_ip && c.public_key == skey) {
                                    // todo: reduce repeated code blocks count
                                    let internal_ip = IpAddr::V4(handshake.request_ip);
                                    info!("Accepted client from {:?}", handshake.request_ip);
                                    let mut k = [0u8; 32];
                                    for (&x, p) in handshake.public_key.iter().zip(k.iter_mut()) {
                                        *p = x;
                                    }
                                    let static_secret = BASE64_STANDARD.decode(&server_config.interface.private_key).unwrap();
                                    let mut k1 = [0u8; 32];
                                    for (&x, p) in static_secret.iter().zip(k1.iter_mut()) {
                                        *p = x;
                                    }
                                    let shared_secret = StaticSecret::from(k1)
                                        .diffie_hellman(&PublicKey::from(k));
                                    mp.insert(internal_ip, UDPeer { addr, shared_secret: *shared_secret.as_bytes() });
                                    
                                    let handshake_response = UDPVpnHandshake{ public_key: BASE64_STANDARD.decode(&server_config.interface.public_key).unwrap(), request_ip: handshake.request_ip };
    
                                    let _ = send2hnd_ssr.send((handshake_response.serialize(), addr));
                                } else {
                                    info!("Bad handshake");
                                    //plp.iter().for_each(|c| info!("ip: {:?}; pkey: {:?}", c.ip, c.public_key));
                                }
                            }, // handshake
                            1 => {
                                let packet = UDPVpnPacket::deserialize(&(buf[..len].to_vec()));
                                mp.values().filter(| p | p.addr == addr).for_each(|p| {
                                    let aes = Aes256Gcm::new(&p.shared_secret.into());
                                    let nonce = Nonce::clone_from_slice(&packet.nonce[..]);
                                    match aes.decrypt(&nonce, &packet.data[..]) {
                                        Ok(decrypted) => { 
                                            //info!("Start of packet: {:#?} from {}", decrypted[..12].to_vec(), addr);
                                            let _ = send2tun.send(decrypted); 
                                        },
                                        Err(error) => error!("Decryption error! {:?}", error)
                                    }
                                });
                            }, // payload
                            2 => { }, // got keepalive packet
                            _ => error!("Unexpected header value.")
                        }
                    },
                    None => error!("There is no header")
                }
                drop(plp);
                drop(mp);
            }
        }
    });
    
    // should be refactored
    let _ = tokio::join!(tun_reader_task, sock_reader_task, sock_writer_task, tun_writer_task, alive_task, tun_rdr_task);
}

struct UDPeer {
    addr: SocketAddr,
    shared_secret: [u8; 32]
}