#![warn(clippy::unwrap_used)]
#![warn(clippy::expect_used)]

pub mod accounting;
pub mod config;
mod device;
mod dns;
pub mod errors;
mod isolation;
mod parser;
mod proxy;
mod runtime;
pub mod scaffolding;
mod socket;
mod tuntap;

use arti_client::status::BootstrapEvents;
pub use arti_client::status::BootstrapStatus;
pub use arti_client::CountryCode;
use arti_client::TorClient;
use device::VirtualDevice;
use dns::{DnsManager, LockedDnsCookies};
use futures::future::Either;
use futures::stream::StreamExt;
use onionmasq_device_testing::{AsyncDevice, PollableAsyncDevice};
use proxy::ArtiProxy;
use runtime::OnionTunnelArtiRuntime;
use smoltcp::{
    iface::Interface,
    phy::Medium,
    wire::{IpAddress, IpCidr, Ipv4Address},
};
use std::collections::HashMap;
#[cfg(feature = "pluggable-transports")]
use std::fs;
use std::fs::File;
#[cfg(feature = "pluggable-transports")]
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::pin::Pin;
use std::sync::Mutex;
use std::{io, sync::Arc};
use tracing::{debug, error, trace, warn};

use crate::tuntap::AsyncTunTapInterface;
use crate::{parser::Parser, socket::TcpSocket};

use crate::config::TunnelConfig;
use crate::errors::{TunnelError, TunnelResult};
use crate::isolation::OnionIsolationKey;
use crate::parser::ParseResult;
use crate::scaffolding::TunnelCommand;
pub use crate::scaffolding::TunnelScaffolding;
pub use arti_client::config as tor_config;
#[cfg(feature = "pluggable-transports")]
use arti_client::config::CfgPath;
pub use arti_client::{config::TorClientConfigBuilder, TorClientConfig};
use futures::FutureExt;
#[cfg(feature = "pluggable-transports")]
use onionmasq_pt_wrapper::TransportKind;
use rand::RngCore;
use smoltcp::iface::{SocketHandle, SocketSet};
use smoltcp::phy::{PcapMode, PcapWriter};
use smoltcp::socket::tcp::State as TcpState;
use smoltcp::socket::Socket;
use smoltcp::time::{Duration, Instant};
pub use smoltcp::wire::IpEndpoint;
#[cfg(feature = "pluggable-transports")]
use std::net::TcpListener;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::sync::Notify;
use tokio::task::JoinHandle;
use tor_geoip::HasCountryCode;
#[cfg(feature = "pluggable-transports")]
use tor_ptmgr::ipc::PtClientMethod;
#[cfg(feature = "pluggable-transports")]
use tor_socksproto::SocksVersion;

const ASYNC_DEVICE_BUFFER_SIZE: usize = 1500 * 4;

/// The type of the device smoltcp uses to send and receive packets.
///
/// Uses `VirtualDevice` so that we can inspect the packets before passing them on.
type TunnelDevice = VirtualDevice<AsyncDevice<ASYNC_DEVICE_BUFFER_SIZE>>;

/// A wrapper for the `TunnelDevice` that might wrap it in a `PcapWriter`, or not.
enum TunnelDeviceWrapper {
    Pcap(PcapWriter<TunnelDevice, File>),
    Raw(TunnelDevice),
}

impl TunnelDeviceWrapper {
    fn inner(&mut self) -> &mut TunnelDevice {
        match self {
            TunnelDeviceWrapper::Raw(ref mut d) => d,
            TunnelDeviceWrapper::Pcap(ref mut d) => d.get_mut(),
        }
    }
}

/// Information about a proxy task.
struct ProxyInformation {
    /// The join handle for the proxy task.
    handle: JoinHandle<()>,
    /// The isolation key used for the connection.
    isolation_key: u64,
}

pub struct OnionTunnel<S = (), D = AsyncTunTapInterface>
where
    S: TunnelScaffolding,
{
    /// The smoltcp interface type.
    interface: Interface,
    /// The smoltcp socket set.
    ///
    /// This is wrapped inside a mutex to allow other tasks to send and receive on the sockets.
    socket_set: Arc<Mutex<SocketSet<'static>>>,
    /// The device smoltcp uses to send and receive packets.
    device: TunnelDeviceWrapper,
    /// The pollable async half of the `AsyncDevice`.
    pollable_adev: PollableAsyncDevice<D, ASYNC_DEVICE_BUFFER_SIZE>,
    arti: TorClient<OnionTunnelArtiRuntime<S>>,
    dns_cookies: Option<LockedDnsCookies>,
    scaffolding: Arc<S>,
    /// A `Notify` used to make the main loop do a poll.
    poll_notify: Arc<Notify>,
    /// An index of smoltcp socket handles to the `proxy()` task join handle handling that socket.
    proxy_handles: HashMap<SocketHandle, ProxyInformation>,
    /// The join handle for the `dns()` task.
    dns_join_handle: Option<JoinHandle<()>>,
    /// The current 'global isolation epoch', which we increment every time a new identity is
    /// requested.
    ///
    /// NOTE(eta): This should be incremented with wrapping logic, to prevent panic on overflow.
    global_epoch: u16,
    /// A map of isolation keys to their individual 'app epoch', which works like the global one
    /// except per-app. This is only touched when we actually increment the epoch, to save space.
    ///
    /// NOTE(eta): This should be incremented with wrapping logic, to prevent panic on overflow.
    app_epochs: HashMap<u64, u16>,
}

impl<S> OnionTunnel<S, AsyncTunTapInterface>
where
    S: TunnelScaffolding,
{
    /// Create an `OnionTunnel` from a file descriptor, representing a TUN device.
    pub async fn create_with_fd(
        scaffolding: S,
        fd: i32,
        config: TunnelConfig,
    ) -> TunnelResult<Self> {
        debug!("OnionTunnel::create_with_fd({fd})");
        let tuntap = match AsyncTunTapInterface::new_from_fd(fd) {
            Ok(v) => {
                debug!("successfully created tun interface");
                v
            }
            Err(e) => {
                warn!("couldn't create tun interface: {}", e);
                return Err(TunnelError::TunSetup(e));
            }
        };

        Self::create_custom(scaffolding, tuntap, config).await
    }

    /// Create an `OnionTunnel` from a TUN device name.
    #[cfg(any(target_os = "linux", target_os = "android"))]
    pub async fn new(scaffolding: S, iface_name: &str, config: TunnelConfig) -> TunnelResult<Self> {
        debug!("OnionTunnel::new({iface_name:?})");
        let tun =
            AsyncTunTapInterface::new(iface_name, Medium::Ip).map_err(TunnelError::TunSetup)?;

        Self::create_custom(scaffolding, tun, config).await
    }
}

impl<S, D> OnionTunnel<S, D>
where
    S: TunnelScaffolding,
    D: AsyncRead + AsyncWrite,
{
    /// Create an `OnionTunnel` with a completely custom device `D`.
    ///
    /// # Warning
    ///
    /// The device must behave in accordance with the restrictions on `AsyncDevice::new`.
    /// In particular, it MUST read and write whole packets at a time (i.e. one call to `poll_read`
    /// must read a whole packet or nothing, and one call to `poll_write` must write a whole packet
    /// or nothing). Splitting reads and writes will break the tunnel.
    pub async fn create_custom(
        scaffolding: S,
        device: D,
        config: TunnelConfig,
    ) -> TunnelResult<Self> {
        // FIXME(eta): MTU?
        let (adev, pollable_adev) = AsyncDevice::new(device, Medium::Ip, 1500);

        let scaffolding = Arc::new(scaffolding);

        let rt = runtime::make_runtime(Arc::clone(&scaffolding)).await;
        let arti_config = config.arti_config()?;
        let arti_builder = TorClient::with_runtime(rt).config(arti_config);
        let arti = arti_builder
            .create_unbootstrapped()
            .map_err(TunnelError::ArtiSetup)?;

        #[cfg(feature = "pluggable-transports")]
        {
            trace!("Attempting to enable pluggable transports.");

            let pt_scaff = scaffolding.clone();
            onionmasq_pt_wrapper::set_protect_callback(move |fd| {
                debug!("PT protect() for {fd}");
                // FIXME(eta): either provide the actual address here, or just change the interface.
                //             I don't think this is actually used any more now that we have the
                //             whole fwmark solution thing.
                let fake = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0);
                pt_scaff.protect(fd, &fake)
            });

            let pt_dir = match config.pt_dir {
                Some(v) => v,
                None => {
                    if let Some(state_dir) = config.state_dir {
                        state_dir.join("pt_state")
                    } else {
                        // Attempt to figure out the Arti state directory.
                        let state_dir = CfgPath::new("${ARTI_LOCAL_DATA}".into())
                            .path()
                            .map_err(TunnelError::ArtiDefaultState)?;
                        state_dir.join("pt_state")
                    }
                }
            };
            fs::create_dir_all(&pt_dir).map_err(|e| TunnelError::Directory {
                dir: pt_dir.clone(),
                err: e,
            })?;
            let pt_dir = pt_dir.to_string_lossy();

            // HACK(eta): Find a free port by binding a TCP listener, getting the port, and then
            //            immediately dropping it and passing that address to the PT.
            //            This is very racy, but we don't have a way to get the bound port out of
            //            the PT yet.
            let listener = TcpListener::bind("127.0.0.1:0").map_err(TunnelError::PtSetup)?;
            let socket_addr = listener.local_addr().map_err(TunnelError::PtSetup)?;
            drop(listener);
            onionmasq_pt_wrapper::start_transport(TransportKind::Obfs4, socket_addr, &pt_dir)
                .map_err(TunnelError::PtSetupGo)?;

            #[allow(clippy::unwrap_used)]
            arti.ptmgr().add_unmanaged(
                "obfs4".parse().unwrap(),
                PtClientMethod::new(SocksVersion::V5, socket_addr),
            );

            trace!("Registered obfs4 at {socket_addr}");
        }

        let mut device = VirtualDevice::new(adev);

        let mut iface_config = smoltcp::iface::Config::new(smoltcp::wire::HardwareAddress::Ip);
        iface_config.random_seed = rand::thread_rng().next_u64();

        let mut iface = Interface::new(iface_config, &mut device, Instant::now());

        iface.set_any_ip(true);

        // FIXME(eta): I think we just need to insert these to make packets flow, and they aren't
        //             actually relevant for anything. This assumption should probably be checked.
        #[allow(clippy::expect_used)]
        iface
            .routes_mut()
            .add_default_ipv4_route(Ipv4Address::new(0, 0, 0, 1))
            .expect("smoltcp routing table full?");

        iface.update_ip_addrs(|ip_addrs| {
            #[allow(clippy::expect_used)]
            ip_addrs
                .push(IpCidr::new(IpAddress::v4(0, 0, 0, 1), 0))
                .expect("smoltcp IP address vector not big enough?")
        });

        let socket_set = Arc::new(Mutex::new(SocketSet::new(vec![])));

        let device = match config.pcap_path {
            Some(v) => {
                error!("*** WARNING WARNING WARNING ***");
                error!("All tunnel traffic is being logged to {}. Do not use the tunnel for anything sensitive!", v.to_string_lossy());
                error!("*** WARNING WARNING WARNING ***");

                let file = File::create(v).map_err(TunnelError::PcapSetup)?;
                let pcap = PcapWriter::new(device, file, PcapMode::Both);
                TunnelDeviceWrapper::Pcap(pcap)
            }
            None => TunnelDeviceWrapper::Raw(device),
        };

        Ok(Self {
            interface: iface,
            device,
            socket_set,
            arti,
            pollable_adev,
            dns_cookies: None,
            scaffolding,
            poll_notify: Arc::new(Notify::new()),
            proxy_handles: HashMap::new(),
            dns_join_handle: None,
            global_epoch: 0,
            app_epochs: Default::default(),
        })
    }

    fn proxy(
        &mut self,
        socket: TcpSocket,
        source: IpEndpoint,
        dest: IpEndpoint,
    ) -> TunnelResult<ProxyInformation> {
        let app_id = match self.scaffolding.isolate(source, dest, 6) {
            Ok(v) => v,
            Err(e) => {
                socket.abort();
                return Err(TunnelError::ScaffoldingIsolation(e));
            }
        };
        let country_code = self.scaffolding.locate(source, dest, app_id);
        let isolation = OnionIsolationKey {
            app_id,
            app_epoch: self.app_epochs.get(&app_id).copied().unwrap_or(0),
            global_epoch: self.global_epoch,
        };
        let handle = socket.handle();
        trace!(
            "socket {}: new proxied connection {}->{}, isolation {:?}, country {:?}",
            handle,
            source,
            dest,
            isolation,
            country_code
        );
        let scaff = Arc::clone(&self.scaffolding);
        #[allow(clippy::unwrap_used)]
        let mut proxy = ArtiProxy {
            socket: socket.clone(),
            arti: self.arti.clone(),
            cookies: self.dns_cookies.as_ref().unwrap().clone(),
            isolation_key: isolation,
            country_code,
            scaffolding: Arc::clone(&scaff),
            source,
            dest,
        };
        let handle = tokio::spawn(async move {
            match proxy.start().await {
                Ok(_) => {
                    trace!("socket {}: connection closed", handle);
                    socket.close();
                    scaff.on_socket_close(source, dest, None);
                }
                Err(e) => {
                    trace!("socket {}: connection failed: {}", handle, e);
                    socket.abort();
                    scaff.on_socket_close(source, dest, Some(e));
                }
            }
        });
        Ok(ProxyInformation {
            handle,
            isolation_key: app_id,
        })
    }

    fn dns(&mut self) -> TunnelResult<JoinHandle<()>> {
        // Spawn the DNS manager service.
        let mut dns_manager = DnsManager::new(self)?;
        self.dns_cookies = Some(dns_manager.cookies());
        Ok(tokio::spawn(async move { dns_manager.start().await }))
    }

    fn poll(&mut self) -> bool {
        #[allow(clippy::unwrap_used)]
        let mut socket_set = self.socket_set.lock().unwrap();
        match self.device {
            TunnelDeviceWrapper::Pcap(ref mut d) => {
                self.interface.poll(Instant::now(), d, &mut socket_set)
            }
            TunnelDeviceWrapper::Raw(ref mut d) => {
                self.interface.poll(Instant::now(), d, &mut socket_set)
            }
        }
    }

    fn poll_delay(&mut self) -> Option<Duration> {
        #[allow(clippy::unwrap_used)]
        let socket_set = self.socket_set.lock().unwrap();
        self.interface.poll_delay(Instant::now(), &socket_set)
    }

    fn handle_command(&mut self, cmd: TunnelCommand) {
        match cmd {
            TunnelCommand::RefreshCircuits => {
                trace!("Refreshing circuits.");
                self.global_epoch = self.global_epoch.wrapping_add(1);
                #[allow(clippy::unwrap_used)]
                for (_, s) in self.socket_set.lock().unwrap().iter_mut() {
                    if let Socket::Tcp(s) = s {
                        s.abort();
                    }
                }
            }
            TunnelCommand::RefreshCircuitsForApp { isolation_key } => {
                trace!("Refreshing circuits for app {isolation_key}.");
                let entry = self.app_epochs.entry(isolation_key).or_insert(0);
                *entry = entry.wrapping_add(1);
                #[allow(clippy::unwrap_used)]
                for (h, s) in self.socket_set.lock().unwrap().iter_mut() {
                    if let Some(info) = self.proxy_handles.get(&h) {
                        if let Socket::Tcp(s) = s {
                            if info.isolation_key == isolation_key {
                                s.abort();
                            }
                        }
                    }
                }
            }
        }
    }

    pub async fn run(&mut self) -> TunnelResult<()> {
        debug!("Starting tunnel...");
        let tc = self.arti.clone();
        let mut bootstrap_task = tokio::spawn(async move { tc.bootstrap().await }).fuse();

        // Spawn the DNS manager service.
        self.dns_join_handle = Some(self.dns()?);

        // Get a command stream, if there is one.
        let mut command_stream = Pin::from(
            self.scaffolding
                .command_stream()
                .unwrap_or_else(|| Box::new(futures::stream::pending())),
        );

        // reusable buffer
        let mut sockets = vec![];

        loop {
            // Do one poll to make sure we enqueue packets in the `VirtualDevice`'s queue.
            self.poll();

            // Get all the enqueued packets, extract things that look like TCP SYN, and
            // make `TcpSocket`s for those packets.
            self.device.inner().inspect_packets(|packet| {
                if let Some(tcp_socket) = Parser::parse(packet) {
                    sockets.push(tcp_socket);
                }
            });

            // Also, make sure we remove old sockets that are closed.
            {
                #[allow(clippy::unwrap_used)]
                let mut socket_set = self.socket_set.lock().unwrap();

                let handles_to_remove = socket_set
                    .iter()
                    .filter_map(|(h, s)| {
                        if let Socket::Tcp(tsock) = s {
                            if tsock.state() == TcpState::Closed {
                                if let Some(task) = self.proxy_handles.get(&h) {
                                    if task.handle.is_finished() {
                                        trace!("socket {}: garbage collecting closed socket", h);
                                        self.proxy_handles.remove(&h);
                                    } else {
                                        trace!("socket {}: waiting for task to finish still", h);
                                        return None;
                                    }
                                } else {
                                    trace!(
                                        "socket {}: garbage collecting closed socket (no task)",
                                        h
                                    );
                                }
                                return Some(h);
                            }
                        }
                        None
                    })
                    .collect::<Vec<_>>();

                for hdl in handles_to_remove {
                    socket_set.remove(hdl);
                }
            }

            // Start proxying all the sockets we made. This registers them with `smoltcp`, ensuring
            // there's a corresponding listening socket for every TCP SYN we observed.`
            for ParseResult {
                source,
                dest,
                socket,
            } in sockets.drain(..)
            {
                let socket = TcpSocket::new(self, socket);
                let handle = socket.handle();
                match self.proxy(socket, source, dest) {
                    Ok(v) => {
                        if self.proxy_handles.insert(handle, v).is_some() {
                            // I think smoltcp's socket handles are based on the index of the socket
                            // in smoltcp's internal backing store, so they could overlap like this.
                            // Let's just log this for now.
                            error!(
                                "{}->{}: displaced an older socket in proxy_handles!",
                                source, dest
                            );
                        }
                    }
                    Err(e) => {
                        trace!("{}->{}: failed to proxy: {}", source, dest, e);
                        // TODO(eta): maybe this error type should be `TunnelError`
                        self.scaffolding.on_socket_close(
                            source,
                            dest,
                            Some(io::Error::new(io::ErrorKind::Other, e.to_string())),
                        );
                    }
                }
            }

            // Now we've registered the sockets, we can let smoltcp have the packets we just
            // looked at.
            self.device.inner().done_inspecting();

            self.poll();

            let poll_delay = if self.device.inner().check_exhausted() {
                // Our device's transmit buffer is full, so there's no point in asking smoltcp
                // for a poll delay -- we need to wait for the buffer to clear first.
                //
                // If we don't check this, we can get into really sad situations where smoltcp
                // has data to send -- so will always return 0 -- but there's no point in polling
                // again until the AsyncDevice shifts some packets around (if we do, we just burn
                // CPU, which isn't ideal on a mobile device...)
                Either::Right(futures::future::pending::<()>())
            } else {
                self.poll_delay()
                    .map(|x| Either::Left(tokio::time::sleep(std::time::Duration::from(x))))
                    .unwrap_or_else(|| Either::Right(futures::future::pending::<()>()))
            };

            tokio::select! {
                _ = poll_delay => {
                    // poll delay elapsed
                }
                _ = self.poll_notify.notified() => {
                    // some other task wanted us to do a poll
                }
                cmd = command_stream.next() => {
                    match cmd {
                        Some(cmd) => self.handle_command(cmd),
                        None => {
                            debug!("Command stream hung up; stopping proxy.");
                            return Ok(());
                        }
                    }
                }
                res = self.pollable_adev.next() => {
                    // PollableAsyncDevice did something; check for an I/O error.
                    // If it returns None, panic; this can only happen if we polled it after
                    // it returned an error (impossible).
                    #[allow(clippy::expect_used)]
                    res
                        .expect("PollableAsyncDevice returned None")
                        .map_err(TunnelError::TunIo)?;
                }
                bootstrap_ret = &mut bootstrap_task => {
                    match bootstrap_ret {
                        Ok(Ok(_)) => {
                            debug!("Arti bootstrapping complete.");
                            self.scaffolding.on_bootstrapped();

                            let netdir = match self.arti.dirmgr().timely_netdir() {
                                Ok(v) => v,
                                Err(e) => {
                                    warn!("Couldn't get a netdir after completed bootstrap: {e}");
                                    continue;
                                }
                            };

                            let mut country_map: HashMap<String, usize> = HashMap::new();
                            let mut count = 0;
                            for relay in netdir.relays() {
                                // Use port 443 as a proxy for whether this is an exit at all.
                                if !relay.supports_exit_port_ipv4(443) {
                                    continue;
                                }
                                if let Some(cc) = relay.country_code() {
                                    *country_map.entry(cc.to_string()).or_insert(0) += 1;
                                    count += 1;
                                }
                            }
                            debug!("Found {count} relays with country codes attached");
                            self.scaffolding.on_directory(country_map);
                        },
                        Ok(Err(e)) => return Err(TunnelError::BootstrapFailure(e)),
                        Err(e) => {
                            // Propagate a panic in the bootstrapping to this thread.
                            if e.is_panic() {
                                std::panic::resume_unwind(e.into_panic());
                            }
                        }
                    }
                }
            }
        }
    }

    pub fn get_bootstrap_events(&mut self) -> BootstrapEvents {
        self.arti.bootstrap_events()
    }
}

impl<S, D> Drop for OnionTunnel<S, D>
where
    S: TunnelScaffolding,
{
    fn drop(&mut self) {
        debug!(
            "Shutting down proxy on drop ({} handles)",
            self.proxy_handles.len()
        );
        if let Some(jh) = self.dns_join_handle.take() {
            jh.abort();
        }
        for (_, task) in self.proxy_handles.drain() {
            task.handle.abort()
        }
    }
}
