//! A [`TunnelScaffolding`] implementation for the Android platform.

use crate::ffi_event::{RelayDetails, TunnelEvent};
use crate::OnionmasqMobile;
use futures::Stream;
use jni::objects::JValue;
use jni::sys::{jboolean, jint};
use onion_tunnel::accounting::BandwidthCounter;
use onion_tunnel::scaffolding::{ConnectionDetails, FailedConnectionDetails, TunnelCommand};
use onion_tunnel::{CountryCode, IpEndpoint, TunnelScaffolding};
use simple_proc_net::ProcNetEntry;
use std::collections::HashMap;
use std::io;
use std::net::{IpAddr, SocketAddr};
use std::os::fd::RawFd;
use std::sync::{Arc, Mutex, RwLock};
use tokio::sync::mpsc;
use tokio_stream::wrappers::UnboundedReceiverStream;
use tracing::{debug, error, trace, warn};

pub struct AndroidScaffolding {
    /// This is only mpsc so it can be Sync.
    pub(crate) bootstrap_done: mpsc::UnboundedSender<()>,
    pub(crate) command_receiver: Mutex<Option<mpsc::UnboundedReceiver<TunnelCommand>>>,
    pub(crate) counters: super::BandwidthCounterMap,
    pub(crate) country_code: Arc<RwLock<Option<CountryCode>>>,
}

impl AndroidScaffolding {
    fn on_tunnel_event(evt: TunnelEvent) {
        if let Err(e) = OnionmasqMobile::get().on_tunnel_event(evt) {
            warn!("Failed to communicate tunnel event: {e}");
        }
    }
}

impl TunnelScaffolding for AndroidScaffolding {
    fn protect(&self, socket: RawFd, _: &SocketAddr) -> io::Result<()> {
        debug!("AndroidScaffolding::protect() for fd {}", socket);
        let mobile = OnionmasqMobile::get();

        match mobile.call_jni_class_method::<jboolean>("protect", "(I)Z", &[JValue::Int(socket)]) {
            Ok(1) => {
                // 1 == JNI_TRUE; i.e. successful call
            }
            Ok(_) => {
                // 0 (or other); i.e. failed call
                warn!("AndroidScaffolding::protect(): got error from Java side");
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    "JNI protect() returned error".to_string(),
                ));
            }
            Err(e) => {
                error!("AndroidScaffolding::protect(): FFI failed: {}", e);
                return Err(io::Error::new(io::ErrorKind::Other, e.to_string()));
            }
        }
        Ok(())
    }

    fn isolate(&self, src: IpEndpoint, dst: IpEndpoint, ip_proto: u8) -> io::Result<u64> {
        // NOTE: using trace level logging here to avoid exposing `src` and `dst` to logcat

        if OnionmasqMobile::get().get_android_api().unwrap() >= 29 {
            trace!("isolate(): calling getConnectionOwnerUid({ip_proto}, {src:?}, {dst:?})");

            let ret = get_connection_owner_uid(
                ip_proto as i32,
                src.addr.into(),
                src.port.into(),
                dst.addr.into(),
                dst.port.into(),
            );

            trace!("isolate(): got connection result of {:?}", ret);
            ret
        } else {
            // Android API levels < 29 can't use ConnectivityManager#getConnectionOwnerUid, so we
            // read `/proc` instead.

            trace!("isolate(): reading /proc for {src} -> {dst} proto {ip_proto}");

            let iter = match ip_proto {
                6 => ProcNetEntry::tcp4()?.chain(ProcNetEntry::tcp6()?),
                17 => ProcNetEntry::udp4()?.chain(ProcNetEntry::udp6()?),
                _ => panic!("unknown ip_proto in isolate(): {ip_proto}"),
            };

            for entry in iter {
                match entry? {
                    Ok(v) => {
                        if IpEndpoint::from(v.local_addr) == src
                            && IpEndpoint::from(v.remote_addr) == dst
                        {
                            if v.uid == 0 {
                                // For some reason, we seem to get this in /proc/net/udp etc for connections that
                                // aren't actually from root. Sigh.
                                trace!(
                                    "isolate(): ignoring probably erroneous connection owner UID 0"
                                );
                                continue;
                            }
                            trace!(
                                "isolate(): isolated {src} -> {dst} proto {ip_proto} as {}",
                                v.uid
                            );
                            return Ok(v.uid as u64);
                        }
                    }
                    Err(e) => {
                        // This is potentially a privacy leak, but also on this API level any app
                        // can just directly read /proc/net itself, so I don't think it's worth
                        // worrying about.
                        warn!(
                            "encountered error '{}' parsing /proc/net line: {}",
                            e.error, e.offending_line
                        );
                    }
                }
            }

            Err(io::Error::new(
                io::ErrorKind::Other,
                "got probably erroneous connection owner UID 0",
            ))
        }
    }

    fn locate(&self, _: IpEndpoint, _: IpEndpoint, _: u64) -> Option<CountryCode> {
        *self.country_code.read().unwrap()
    }

    fn on_bootstrapped(&self) {
        let _ = self.bootstrap_done.send(());
    }

    fn on_directory(&self, relays_by_country: HashMap<String, usize>) {
        Self::on_tunnel_event(TunnelEvent::NewDirectory { relays_by_country })
    }

    fn on_established(&self, details: ConnectionDetails<'_>) {
        Self::on_tunnel_event(TunnelEvent::NewConnection {
            proxy_src: SocketAddr::new(details.proxy_src.addr.into(), details.proxy_src.port),
            proxy_dst: SocketAddr::new(details.proxy_dst.addr.into(), details.proxy_dst.port),
            tor_dst: details.tor_dst.to_string(),
            app_id: details.isolation_key,
            circuit: details
                .circuit_relays()
                .into_iter()
                .map(|x| RelayDetails {
                    rsa_identity: x.rsa_identity.map(|x| x.to_string()),
                    ed_identity: x.ed_identity.map(|x| x.to_string()),
                    addresses: x.addresses,
                    country_code: x.country_code.map(|x| x.to_string()),
                })
                .collect(),
        });
    }

    fn on_arti_failure(&self, details: FailedConnectionDetails) {
        Self::on_tunnel_event(TunnelEvent::FailedConnection {
            proxy_src: SocketAddr::new(details.proxy_src.addr.into(), details.proxy_src.port),
            proxy_dst: SocketAddr::new(details.proxy_dst.addr.into(), details.proxy_dst.port),
            tor_dst: details.tor_dst.to_string(),
            app_id: details.isolation_key,
            error: details.error.to_string(),
        })
    }

    fn on_socket_close(&self, src: IpEndpoint, dst: IpEndpoint, err: Option<io::Error>) {
        Self::on_tunnel_event(TunnelEvent::ClosedConnection {
            proxy_src: SocketAddr::new(src.addr.into(), src.port),
            proxy_dst: SocketAddr::new(dst.addr.into(), dst.port),
            error: err.map(|x| x.to_string()),
        })
    }

    fn command_stream(&self) -> Option<Box<dyn Stream<Item = TunnelCommand> + Send + Sync>> {
        if let Some(rx) = self.command_receiver.lock().unwrap().take() {
            Some(Box::new(UnboundedReceiverStream::new(rx)))
        } else {
            panic!("AndroidScaffolding::command_stream() called more than once");
        }
    }

    fn get_bandwidth_counter(&self, isolation_key: u64) -> Option<Arc<BandwidthCounter>> {
        Some(
            self.counters
                .entry(isolation_key)
                .or_insert_with(|| Arc::new(BandwidthCounter::new()))
                .clone(),
        )
    }
}

pub fn get_connection_owner_uid(
    protocol: i32,
    source_address: IpAddr,
    source_port: i32,
    destination_address: IpAddr,
    destination_port: i32,
) -> io::Result<u64> {
    let mobile = OnionmasqMobile::get();
    let guard = mobile.get_jni_env().expect("failed to get JNI env");

    let source_ip_bytes = match source_address {
        IpAddr::V4(source_address) => source_address.octets().to_vec(),
        IpAddr::V6(source_address) => source_address.octets().to_vec(),
    };
    let destination_ip_bytes = match destination_address {
        IpAddr::V4(destination_address) => destination_address.octets().to_vec(),
        IpAddr::V6(destination_address) => destination_address.octets().to_vec(),
    };

    let source_address_jbyte_array = guard
        .byte_array_from_slice(source_ip_bytes.as_slice())
        .unwrap();
    let destination_address_jbyte_array = guard
        .byte_array_from_slice(destination_ip_bytes.as_slice())
        .unwrap();

    let ret = mobile
        .call_jni_class_method::<jint>(
            "getConnectionOwnerUid",
            "(I[BI[BI)I",
            &[
                JValue::from(protocol),
                JValue::from(&source_address_jbyte_array),
                JValue::from(source_port),
                JValue::from(&destination_address_jbyte_array),
                JValue::from(destination_port),
            ],
        )
        .map_err(|e| io::Error::new(io::ErrorKind::Other, format!("JNI call failed: {e}")))?;

    match u64::try_from(ret) {
        Ok(v) => Ok(v),
        Err(_) => Err(io::Error::new(
            io::ErrorKind::Other,
            "got negative response from JNI",
        )),
    }
}
