//! The `onionmasq-apple` crate provides a C binding for the Rust `onion-tunnel` crate.

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

use crate::errors::{CountryCodeError, ProxyStoppedError};
use crate::ffi::{EventCallback, LoggingCallback, ReaderCallback, WriterCallback};
use crate::nepackettunnelflow::NePacketTunnelFlow;
use crate::scaffolding::AppleScaffolding;
use anyhow::{Context, Result};
use dashmap::DashMap;
use ffi_event::TunnelEvent;
use futures::StreamExt;
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::ffi::{c_char, CStr, CString};
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::{Layer, Subscriber};
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;

static ONIONMASQ_APPLE_SINGLETON: OnceCell<OnionmasqApple> = OnceCell::new();

// We want to avoid logging detailed connection-level stuff in production builds to avoid
// embarrassing privacy issues.
#[cfg(feature = "verbose")]
static LOG_FILTER: &str = "info,onionmasq_apple=trace,onion_tunnel=trace,arti_client=debug,tor_chanmgr=debug,tor_proto=debug,smoltcp=debug";
#[cfg(not(feature = "verbose"))]
static LOG_FILTER: &str = "info,onionmasq_apple=debug";

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

/// A wrapper to run and control `onion-tunnel` from Apple platforms.
///
/// This `struct` is the main entrypoint for FFI interactions to and from the C interface.
pub struct OnionmasqApple {
    /// Event callback.
    event_fn: EventCallback,
    /// A tokio Runtime.
    rt: Runtime,
    /// 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 OnionmasqApple {
    /// Initialise the FFI layer, creating a singleton `OnionmasqApple` 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(event_fn: EventCallback, log_fn: LoggingCallback) -> Result<&'static Self> {
        panic_handling::register_panic_hook();

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

        // FIXME(eta): We could just pass the C thing in directly maybe.
        let log_fn = Arc::new(move |log: *const c_char| {
            (log_fn)(log);
        });
        let log = Layer::new().with_writer(move || CallbackWriter::new(log_fn.clone()));

        // Set up logging via log callback:
        Subscriber::builder()
            .with_env_filter(LOG_FILTER)
            .finish()
            .with(log)
            .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")?;

        let ret = Self {
            event_fn,
            rt,
            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_APPLE_SINGLETON.set(ret).is_err() {
            panic!("OnionmasqApple::new() called concurrently somehow!");
        }
        Ok(Self::get())
    }

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

    /// Run the `onion-tunnel` proxy.
    pub fn run_proxy(
        &self,
        cache_dir: &str,
        data_dir: &str,
        reader: ReaderCallback,
        writer: WriterCallback,
    ) -> Result<()> {
        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(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_custom(
                AppleScaffolding {
                    bootstrap_done: bootstrap_done_sender,
                    command_receiver: Mutex::new(Some(command_receiver)),
                    counters,
                    country_code,
                },
                NePacketTunnelFlow::new(reader, writer),
                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 = OnionmasqApple::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 = OnionmasqApple::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: *const c_char) -> Result<()> {
        let cc = if cc.is_null() {
            None
        } else {
            let str = unsafe { CStr::from_ptr(cc) }
                .to_str()
                .context("failed to get country code string")?;

            Some(str.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 C 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 cstr = CString::new(serialized)
            .context("failed to create CString from serialized TunnelEvent")?;
        (self.event_fn)(cstr.as_ptr());

        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;
    }
}

/// This implements `Write` and forwards all written data to a C callback function.
///
// NOTE(eta): I'm pretty sure I wrote this during the Brussels hacking session over tla's
//            shoulder. Looking back at it now, there's a few issues with it...
#[derive(Clone)]
struct CallbackWriter<F> {
    func: Arc<F>,
}

impl<F> CallbackWriter<F>
where
    F: Fn(*const c_char) + Send + Sync + 'static,
{
    pub fn new(callback: Arc<F>) -> Self {
        CallbackWriter { func: callback }
    }
}

impl<F> std::io::Write for CallbackWriter<F>
where
    F: Fn(*const c_char) + Send + Sync + 'static,
{
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        // FIXME(eta): should we not buffer each line?

        let cstr = CString::new(buf.to_owned())
            .context("failed to create CString from log string")
            .unwrap();

        (self.func)(cstr.as_ptr());

        Ok(buf.len())
    }

    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}
