//! The `onionmasq-mobile` crate provides the Rust side of the Java `OnionMasqJni` class, allowing
//! Android applications to access the functionality provided by the Rust `onion-tunnel` crate.

mod errors;
mod ffi;
mod ffi_event;
mod panic_handling;
mod scaffolding;

use crate::errors::{CountryCodeError, ProxyStoppedError};
use crate::scaffolding::AndroidScaffolding;
use anyhow::{Context, Result};
use dashmap::DashMap;
use ffi_event::TunnelEvent;
use futures::StreamExt;
use jni::objects::{GlobalRef, JString};
use jni::objects::{JValue, JValueOwned};
use jni::{AttachGuard, JNIEnv, JavaVM};
use once_cell::sync::OnceCell;
use onion_tunnel::accounting::BandwidthCounter;
use onion_tunnel::config::TunnelConfig;
use onion_tunnel::scaffolding::TunnelCommand;
use onion_tunnel::{CountryCode, OnionTunnel};
use std::mem;
use std::str::FromStr;
use std::sync::{Arc, Mutex, RwLock};
use tokio::runtime::Runtime;
use tokio::sync::mpsc;
use tracing::{debug, error, info, warn};
use tracing_subscriber::fmt::Subscriber;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;

static ONIONMASQ_MOBILE_SINGLETON: OnceCell<OnionmasqMobile> = OnceCell::new();

// We want to avoid logging detailed connection-level stuff in production builds to avoid
// embarrassing privacy issues (note that the Android logcat is potentially readable by any
// app with the required permissions!)
#[cfg(feature = "verbose")]
static LOG_FILTER: &str = "info,onionmasq_mobile=trace,onion_tunnel=trace,arti_client=debug,tor_chanmgr=debug,tor_proto=debug,smoltcp=debug,onionmasq_pt_wrapper=trace";
#[cfg(not(feature = "verbose"))]
static LOG_FILTER: &str = "info,onionmasq_mobile=debug";

pub(crate) type BandwidthCounterMap = Arc<DashMap<u64, Arc<BandwidthCounter>>>;

/// A wrapper to run and control `onion-tunnel` from Android.
///
/// This `struct` is the main entrypoint for FFI interactions to and from the Android JNI interface.
/// It is initialised once from Java, and other JNI native methods are translated to calls on a
/// singleton instance of this `struct`.
pub struct OnionmasqMobile {
    /// A handle to the Java VM.
    jvm: JavaVM,
    /// A tokio Runtime.
    rt: Runtime,
    /// A reference to the `org.torproject.onionmasq.OnionMasqJni` Java class.
    class_ref: GlobalRef,
    /// If present, a path we should configure the proxy to log pcap files to.
    pcap_path: Mutex<Option<String>>,
    /// If present, a sender we can use to control a currently running proxy.
    ///
    /// (We can also drop this in order to stop the proxy.)
    command_sender: Mutex<Option<mpsc::UnboundedSender<TunnelCommand>>>,
    /// Map of per-app bandwidth counters (shared with the scaffolding).
    bandwidth_counters: BandwidthCounterMap,
    /// Current global country code for traffic (shared with the scaffolding).
    country_code: Arc<RwLock<Option<CountryCode>>>,
}

impl OnionmasqMobile {
    /// Initialise the FFI layer, creating a singleton `OnionmasqMobile` object and returning it.
    /// This should be called only once at the start of the program before doing anything else.
    ///
    /// # Panics
    ///
    /// Panics if called twice, unless the previous invocations failed with `Err`.
    pub fn new(env: &mut JNIEnv) -> Result<&'static Self> {
        panic_handling::register_panic_hook();

        if ONIONMASQ_MOBILE_SINGLETON.get().is_some() {
            panic!("OnionmasqMobile::new() called twice!");
        }

        // Set up logging via `tracing_android`:
        Subscriber::builder()
            .with_env_filter(LOG_FILTER)
            .finish()
            .with(tracing_android::layer("onionmasq").unwrap())
            .init();

        info!("Hello from Rust!");

        if cfg!(feature = "verbose") {
            warn!("Detailed debug logging enabled -- do not use in production builds!");
        }

        // Make a tokio runtime.
        let rt = Runtime::new().context("couldn't create Tokio runtime")?;

        // `JNIEnv` cannot be sent across thread boundaries. To be able to use JNI
        // functions in other threads, we must first obtain the `JavaVM` interface
        // which, unlike `JNIEnv` is `Send`.
        let jvm = env.get_java_vm().context("couldn't get JavaVM")?;

        // We need to find the `OnionMasqJni` class, since we'll use some of its methods like
        // `protect()` in other code.
        debug!("Getting reference to \"org.torproject.onionmasq.OnionMasqJni\" class...");
        let class = env
            .find_class("org/torproject/onionmasq/OnionMasqJni")
            .context("failed to find OnionMasqJni class")?;

        let reference = env
            .new_global_ref(class)
            .context("failed to make global ref")?;

        let ret = Self {
            jvm,
            rt,
            class_ref: reference,
            pcap_path: Mutex::new(None),
            command_sender: Mutex::new(None),
            bandwidth_counters: Arc::new(DashMap::new()),
            country_code: Arc::new(RwLock::new(Default::default())),
        };
        if ONIONMASQ_MOBILE_SINGLETON.set(ret).is_err() {
            panic!("OnionmasqMobile::new() called concurrently somehow!");
        }
        Ok(Self::get())
    }

    /// Get the singleton `OnionmasqMobile` instance, assuming one is present.
    ///
    /// # Panics
    ///
    /// Panics if `init()` has not successfully been called yet.
    pub fn get() -> &'static Self {
        ONIONMASQ_MOBILE_SINGLETON
            .get()
            .expect("OnionmasqMobile::get() called without being initialised (did you forget to call init?)")
    }

    /// Get the Android API version.
    pub fn get_android_api(&self) -> Result<i32> {
        self.call_jni_class_method("getAndroidAPI", "()I", &[])
    }

    /// Run the `onion-tunnel` proxy.
    pub fn run_proxy(
        &self,
        env: &mut JNIEnv,
        cache_dir: JString,
        data_dir: JString,
        fd: i32,
        bridge_lines: JString,
    ) -> Result<()> {
        let cache_dir = env
            .get_string(&cache_dir)
            .context("failed to get string cache_dir")?;
        let cache_dir = cache_dir.to_string_lossy();

        let data_dir = env
            .get_string(&data_dir)
            .context("failed to get string for data_dir")?;
        let data_dir = data_dir.to_string_lossy();

        let bridge_lines = if !bridge_lines.is_null() {
            let bridge_lines = env
                .get_string(&bridge_lines)
                .context("failed to get string bridge_lines")?;
            Some(bridge_lines.to_string_lossy().into_owned())
        } else {
            None
        };

        let api_version = self.get_android_api()?;
        debug!("Onionmasq_runProxy on Android API {}", api_version);

        let (command_sender, command_receiver) = mpsc::unbounded_channel();
        let (bootstrap_done_sender, mut bootstrap_done_receiver) = mpsc::unbounded_channel();

        *self.command_sender.lock().unwrap() = Some(command_sender);

        let pcap_path = self.pcap_path.lock().unwrap().clone();

        let mut config = TunnelConfig::default();

        config.state_dir = Some(format!("{data_dir}/arti-data").into());
        config.cache_dir = Some(format!("{cache_dir}/arti-cache").into());
        config.pt_dir = Some(format!("{cache_dir}/arti-pts").into());

        if let Some(bridge_lines) = bridge_lines {
            config.bridge_lines = bridge_lines
                .lines()
                .filter(|x| !x.is_empty())
                .map(|x| x.to_owned())
                .collect();
        }

        if let Some(pp) = pcap_path {
            config.pcap_path = Some(pp.into());
        }

        let counters = self.bandwidth_counters.clone();
        let country_code = self.country_code.clone();

        let ret = self.rt.block_on(async move {
            debug!("creating onion_tunnel...");

            let mut onion_tunnel = OnionTunnel::create_with_fd(
                AndroidScaffolding {
                    bootstrap_done: bootstrap_done_sender,
                    command_receiver: Mutex::new(Some(command_receiver)),
                    counters,
                    country_code,
                },
                fd,
                config,
            )
            .await
            .context("couldn't start onionmasq proxy")?;

            debug!("successfully created tun interface");

            let mut events = onion_tunnel.get_bootstrap_events();
            tokio::spawn(async move {
                debug!("starting bootstrap event listening ...");
                loop {
                    tokio::select! {
                        event_result = events.next() => {
                            match event_result {
                                Some(bootstrap_status) => {
                                    let mobile = OnionmasqMobile::get();
                                    let event = TunnelEvent::bootstrap(bootstrap_status);
                                    if let Err(e) = mobile.on_tunnel_event(event) {
                                        error!("Failed to update bootstrap status: {e}");
                                    }
                                },
                                None => return,
                            }
                        },
                        Some(_) = bootstrap_done_receiver.recv() => {
                            info!("Bootstrapping completed.");
                            let mobile = OnionmasqMobile::get();
                            let event = TunnelEvent::bootstrap_done();
                            if let Err(e) = mobile.on_tunnel_event(event) {
                                error!("Failed to update bootstrap status: {e}");
                            }
                            return;
                        }
                    }
                }
            });

            debug!("starting onionmasq...");

            // (This call is blocking, and only returns in case of error or external shutdown.)
            let result = onion_tunnel.run().await;
            if let Err(ref e) = result {
                error!("OnionTunnel failed: {e}");
            }
            result?;

            debug!("stopped onionmasq...");
            anyhow::Ok(())
        });

        // Make sure to clear out the command sender after a stop, even in a failure condition.
        let _ = self.command_sender.lock().unwrap().take();

        ret
    }

    fn stop_proxy(&self) -> Result<()> {
        debug!("closing proxy...");

        if let Some(sender) = self.command_sender.lock().unwrap().take() {
            // This hangs up the command control channel, which should cause the proxy to quit.
            mem::drop(sender);
        } else {
            warn!("Attempted to stop proxy when it wasn't running!");
        }

        Ok(())
    }

    /// Send a `TunnelCommand` to the running proxy.
    fn send_command(&self, cmd: TunnelCommand) -> Result<()> {
        if let Some(sender) = self.command_sender.lock().unwrap().as_mut() {
            if sender.send(cmd).is_err() {
                // This is a potentially normal occurrence if a stop command or failure races with the command.
                warn!("Sending tunnel command to running proxy failed!");
                return Err(ProxyStoppedError.into());
            }
        } else {
            warn!("Attempted to send {cmd:?} while the proxy was stopped!");
            return Err(ProxyStoppedError.into());
        }
        Ok(())
    }

    /// Set the country code that proxied connections should use.
    ///
    /// You can clear it back to "no country code" by passing in `null`.
    fn set_country_code(&self, cc: JString<'_>) -> Result<()> {
        let cc = if cc.is_null() {
            None
        } else {
            let mut local_env = self
                .jvm
                .attach_current_thread()
                .context("failed to attach current thread")?;

            let java_str = local_env
                .get_string(&cc)
                .context("failed to get country code string")?;

            Some(java_str.to_string_lossy().to_string())
        };

        let parsed = cc
            .as_ref()
            .map(|x| {
                CountryCode::from_str(x).map_err(|_| CountryCodeError {
                    code: x.to_string(),
                })
            })
            .transpose()?;

        debug!("updating country code to {parsed:?}");

        *self.country_code.write().unwrap() = parsed;

        Ok(())
    }

    /// Notify the Java layer about a new asynchronous event.
    fn on_tunnel_event(&self, evt: TunnelEvent) -> Result<()> {
        let serialized = serde_json::to_string(&evt).context("failed to serialize TunnelEvent")?;

        let local_env = self
            .jvm
            .attach_current_thread()
            .context("failed to attach current thread")?;

        let jstring = local_env
            .new_string(serialized)
            .context("JString creation failed")?;

        self.call_jni_class_method::<()>(
            "postEvent",
            "(Ljava/lang/String;)V",
            &[JValue::from(&jstring)],
        )
        .context("postEvent call failed")?;
        Ok(())
    }

    /// Get the number of bytes sent by an app ID since the last counter reset.
    ///
    /// None if the app doesn't have a counter yet.
    pub fn app_tx_counter(&self, app_id: u64) -> u64 {
        self.bandwidth_counters
            .get(&app_id)
            .map(|x| x.bytes_tx())
            .unwrap_or(0)
    }

    /// Get the number of bytes received by an app ID since the last counter reset.
    pub fn app_rx_counter(&self, app_id: u64) -> u64 {
        self.bandwidth_counters
            .get(&app_id)
            .map(|x| x.bytes_rx())
            .unwrap_or(0)
    }

    /// Reset the global bandwidth counter, and all per-app bandwidth counters.
    pub fn reset_counters(&self) {
        BandwidthCounter::global().reset();
        for entry in self.bandwidth_counters.iter() {
            entry.value().reset();
        }
    }

    /// Set a path that the tunnel should log `.pcap` files of all traffic to.
    ///
    /// This only takes effect when starting the tunnel afresh.
    pub fn set_pcap_path(&self, path: Option<String>) {
        *self.pcap_path.lock().unwrap() = path;
    }

    /// Get a JNI environment by attaching the current thread to the JVM, if it isn't
    /// already attached.
    pub fn get_jni_env(&self) -> Result<AttachGuard> {
        let guard = self
            .jvm
            .attach_current_thread()
            .context("failed to attach current thread")?;
        Ok(guard)
    }

    /// Call a method on the `org.torproject.onionmasq.OnionMasqJni` class, and try to convert the
    /// returned Java type into the Rust type `T`.
    //
    // TODO(eta): I'm pretty sure the lifetimes here are slightly unsound.
    pub fn call_jni_class_method<T>(
        &self,
        method: &str,
        signature: &str,
        args: &[JValue],
    ) -> Result<T>
    where
        T: for<'a> TryFrom<JValueOwned<'a>, Error = jni::errors::Error>,
    {
        let mut guard = self.get_jni_env()?;
        let result = guard
            .call_static_method(&self.class_ref, method, signature, args)
            .with_context(|| format!("calling method {method} on OnionMasqJni failed"))?;
        let ret = result.try_into().with_context(|| {
            // Include the type name to help figure out what went wrong.
            format!(
                "failed to convert return value of {method} on OnionMasqJni to {}",
                std::any::type_name::<T>()
            )
        })?;
        Ok(ret)
    }
}
