//! The client side: a connection to another device's app on its own thread,
//! reconnecting when it drops. Messages from the host reach the workspace as
//! [`ClientEvent`]s.
//!
//! A host can have several routes. The client connects over the first one,
//! fastest first, that answers with the right key, and while it's on a
//! slower one it keeps trying the faster ones and moves as soon as one works.

use std::{
    collections::HashMap,
    io::ErrorKind,
    net::{TcpStream, ToSocketAddrs},
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicU64, Ordering},
        mpsc,
    },
    thread,
    time::{Duration, Instant},
};

use tungstenite::{Message, WebSocket};

use super::{
    ClientMessage, DEFAULT_PORT, PROTOCOL, RemoteFileDiff, Route, RouteKind, ServerMessage, tailnet_peers, tailnet_routes,
};

/// How long to wait for a route that goes over the internet or a VPN.
const CONNECT_TIMEOUT: Duration = Duration::from_secs(6);
/// A LAN address answers quickly when it's there; away from home it isn't.
const LAN_TIMEOUT: Duration = Duration::from_secs(2);
const POLL: Duration = Duration::from_millis(40);

/// How the connection checks on itself.
#[derive(Clone, Copy)]
struct Timing {
    /// How often to ping the host, which answers even when idle.
    ping: Duration,
    /// Hearing nothing this long means the route is gone, as when this
    /// device left the network it went over.
    silence: Duration,
    /// How often to try the faster routes while on a slower one.
    probe: Duration,
}

const TIMING: Timing =
    Timing { ping: Duration::from_secs(15), silence: Duration::from_secs(45), probe: Duration::from_secs(60) };

#[derive(Clone, Debug, PartialEq)]
pub enum ConnectionStatus {
    Connecting,
    /// Connected to the host of that name.
    Connected(String),
    Failed(String),
    /// The host refused the key; no more retries until the settings change.
    Denied(String),
}

pub enum ClientEvent {
    Status(ConnectionStatus),
    /// The route the connection goes over, sent before each `Connected`.
    Route(Route),
    Message(ServerMessage),
}

pub struct RemoteClient {
    outbox: mpsc::Sender<String>,
    stop: Arc<AtomicBool>,
    routes: Arc<Mutex<Vec<Route>>>,
    pub events: smol::channel::Receiver<ClientEvent>,
    /// Diffs asked for and not answered yet, by request.
    diffs: Mutex<HashMap<u64, smol::channel::Sender<Result<Vec<RemoteFileDiff>, String>>>>,
    next_request: AtomicU64,
}

/// A way to the host that works off the UI thread.
#[derive(Clone)]
pub struct ClientSender(mpsc::Sender<String>);

impl ClientSender {
    pub fn send(&self, message: &ClientMessage) {
        if let Ok(text) = serde_json::to_string(message) {
            let _ = self.0.send(text);
        }
    }
}

impl RemoteClient {
    pub fn sender(&self) -> ClientSender {
        ClientSender(self.outbox.clone())
    }

    /// Connects over `routes`, tried in order, and keeps reconnecting until
    /// dropped. When none of them answers, it looks for a device called one
    /// of `names` on this device's tailnet.
    pub fn connect(routes: Vec<Route>, key: String, name: String, names: Vec<String>) -> Self {
        Self::connect_timed(routes, key, name, names, TIMING)
    }

    fn connect_timed(routes: Vec<Route>, key: String, name: String, names: Vec<String>, timing: Timing) -> Self {
        let stop = Arc::new(AtomicBool::new(false));
        let routes = Arc::new(Mutex::new(routes));
        let (outbox, inbox) = mpsc::channel::<String>();
        let (events_tx, events) = smol::channel::unbounded();
        {
            let stop = stop.clone();
            let routes = routes.clone();
            let _ = thread::Builder::new().name("remote-client".into()).spawn(move || {
                let hello = Hello { key, name, names, routes };
                let mut backoff = Duration::from_secs(2);
                while !stop.load(Ordering::Relaxed) {
                    let _ = events_tx.send_blocking(ClientEvent::Status(ConnectionStatus::Connecting));
                    let attempt = match open_any(&hello.routes(), &hello) {
                        Err(Failure::Other(reason)) => discover(&hello).ok_or(Failure::Other(reason)),
                        attempt => attempt,
                    };
                    match attempt {
                        Ok(link) => {
                            backoff = Duration::from_secs(2);
                            // Whatever was asked while offline is stale now.
                            while inbox.try_recv().is_ok() {}
                            announce(&events_tx, &link);
                            let reason = run(link, &hello, &inbox, &events_tx, &stop, timing);
                            if stop.load(Ordering::Relaxed) {
                                break;
                            }
                            let _ = events_tx.send_blocking(ClientEvent::Status(ConnectionStatus::Failed(reason)));
                        }
                        Err(Failure::Denied(reason)) => {
                            let _ = events_tx.send_blocking(ClientEvent::Status(ConnectionStatus::Denied(reason)));
                            break;
                        }
                        Err(Failure::Other(reason)) => {
                            let _ = events_tx.send_blocking(ClientEvent::Status(ConnectionStatus::Failed(reason)));
                        }
                    }
                    // Wait before retrying, a little longer each time.
                    let mut waited = Duration::ZERO;
                    while waited < backoff && !stop.load(Ordering::Relaxed) {
                        thread::sleep(Duration::from_millis(200));
                        waited += Duration::from_millis(200);
                    }
                    backoff = (backoff * 2).min(Duration::from_secs(30));
                }
            });
        }
        Self { outbox, stop, routes, events, diffs: Mutex::default(), next_request: AtomicU64::new(1) }
    }

    /// Asks the host for the diff between two of a thread's snapshots. The
    /// answer arrives on the returned channel once [`Self::resolve_diff`]
    /// gets it; older hosts never answer.
    pub fn request_diff(
        &self,
        thread: String,
        base: String,
        head: String,
        path: Option<String>,
        ignore_whitespace: bool,
    ) -> (u64, smol::channel::Receiver<Result<Vec<RemoteFileDiff>, String>>) {
        let request = self.next_request.fetch_add(1, Ordering::Relaxed);
        let (tx, rx) = smol::channel::bounded(1);
        self.diffs.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).insert(request, tx);
        self.send(&ClientMessage::Diff { request, thread, base, head, path, ignore_whitespace });
        (request, rx)
    }

    /// Hands a diff reply to the request waiting for it.
    pub fn resolve_diff(&self, request: u64, result: Result<Vec<RemoteFileDiff>, String>) {
        let waiting = self.diffs.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).remove(&request);
        if let Some(waiting) = waiting {
            let _ = waiting.try_send(result);
        }
    }

    /// Forgets a request that stopped waiting.
    pub fn forget_diff(&self, request: u64) {
        self.diffs.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).remove(&request);
    }

    /// Changes the routes without dropping the connection; the next try,
    /// and the next look for a faster route, use them.
    pub fn set_routes(&self, routes: Vec<Route>) {
        *self.routes.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) = routes;
    }

    pub fn send(&self, message: &ClientMessage) {
        if let Ok(text) = serde_json::to_string(message) {
            let _ = self.outbox.send(text);
        }
    }
}

impl Drop for RemoteClient {
    fn drop(&mut self) {
        self.stop.store(true, Ordering::Relaxed);
    }
}

enum Failure {
    Denied(String),
    Other(String),
}

/// What every try says, and where it may go.
#[derive(Clone)]
struct Hello {
    key: String,
    name: String,
    /// What the host is called, to look it up on the tailnet by.
    names: Vec<String>,
    routes: Arc<Mutex<Vec<Route>>>,
}

impl Hello {
    fn routes(&self) -> Vec<Route> {
        self.routes.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).clone()
    }
}

/// A connection that said hello and was welcomed.
struct Link {
    socket: WebSocket<TcpStream>,
    host: String,
    route: Route,
    /// Where the host says it listens.
    routes: Vec<Route>,
    /// The host's app version, and whether it can update itself.
    version: String,
    updatable: bool,
}

/// Tells the workspace about a new connection.
fn announce(events: &smol::channel::Sender<ClientEvent>, link: &Link) {
    let _ = events.send_blocking(ClientEvent::Route(link.route.clone()));
    let _ = events.send_blocking(ClientEvent::Message(ServerMessage::Routes { routes: link.routes.clone() }));
    let _ = events.send_blocking(ClientEvent::Message(ServerMessage::Welcome {
        host: link.host.clone(),
        version: link.version.clone(),
        protocol: PROTOCOL,
        routes: link.routes.clone(),
        updatable: link.updatable,
    }));
    let _ = events.send_blocking(ClientEvent::Status(ConnectionStatus::Connected(link.host.clone())));
}

/// Tries each route in turn. It's denied only when every route that
/// answered denied it: an address on another network can be another
/// device's.
fn open_any(routes: &[Route], hello: &Hello) -> Result<Link, Failure> {
    let mut denied = None;
    let mut failures = Vec::new();
    for route in routes {
        match open(route, hello) {
            Ok(link) => return Ok(link),
            Err(Failure::Denied(reason)) => denied = denied.or(Some(reason)),
            Err(Failure::Other(reason)) if routes.len() == 1 => failures.push(reason),
            Err(Failure::Other(reason)) => failures.push(format!("{}: {reason}", route.kind.label())),
        }
    }
    match denied {
        Some(reason) if failures.is_empty() => Err(Failure::Denied(reason)),
        _ if routes.is_empty() => Err(Failure::Other("no address to connect to".into())),
        _ => Err(Failure::Other(failures.join("; "))),
    }
}

/// Looks for the host on this device's tailnet by name when no saved route
/// answers, as when it was only ever reached over the LAN and couldn't
/// report its Tailscale address yet. A route found this way is kept like
/// one the host reported.
fn discover(hello: &Hello) -> Option<Link> {
    if hello.names.is_empty() {
        return None;
    }
    let routes = hello.routes();
    let port = routes.first().and_then(|route| route.address.rsplit_once(':')?.1.parse().ok()).unwrap_or(DEFAULT_PORT);
    let found: Vec<Route> = tailnet_routes(&tailnet_peers(), &hello.names, port)
        .into_iter()
        .filter(|route| !routes.iter().any(|known| known.address == route.address))
        .collect();
    if found.is_empty() {
        return None;
    }
    let mut link = open_any(&found, hello).ok()?;
    // An older host reports no routes; the workspace saves what it hears.
    if !link.routes.iter().any(|route| route.address == link.route.address) {
        link.routes.push(link.route.clone());
    }
    hello.routes.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).push(link.route.clone());
    Some(link)
}

/// Connects, says hello, and waits for the welcome.
fn open(route: &Route, hello: &Hello) -> Result<Link, Failure> {
    let other = |err: &dyn std::fmt::Display| Failure::Other(err.to_string());
    let address = &route.address;
    let target = address.to_socket_addrs().map_err(|err| other(&err))?.next().ok_or_else(|| other(&"no such host"))?;
    let timeout = match route.kind {
        RouteKind::Loopback | RouteKind::Lan => LAN_TIMEOUT,
        _ => CONNECT_TIMEOUT,
    };
    let stream = TcpStream::connect_timeout(&target, timeout).map_err(|err| other(&err))?;
    let _ = stream.set_nodelay(true);
    stream.set_read_timeout(Some(Duration::from_secs(15))).map_err(|err| other(&err))?;
    let (mut socket, _) = tungstenite::client(format!("ws://{address}/"), stream).map_err(|err| other(&err))?;
    let message = ClientMessage::Hello { key: hello.key.clone(), name: hello.name.clone(), protocol: PROTOCOL };
    socket
        .send(Message::text(serde_json::to_string(&message).unwrap_or_default()))
        .map_err(|err| other(&err))?;
    loop {
        match socket.read().map_err(|err| other(&err))? {
            Message::Text(text) => {
                return match serde_json::from_str::<ServerMessage>(text.as_str()) {
                    Ok(ServerMessage::Welcome { host, routes, version, updatable, .. }) => {
                        let _ = socket.get_mut().set_read_timeout(Some(POLL));
                        Ok(Link { socket, host, route: route.clone(), routes, version, updatable })
                    }
                    Ok(ServerMessage::Denied { reason }) => Err(Failure::Denied(reason)),
                    _ => Err(other(&"the host answered something unexpected")),
                };
            }
            Message::Close(_) => return Err(Failure::Denied("The host closed the connection.".into())),
            _ => {}
        }
    }
}

/// Trades messages until the connection drops; returns why it did. Moves
/// to a faster route when one answers.
fn run(
    mut link: Link,
    hello: &Hello,
    inbox: &mpsc::Receiver<String>,
    events: &smol::channel::Sender<ClientEvent>,
    stop: &AtomicBool,
    timing: Timing,
) -> String {
    let (mut pinged, mut heard, mut probed) = (Instant::now(), Instant::now(), Instant::now());
    let mut probe: Option<mpsc::Receiver<Link>> = None;
    let reason = 'connection: {
        while !stop.load(Ordering::Relaxed) {
            while let Ok(text) = inbox.try_recv() {
                if let Err(err) = link.socket.send(Message::text(text)) {
                    break 'connection err.to_string();
                }
            }
            if pinged.elapsed() >= timing.ping {
                pinged = Instant::now();
                if let Err(err) = link.socket.send(Message::Ping(Default::default())) {
                    break 'connection err.to_string();
                }
            }
            if heard.elapsed() > timing.silence {
                break 'connection "The host stopped answering.".into();
            }
            if probe.is_none() && probed.elapsed() >= timing.probe {
                probed = Instant::now();
                let faster: Vec<Route> =
                    hello.routes().into_iter().take_while(|route| route.address != link.route.address).collect();
                if !faster.is_empty() {
                    let (found, receiver) = mpsc::channel();
                    let hello = hello.clone();
                    let _ = thread::Builder::new().name("remote-probe".into()).spawn(move || {
                        if let Some(link) = faster.iter().find_map(|route| open(route, &hello).ok()) {
                            let _ = found.send(link);
                        }
                    });
                    probe = Some(receiver);
                }
            }
            if let Some(receiver) = &probe {
                match receiver.try_recv() {
                    Ok(faster) => {
                        let _ = link.socket.close(None);
                        let _ = link.socket.flush();
                        link = faster;
                        heard = Instant::now();
                        probe = None;
                        announce(events, &link);
                    }
                    Err(mpsc::TryRecvError::Disconnected) => probe = None,
                    Err(mpsc::TryRecvError::Empty) => {}
                }
            }
            match link.socket.read() {
                Ok(message) => {
                    heard = Instant::now();
                    if let Message::Text(text) = message {
                        match serde_json::from_str::<ServerMessage>(text.as_str()) {
                            Ok(message) => {
                                if events.send_blocking(ClientEvent::Message(message)).is_err() {
                                    break 'connection "closed".into();
                                }
                            }
                            Err(err) => eprintln!("vectrum-code: unreadable message from host: {err}"),
                        }
                    } else if let Message::Close(_) = message {
                        break 'connection "The host closed the connection.".into();
                    }
                }
                Err(tungstenite::Error::Io(err)) if matches!(err.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) => {}
                Err(err) => break 'connection err.to_string(),
            }
        }
        "closed".into()
    };
    let _ = link.socket.close(None);
    reason
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::remote::{new_key, server::RemoteServer};

    fn next(client: &RemoteClient) -> Option<ClientEvent> {
        smol::block_on(async {
            smol::future::or(async { client.events.recv().await.ok() }, async {
                smol::Timer::after(Duration::from_secs(10)).await;
                None
            })
            .await
        })
    }

    /// The route the next connection goes over.
    fn connected_route(client: &RemoteClient) -> String {
        let mut route = None;
        loop {
            match next(client) {
                Some(ClientEvent::Route(found)) => route = Some(found.address),
                Some(ClientEvent::Status(ConnectionStatus::Connected(_))) => return route.expect("no route before connecting"),
                Some(ClientEvent::Status(ConnectionStatus::Denied(reason))) => panic!("denied: {reason}"),
                Some(_) => {}
                None => panic!("didn't connect"),
            }
        }
    }

    #[::core::prelude::v1::test]
    fn moves_to_a_faster_route() {
        let port = 48_000 + (std::process::id() % 90) as u16;
        let key = new_key();
        let _server = RemoteServer::start(port, key.clone(), "test-host".into()).unwrap();
        // Nothing listens on the LAN route at first, as away from home.
        let gone = Route { address: format!("127.0.0.1:{}", port + 100), kind: RouteKind::Lan, learned: true };
        let tailnet = Route { address: format!("127.0.0.1:{port}"), kind: RouteKind::Tailnet, learned: true };
        let timing = Timing { ping: Duration::from_millis(200), silence: Duration::from_secs(5), probe: Duration::from_millis(300) };
        let client = RemoteClient::connect_timed(vec![gone, tailnet.clone()], key.clone(), "test-client".into(), Vec::new(), timing);
        assert_eq!(connected_route(&client), tailnet.address);

        // Back home: a faster route answers, and the client moves to it.
        let home = Route { address: format!("127.0.0.2:{port}"), kind: RouteKind::Lan, learned: true };
        client.set_routes(vec![home.clone(), tailnet.clone()]);
        assert_eq!(connected_route(&client), home.address);

        // A route that's another device's (a wrong key) is skipped, not fatal.
        let other_port = port + 200;
        let _other = RemoteServer::start(other_port, new_key(), "someone-else".into()).unwrap();
        let theirs = Route { address: format!("127.0.0.1:{other_port}"), kind: RouteKind::Lan, learned: true };
        let client = RemoteClient::connect_timed(vec![theirs, tailnet.clone()], key, "test-client".into(), Vec::new(), timing);
        assert_eq!(connected_route(&client), tailnet.address);
    }
}