//! WebSocket tunnel client with auto-reconnect.
//!
//! Manages a persistent WebSocket connection to a tunnel server,
//! bridging RPC messages between a TunnelEndpoint and the remote side.
//! Automatically reconnects with exponential backoff on disconnection.

use crate::state::LivenessToken;
use crate::tunnel::{create_endpoint, Dispatch, TunnelCaller, TunnelEndpoint};
use serde_json::Value;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tokio::sync::{mpsc, watch};
use tokio_tungstenite::tungstenite::http;
use tracing::{error, info, warn};

/// Configuration for tunnel client connection.
pub struct TunnelConfig {
    /// WebSocket URL to connect to (e.g., "ws://host:port/tunnel")
    pub url: String,
    /// Optional HTTP headers for the WebSocket handshake
    pub headers: Vec<(String, String)>,
    /// RPC call timeout
    pub rpc_timeout: Duration,
    /// Reconnect policy
    pub reconnect: ReconnectPolicy,
    /// WebSocket ping interval for keepalive
    pub ping_interval: Duration,
}

/// Reconnect policy for the tunnel client.
pub struct ReconnectPolicy {
    /// Initial delay before first reconnect attempt
    pub initial_delay: Duration,
    /// Maximum delay between reconnect attempts
    pub max_delay: Duration,
    /// Multiplier for exponential backoff
    pub backoff_multiplier: f64,
}

impl Default for ReconnectPolicy {
    fn default() -> Self {
        Self {
            initial_delay: Duration::from_secs(2),
            max_delay: Duration::from_secs(60),
            backoff_multiplier: 2.0,
        }
    }
}

/// Managed WebSocket tunnel client.
///
/// Maintains a persistent connection, auto-reconnects on failure,
/// and bridges RPC messages through a TunnelEndpoint.
pub struct TunnelClient {
    endpoint: Arc<TunnelEndpoint>,
    shutdown_tx: watch::Sender<bool>,
    /// Tracks whether the tunnel connection is alive.
    /// Set to `false` when the relay task exits (disconnect or WebSocket failure).
    connected: Arc<AtomicBool>,
}

impl TunnelClient {
    /// Connect to a WebSocket tunnel server.
    ///
    /// - `config`: connection configuration
    /// - `dispatchers`: local RPC handlers for incoming requests
    /// - `connected_tx`: notified each time a connection is established (including reconnects)
    pub async fn connect(
        config: TunnelConfig,
        dispatchers: Vec<Box<dyn Dispatch>>,
        connected_tx: mpsc::UnboundedSender<()>,
    ) -> Result<Arc<Self>, String> {
        let (shutdown_tx, shutdown_rx) = watch::channel(false);

        let (endpoint, ws_to_ep_tx, ep_to_ws_rx) =
            create_endpoint(dispatchers, config.rpc_timeout);

        let connected = Arc::new(AtomicBool::new(true));

        tokio::spawn(relay_task(
            config,
            ws_to_ep_tx,
            ep_to_ws_rx,
            connected_tx,
            shutdown_rx,
            connected.clone(),
        ));

        Ok(Arc::new(Self {
            endpoint,
            shutdown_tx,
            connected,
        }))
    }

    /// Connect to an HTTP server's WebSocket tunnel endpoint.
    ///
    /// Converts `http(s)://` URL to `ws(s)://` and appends the tunnel path.
    pub async fn connect_to_http(
        http_url: &str,
        ws_path: &str,
        dispatchers: Vec<Box<dyn Dispatch>>,
        headers: Vec<(String, String)>,
        connected_tx: mpsc::UnboundedSender<()>,
    ) -> Result<Arc<Self>, String> {
        let ws_url = http_to_ws_url(http_url, ws_path)?;
        let config = TunnelConfig {
            url: ws_url,
            headers,
            rpc_timeout: Duration::from_secs(30),
            reconnect: ReconnectPolicy::default(),
            ping_interval: Duration::from_secs(30),
        };
        Self::connect(config, dispatchers, connected_tx).await
    }

    /// Disconnect from the tunnel server.
    ///
    /// Signals the relay task to stop. The WebSocket connection will be closed.
    pub fn disconnect(&self) {
        let _ = self.shutdown_tx.send(true);
    }

    /// Returns a [`LivenessToken`] for this connection.
    ///
    /// The token becomes dead (`is_alive() == false`) when the relay task exits,
    /// which happens on `disconnect()` or when the WebSocket connection is lost.
    pub fn liveness(&self) -> LivenessToken {
        LivenessToken::new(self.connected.clone())
    }
}

#[async_trait::async_trait]
impl TunnelCaller for TunnelClient {
    async fn call(&self, method: &str, params: Value) -> Result<Value, crate::tunnel::TunnelError> {
        self.endpoint.call(method, params).await
    }
}

// ── Relay task (manages reconnection loop) ────────────────────

async fn relay_task(
    config: TunnelConfig,
    ws_to_ep_tx: mpsc::UnboundedSender<String>,
    ep_to_ws_rx: mpsc::UnboundedReceiver<String>,
    connected_tx: mpsc::UnboundedSender<()>,
    shutdown_rx: watch::Receiver<bool>,
    liveness_flag: Arc<AtomicBool>,
) {
    let mut ep_to_ws_rx = ep_to_ws_rx;
    let mut delay = config.reconnect.initial_delay;

    loop {
        // Check shutdown before connecting
        if *shutdown_rx.borrow() {
            info!("[TUNNEL] Shutdown requested, stopping relay task");
            break;
        }

        info!("[TUNNEL] Connecting to {}", config.url);

        match try_connect(&config).await {
            Ok(ws_stream) => {
                info!("[TUNNEL] Connected to {}", config.url);
                delay = config.reconnect.initial_delay; // Reset backoff

                // Notify: connection established
                let _ = connected_tx.send(());

                run_connection(
                    ws_stream,
                    &ws_to_ep_tx,
                    &mut ep_to_ws_rx,
                    &config,
                    shutdown_rx.clone(),
                )
                .await;

                if *shutdown_rx.borrow() {
                    info!("[TUNNEL] Shutdown after connection closed");
                    break;
                }

                warn!("[TUNNEL] Connection lost, will reconnect in {:?}", delay);
            }
            Err(e) => {
                error!("[TUNNEL] Connection failed: {}, retrying in {:?}", e, delay);
            }
        }

        // Wait before reconnecting (interruptible by shutdown)
        let mut shutdown_rx_clone = shutdown_rx.clone();
        tokio::select! {
            _ = tokio::time::sleep(delay) => {}
            _ = shutdown_rx_clone.changed() => {
                if *shutdown_rx_clone.borrow() {
                    info!("[TUNNEL] Shutdown during reconnect wait");
                    break;
                }
            }
        }

        // Exponential backoff
        delay = Duration::from_secs_f64(
            (delay.as_secs_f64() * config.reconnect.backoff_multiplier)
                .min(config.reconnect.max_delay.as_secs_f64()),
        );
    }

    // Mark connection as dead when relay task exits
    liveness_flag.store(false, Ordering::Release);
}

async fn try_connect(
    config: &TunnelConfig,
) -> Result<
    tokio_tungstenite::WebSocketStream<
        tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
    >,
    String,
> {
    use tokio_tungstenite::tungstenite::client::IntoClientRequest;

    let uri: http::Uri = config
        .url
        .parse()
        .map_err(|e| format!("invalid URL: {}", e))?;

    let mut request = uri
        .into_client_request()
        .map_err(|e| format!("request build error: {}", e))?;

    for (key, value) in &config.headers {
        let header_name = http::header::HeaderName::from_bytes(key.as_bytes())
            .map_err(|e| format!("invalid header name '{}': {}", key, e))?;
        let header_value = http::header::HeaderValue::from_str(value)
            .map_err(|e| format!("invalid header value for '{}': {}", key, e))?;
        request.headers_mut().insert(header_name, header_value);
    }

    let (ws_stream, _response) = tokio_tungstenite::connect_async(request)
        .await
        .map_err(|e| format!("WebSocket connect error: {}", e))?;

    Ok(ws_stream)
}

async fn run_connection(
    ws_stream: tokio_tungstenite::WebSocketStream<
        tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
    >,
    ws_to_ep_tx: &mpsc::UnboundedSender<String>,
    ep_to_ws_rx: &mut mpsc::UnboundedReceiver<String>,
    config: &TunnelConfig,
    mut shutdown_rx: watch::Receiver<bool>,
) {
    use futures_util::{SinkExt, StreamExt};
    use tokio_tungstenite::tungstenite::Message;

    let (mut writer, mut reader) = ws_stream.split();

    // Spawn reader task: forward all text messages to endpoint
    let ws_to_ep_tx_clone = ws_to_ep_tx.clone();
    let mut reader_handle = tokio::spawn(async move {
        while let Some(msg) = reader.next().await {
            match msg {
                Ok(Message::Text(text)) => {
                    let _ = ws_to_ep_tx_clone.send(text.to_string());
                }
                Ok(Message::Close(_)) => break,
                Err(e) => {
                    warn!("[TUNNEL] WebSocket read error: {}", e);
                    break;
                }
                _ => {} // Ignore Ping/Pong/Binary
            }
        }
    });

    let ping_interval = config.ping_interval;
    let mut heartbeat = tokio::time::interval(ping_interval);
    heartbeat.tick().await; // Skip first immediate tick

    loop {
        tokio::select! {
            // Outgoing RPC messages from endpoint → WebSocket
            msg = ep_to_ws_rx.recv() => {
                match msg {
                    Some(text) => {
                        if writer.send(Message::Text(text.into())).await.is_err() {
                            break;
                        }
                    }
                    None => break, // Endpoint closed
                }
            }

            // Periodic WebSocket ping for keepalive
            _ = heartbeat.tick() => {
                if writer.send(Message::Ping(vec![].into())).await.is_err() {
                    break;
                }
            }

            // Shutdown signal
            _ = shutdown_rx.changed() => {
                if *shutdown_rx.borrow() {
                    let _ = writer.send(Message::Close(None)).await;
                    break;
                }
            }

            // Reader task finished (connection lost)
            _ = &mut reader_handle => {
                break;
            }
        }
    }
}

// ── URL conversion ─────────────────────────────────────────────

/// Convert HTTP URL to WebSocket URL, appending the tunnel path.
///
/// `http://` → `ws://`, `https://` → `wss://`
#[allow(clippy::string_slice)]
pub fn http_to_ws_url(http_url: &str, ws_path: &str) -> Result<String, String> {
    let trimmed = http_url.trim_end_matches('/');
    let ws_base = if trimmed.starts_with("https://") {
        format!("wss://{}", &trimmed["https://".len()..])
    } else if trimmed.starts_with("http://") {
        format!("ws://{}", &trimmed["http://".len()..])
    } else {
        return Err(format!("unsupported URL scheme: {}", http_url));
    };
    let path = ws_path.trim_start_matches('/');
    Ok(format!("{}/{}", ws_base, path))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    #[allow(clippy::string_slice)]
    fn test_http_to_ws_url() {
        assert_eq!(
            http_to_ws_url("http://localhost:8080", "tunnel").unwrap(),
            "ws://localhost:8080/tunnel"
        );
        assert_eq!(
            http_to_ws_url("https://example.com", "/ws/tunnel").unwrap(),
            "wss://example.com/ws/tunnel"
        );
        assert_eq!(
            http_to_ws_url("http://host:3000/", "tunnel").unwrap(),
            "ws://host:3000/tunnel"
        );
        assert!(http_to_ws_url("ftp://host", "tunnel").is_err());
    }
}