dink-sdk 0.3.1

use async_nats::Client;
use base64::Engine as _;
use bytes::Bytes;
use serde_json::json;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{broadcast, RwLock};
use tracing::{debug, error, info, warn};

use super::reconnect::ConnectionMonitor;

use crate::channel::{Channel, DEFAULT_MAX_PAYLOAD};
use crate::error::{DinkError, Result};
use crate::traits::ServiceHandler;
use crate::types::{EdgeConfig, WebhookRegistration};

// ---------------------------------------------------------------------------
// Key parsing
// ---------------------------------------------------------------------------

/// Returns (server_url, api_key_secret, app_id, edge_id, http_url).
fn parse_edge_key(
    api_key: &str,
    config_url: Option<&str>,
) -> Result<(String, String, String, String, Option<String>)> {
    // Try new format first: base64(JSON{"url":"nats://...", "secret":"dk_..."})
    if let Ok(decoded) = base64::engine::general_purpose::STANDARD.decode(api_key) {
        if let Ok(json_val) = serde_json::from_slice::<serde_json::Value>(&decoded) {
            if let (Some(url), Some(secret)) = (
                json_val.get("url").and_then(|v| v.as_str()),
                json_val.get("secret").and_then(|v| v.as_str()),
            ) {
                let (_, _, app_id, edge_id, _) = parse_legacy_key(secret)?;
                let server_url = config_url.unwrap_or(url).to_string();
                let http_url = json_val
                    .get("httpUrl")
                    .and_then(|v| v.as_str())
                    .map(String::from);
                return Ok((server_url, secret.to_string(), app_id, edge_id, http_url));
            }
        }
    }

    // Legacy: dk_{env}_ed_{appId}_{edgeId}_{random}
    let (_, _, app_id, edge_id, _) = parse_legacy_key(api_key)?;
    let server_url = config_url
        .unwrap_or("nats://connect.dink.cloud:4222")
        .to_string();
    Ok((server_url, api_key.to_string(), app_id, edge_id, None))
}

fn parse_legacy_key(key: &str) -> Result<(String, String, String, String, Option<String>)> {
    let parts: Vec<&str> = key.split('_').collect();
    if parts.len() < 6 || parts[0] != "dk" || parts[2] != "ed" {
        return Err(DinkError::Config(format!(
            "invalid edge key format: {}",
            key
        )));
    }
    // parts: [dk, env, ed, appId, edgeId, random...]
    Ok((
        String::new(), // url placeholder — caller uses config_url
        key.to_string(),
        parts[3].to_string(),
        parts[4].to_string(),
        None, // no http_url in legacy keys
    ))
}

// ---------------------------------------------------------------------------
// Service entry (internal bookkeeping)
// ---------------------------------------------------------------------------

struct ServiceEntry {
    name: String,
    version: String,
    cancel_tx: broadcast::Sender<()>,
}

// ---------------------------------------------------------------------------
// EdgeClient
// ---------------------------------------------------------------------------

pub struct EdgeClient {
    nc: Client,
    app_id: String,
    edge_id: String,
    api_key: String,
    config: EdgeConfig,
    groups: Arc<RwLock<Vec<String>>>,
    services: Arc<RwLock<HashMap<String, ServiceEntry>>>,
    shutdown_tx: Option<broadcast::Sender<()>>,
    monitor: ConnectionMonitor,
}

impl EdgeClient {
    // ------------------------------------------------------------------
    // connect
    // ------------------------------------------------------------------

    pub async fn connect(config: EdgeConfig) -> Result<Self> {
        let (server_url, api_key, app_id, edge_id, _http_url) =
            parse_edge_key(&config.api_key, config.server_url.as_deref())?;

        info!(app_id = %app_id, edge_id = %edge_id, "connecting to dink");

        let monitor = ConnectionMonitor::new();
        let monitor_cb = monitor.clone();

        let mut connect_opts = async_nats::ConnectOptions::with_token(api_key.clone())
            .event_callback(move |event| {
                let mon = monitor_cb.clone();
                async move {
                    match event {
                        async_nats::Event::Disconnected => mon.on_disconnected().await,
                        async_nats::Event::Connected => mon.on_reconnected().await,
                        _ => {}
                    }
                }
            });

        // Wire EdgeConfig reconnect settings
        connect_opts = connect_opts.retry_on_initial_connect();
        connect_opts = connect_opts.connection_timeout(config.timeout);

        let reconnect_wait = config.reconnect_wait;
        connect_opts = connect_opts.reconnect_delay_callback(fixed_reconnect_delay(reconnect_wait));

        let max = map_max_reconnects(config.max_reconnects);
        connect_opts = connect_opts.max_reconnects(max);

        let nc = connect_opts
            .connect(&server_url)
            .await
            .map_err(|e| DinkError::Connection(e.to_string()))?;

        let (shutdown_tx, _) = broadcast::channel::<()>(1);

        let client = Self {
            nc,
            app_id,
            edge_id,
            api_key: api_key.clone(),
            config,
            groups: Arc::new(RwLock::new(Vec::new())),
            services: Arc::new(RwLock::new(HashMap::new())),
            shutdown_tx: Some(shutdown_tx.clone()),
            monitor,
        };

        // Register presence
        client.register_presence().await?;

        // Start heartbeat
        client.spawn_heartbeat(shutdown_tx.subscribe());

        // Subscribe to kick subject
        client.subscribe_kick(shutdown_tx.clone()).await?;

        // Start reconnect handler
        client.spawn_reconnect_handler(shutdown_tx.subscribe());

        info!(edge_id = %client.edge_id, "edge connected");
        Ok(client)
    }

    // ------------------------------------------------------------------
    // expose_service
    // ------------------------------------------------------------------

    pub async fn expose_service(&self, handler: Arc<dyn ServiceHandler>) -> Result<()> {
        let def = handler.definition();
        let name = def.name.to_string();

        // Subscribe to: edge.{appId}.{edgeId}.services.{name}.>
        let subject = format!("edge.{}.{}.services.{}.>", self.app_id, self.edge_id, name);
        debug!(subject = %subject, "subscribing to service");
        let sub = self.nc.subscribe(subject.clone()).await?;

        // Subscribe to stream variant: edge.{appId}.{edgeId}.services.{name}.*.stream
        let stream_subject = format!(
            "edge.{}.{}.services.{}.*.stream",
            self.app_id, self.edge_id, name
        );
        debug!(subject = %stream_subject, "subscribing to stream service");
        let stream_sub = self.nc.subscribe(stream_subject.clone()).await?;

        // Register service via platform
        self.register_service(&name, def.version).await?;

        // Cancel token for this service's tasks
        let (cancel_tx, cancel_rx1) = broadcast::channel::<()>(1);
        let cancel_rx2 = cancel_tx.subscribe();

        // Spawn request handler for unary
        self.spawn_request_handler(name.clone(), handler.clone(), sub, cancel_rx1);

        // Spawn stream handler
        self.spawn_stream_handler(name.clone(), handler.clone(), stream_sub, cancel_rx2);

        // Store entry for later cleanup
        {
            let mut services = self.services.write().await;
            services.insert(
                name.clone(),
                ServiceEntry {
                    name: name.clone(),
                    version: def.version.to_string(),
                    cancel_tx,
                },
            );
        }

        // Auto-register webhooks if handler provides webhook methods
        let webhook_methods = handler.webhook_methods();
        for wm in &webhook_methods {
            if let Err(e) = self
                .register_webhook(&name, &wm.method, "", &wm.secret_mode, &wm.description)
                .await
            {
                warn!(service = %name, method = %wm.method, error = %e, "failed to register webhook");
            }
        }

        info!(service = %name, "service exposed");
        Ok(())
    }

    // ------------------------------------------------------------------
    // expose_channel
    // ------------------------------------------------------------------

    pub async fn expose_channel<F, Fut>(&self, name: &str, version: &str, handler: F) -> Result<()>
    where
        F: Fn(Channel) -> Fut + Send + Sync + 'static,
        Fut: std::future::Future<Output = ()> + Send + 'static,
    {
        let subject = format!(
            "edge.{}.{}.services.{}.*.channel",
            self.app_id, self.edge_id, name
        );

        let mut sub = self.nc.subscribe(subject).await?;
        let nc = self.nc.clone();
        let app_id = self.app_id.clone();
        let edge_id = self.edge_id.clone();
        let handler = Arc::new(handler);

        self.register_service(name, version).await?;

        let (cancel_tx, _) = broadcast::channel::<()>(1);
        let mut cancel_rx = cancel_tx.subscribe();

        {
            let mut services = self.services.write().await;
            services.insert(
                format!("channel:{}", name),
                ServiceEntry {
                    name: name.to_string(),
                    version: version.to_string(),
                    cancel_tx,
                },
            );
        }

        tokio::spawn(async move {
            use futures::StreamExt;
            loop {
                tokio::select! {
                    msg = sub.next() => {
                        let msg = match msg {
                            Some(m) => m,
                            None => break,
                        };

                        let reply = match msg.reply {
                            Some(ref r) => r.clone(),
                            None => continue,
                        };

                        let init: serde_json::Value = match serde_json::from_slice(&msg.payload) {
                            Ok(v) => v,
                            Err(e) => {
                                error!(error = %e, "invalid channel initiation payload");
                                continue;
                            }
                        };

                        let output_subject = match init.get("reply_subject").and_then(|v| v.as_str()) {
                            Some(s) => s.to_string(),
                            None => {
                                error!("channel initiation missing reply_subject");
                                continue;
                            }
                        };

                        let input_subject = match init.get("input_subject").and_then(|v| v.as_str()) {
                            Some(s) => s.to_string(),
                            None => {
                                error!("channel initiation missing input_subject");
                                continue;
                            }
                        };

                        let data = init.get("data").cloned().unwrap_or(serde_json::Value::Null);

                        let cancel_subject = format!(
                            "edge.{}.{}.channel-cancel.{}",
                            app_id, edge_id, uuid::Uuid::new_v4()
                        );

                        let (read_tx, read_rx) = tokio::sync::mpsc::channel::<Bytes>(256);

                        let channel = Channel::new(
                            nc.clone(),
                            output_subject,
                            data,
                            read_rx,
                            DEFAULT_MAX_PAYLOAD,
                        );

                        let mut input_sub = match nc.subscribe(input_subject.clone()).await {
                            Ok(s) => s,
                            Err(e) => {
                                error!(error = %e, "failed to subscribe to channel input");
                                continue;
                            }
                        };

                        let mut cancel_sub = match nc.subscribe(cancel_subject.clone()).await {
                            Ok(s) => s,
                            Err(e) => {
                                error!(error = %e, "failed to subscribe to channel cancel");
                                continue;
                            }
                        };

                        // ACK
                        let ack = serde_json::json!({
                            "status": "channel",
                            "cancelSubject": cancel_subject,
                        });
                        let ack_bytes = serde_json::to_vec(&ack).unwrap_or_default();
                        if let Err(e) = nc.publish(reply, Bytes::from(ack_bytes)).await {
                            error!(error = %e, "failed to send channel ack");
                            continue;
                        }

                        // Spawn input forwarder + cancel listener
                        tokio::spawn(async move {
                            use futures::StreamExt;
                            loop {
                                tokio::select! {
                                    msg = input_sub.next() => {
                                        match msg {
                                            Some(m) => {
                                                if read_tx.send(m.payload).await.is_err() {
                                                    break;
                                                }
                                            }
                                            None => break,
                                        }
                                    }
                                    _ = cancel_sub.next() => {
                                        break;
                                    }
                                }
                            }
                        });

                        let handler = handler.clone();
                        tokio::spawn(async move {
                            handler(channel).await;
                        });
                    }
                    _ = cancel_rx.recv() => {
                        break;
                    }
                }
            }
        });

        Ok(())
    }

    // ------------------------------------------------------------------
    // unexpose_channel
    // ------------------------------------------------------------------

    pub async fn unexpose_channel(&self, name: &str) -> Result<()> {
        let channel_key = format!("channel:{}", name);
        let mut services = self.services.write().await;
        if let Some(entry) = services.remove(&channel_key) {
            let _ = entry.cancel_tx.send(());
        } else {
            return Err(DinkError::Config(format!("channel {} not found", name)));
        }
        info!(service = %name, "channel unexposed");
        Ok(())
    }

    // ------------------------------------------------------------------
    // unexpose_service
    // ------------------------------------------------------------------

    pub async fn unexpose_service(&self, name: &str) -> Result<()> {
        // Deregister from platform
        let service_id = format!("{}.{}", self.edge_id, name);
        let body = json!({ "service_id": service_id });
        self.publish_with_auth("$API.SERVICES.DEREGISTER", serde_json::to_vec(&body)?)
            .await?;

        // Cancel running handlers and remove from map
        let mut services = self.services.write().await;
        if let Some(entry) = services.remove(name) {
            let _ = entry.cancel_tx.send(());
        }

        info!(service = %name, "service unexposed");
        Ok(())
    }

    // ------------------------------------------------------------------
    // call_peer — directed to a specific edge via first group
    // ------------------------------------------------------------------

    pub async fn call_peer(
        &self,
        edge_id: &str,
        service: &str,
        method: &str,
        req: &[u8],
    ) -> Result<Vec<u8>> {
        let groups = self.groups.read().await;
        let group_id = groups
            .first()
            .ok_or_else(|| DinkError::Config("no group joined — cannot call peer".into()))?;
        let subject = format!(
            "edge.{}.group.{}.{}.services.{}.{}",
            self.app_id, group_id, edge_id, service, method
        );
        self.request(&subject, req).await
    }

    // ------------------------------------------------------------------
    // call_group — broadcast to a group
    // ------------------------------------------------------------------

    pub async fn call_group(
        &self,
        group_id: &str,
        service: &str,
        method: &str,
        req: &[u8],
    ) -> Result<Vec<u8>> {
        let subject = format!(
            "edge.{}.group.{}.services.{}.{}",
            self.app_id, group_id, service, method
        );
        self.request(&subject, req).await
    }

    // ------------------------------------------------------------------
    // disconnect
    // ------------------------------------------------------------------

    pub async fn disconnect(&self) -> Result<()> {
        info!("disconnecting edge");

        // Deregister presence first
        self.deregister_presence().await?;

        // Signal heartbeat & handlers to stop
        if let Some(tx) = &self.shutdown_tx {
            let _ = tx.send(());
        }

        // Drain NATS connection
        self.nc
            .drain()
            .await
            .map_err(|e| DinkError::Connection(e.to_string()))?;

        info!("edge disconnected");
        Ok(())
    }

    // ------------------------------------------------------------------
    // accessors
    // ------------------------------------------------------------------

    pub fn app_id(&self) -> &str {
        &self.app_id
    }

    pub fn edge_id(&self) -> &str {
        &self.edge_id
    }

    pub fn connection_monitor(&self) -> &ConnectionMonitor {
        &self.monitor
    }

    // ------------------------------------------------------------------
    // webhook management
    // ------------------------------------------------------------------

    /// Register a webhook endpoint for a service method.
    pub async fn register_webhook(
        &self,
        service: &str,
        method: &str,
        secret: &str,
        secret_mode: &str,
        description: &str,
    ) -> Result<WebhookRegistration> {
        let body = json!({
            "edge_id": self.edge_id,
            "service": service,
            "method": method,
            "secret": secret,
            "secret_mode": secret_mode,
            "description": description,
        });
        let data = self
            .request_with_auth("$API.WEBHOOKS.REGISTER", serde_json::to_vec(&body)?)
            .await?;
        let wh: WebhookRegistration = serde_json::from_slice(&data)?;
        info!(service = %service, method = %method, "webhook registered");
        Ok(wh)
    }

    /// Remove a webhook registration.
    pub async fn deregister_webhook(&self, webhook_id: &str) -> Result<()> {
        let body = json!({ "webhook_id": webhook_id });
        self.request_with_auth("$API.WEBHOOKS.DEREGISTER", serde_json::to_vec(&body)?)
            .await?;
        info!(webhook_id = %webhook_id, "webhook deregistered");
        Ok(())
    }

    /// List webhooks registered for this edge.
    pub async fn list_webhooks(&self) -> Result<Vec<WebhookRegistration>> {
        let body = json!({ "edge_id": self.edge_id });
        let data = self
            .request_with_auth("$API.WEBHOOKS.LIST", serde_json::to_vec(&body)?)
            .await?;
        let webhooks: Vec<WebhookRegistration> = serde_json::from_slice(&data)?;
        Ok(webhooks)
    }

    // ==================================================================
    // Internal helpers
    // ==================================================================

    async fn register_presence(&self) -> Result<()> {
        let body = json!({
            "name": self.config.edge_name.as_deref().unwrap_or(&self.edge_id),
            "type": "lite",
            "labels": self.config.labels,
            "version": self.config.version.as_deref().unwrap_or(""),
            "hostname": "",
            "ip_address": "",
        });

        self.request_with_auth("$EDGE.PRESENCE.REGISTER", serde_json::to_vec(&body)?)
            .await?;

        debug!("presence registered");
        Ok(())
    }

    async fn deregister_presence(&self) -> Result<()> {
        let body = json!({});
        self.request_with_auth("$EDGE.PRESENCE.DEREGISTER", serde_json::to_vec(&body)?)
            .await?;
        debug!("presence deregistered");
        Ok(())
    }

    fn spawn_heartbeat(&self, mut shutdown_rx: broadcast::Receiver<()>) {
        let nc = self.nc.clone();
        let api_key = self.api_key.clone();
        let interval = self.config.heartbeat_interval;

        tokio::spawn(async move {
            let mut tick = tokio::time::interval(interval);
            loop {
                tokio::select! {
                    _ = tick.tick() => {
                        let mut headers = async_nats::HeaderMap::new();
                        headers.insert(
                            "Authorization",
                            format!("Bearer {}", api_key).as_str(),
                        );
                        if let Err(e) = nc.publish_with_headers(
                            "$EDGE.PRESENCE.HEARTBEAT".to_string(),
                            headers,
                            Bytes::from_static(b"{}"),
                        ).await {
                            warn!(error = %e, "heartbeat publish failed");
                        }
                    }
                    _ = shutdown_rx.recv() => {
                        debug!("heartbeat stopped");
                        return;
                    }
                }
            }
        });
    }

    fn spawn_reconnect_handler(&self, mut shutdown_rx: broadcast::Receiver<()>) {
        let monitor = self.monitor.clone();
        let nc = self.nc.clone();
        let api_key = self.api_key.clone();
        let config = self.config.clone();
        let edge_id = self.edge_id.clone();
        let services = self.services.clone();

        tokio::spawn(async move {
            loop {
                tokio::select! {
                    _ = monitor.wait_for_reconnect() => {
                        info!("reconnect detected — re-registering presence and services");

                        // Re-register presence
                        let body = serde_json::json!({
                            "name": config.edge_name.as_deref().unwrap_or(&edge_id),
                            "type": "lite",
                            "labels": config.labels,
                            "version": config.version.as_deref().unwrap_or(""),
                            "hostname": "",
                            "ip_address": "",
                        });
                        let mut headers = async_nats::HeaderMap::new();
                        headers.insert(
                            "Authorization",
                            format!("Bearer {}", api_key).as_str(),
                        );
                        if let Ok(payload) = serde_json::to_vec(&body) {
                            match nc.request_with_headers(
                                "$EDGE.PRESENCE.REGISTER".to_string(),
                                headers.clone(),
                                Bytes::from(payload),
                            ).await {
                                Ok(_) => info!("presence re-registered after reconnect"),
                                Err(e) => error!(error = %e, "failed to re-register presence after reconnect"),
                            }
                        }

                        // Re-register all services
                        let svcs = services.read().await;
                        for (_, entry) in svcs.iter() {
                            let svc_body = serde_json::json!({
                                "name": entry.name,
                                "version": entry.version,
                                "location": "edge",
                                "edge_id": edge_id,
                            });
                            let mut svc_headers = async_nats::HeaderMap::new();
                            svc_headers.insert(
                                "Authorization",
                                format!("Bearer {}", api_key).as_str(),
                            );
                            if let Ok(svc_payload) = serde_json::to_vec(&svc_body) {
                                match nc.request_with_headers(
                                    "$API.SERVICES.REGISTER".to_string(),
                                    svc_headers,
                                    Bytes::from(svc_payload),
                                ).await {
                                    Ok(_) => info!(service = %entry.name, "service re-registered after reconnect"),
                                    Err(e) => error!(service = %entry.name, error = %e, "failed to re-register service after reconnect"),
                                }
                            }
                        }
                    }
                    _ = shutdown_rx.recv() => {
                        debug!("reconnect handler stopped");
                        return;
                    }
                }
            }
        });
    }

    async fn subscribe_kick(&self, shutdown_tx: broadcast::Sender<()>) -> Result<()> {
        let subject = format!("edge.{}.{}.kick", self.app_id, self.edge_id);
        let mut sub = self.nc.subscribe(subject).await?;

        tokio::spawn(async move {
            use futures::StreamExt;
            if let Some(_msg) = sub.next().await {
                warn!("received kick — shutting down");
                let _ = shutdown_tx.send(());
            }
        });
        Ok(())
    }

    async fn register_service(&self, name: &str, version: &str) -> Result<()> {
        let body = json!({
            "name": name,
            "version": version,
            "location": "edge",
            "edge_id": self.edge_id,
        });
        self.request_with_auth("$API.SERVICES.REGISTER", serde_json::to_vec(&body)?)
            .await?;
        debug!(service = %name, "service registered with platform");
        Ok(())
    }

    /// Publish with Authorization header, expect a response (request-reply).
    async fn request_with_auth(&self, subject: &str, payload: Vec<u8>) -> Result<Vec<u8>> {
        let mut headers = async_nats::HeaderMap::new();
        headers.insert("Authorization", format!("Bearer {}", self.api_key).as_str());
        let msg = self
            .nc
            .request_with_headers(subject.to_string(), headers, Bytes::from(payload))
            .await?;

        // Parse platform envelope
        let resp: serde_json::Value = serde_json::from_slice(&msg.payload)?;
        if resp.get("success").and_then(|v| v.as_bool()) == Some(true) {
            let data = resp.get("data").cloned().unwrap_or(serde_json::Value::Null);
            Ok(serde_json::to_vec(&data)?)
        } else {
            let err = resp.get("error").cloned().unwrap_or_default();
            let code = err
                .get("code")
                .and_then(|v| v.as_str())
                .unwrap_or("UNKNOWN");
            let message = err
                .get("message")
                .and_then(|v| v.as_str())
                .unwrap_or("unknown error");
            Err(DinkError::Auth(format!("{}: {}", code, message)))
        }
    }

    /// Publish with Authorization header, fire-and-forget.
    async fn publish_with_auth(&self, subject: &str, payload: Vec<u8>) -> Result<()> {
        let mut headers = async_nats::HeaderMap::new();
        headers.insert("Authorization", format!("Bearer {}", self.api_key).as_str());
        self.nc
            .publish_with_headers(subject.to_string(), headers, Bytes::from(payload))
            .await?;
        Ok(())
    }

    /// Standard NATS request-reply (no auth header).
    async fn request(&self, subject: &str, payload: &[u8]) -> Result<Vec<u8>> {
        let msg = self
            .nc
            .request(subject.to_string(), Bytes::from(payload.to_vec()))
            .await?;

        // Parse service response envelope
        let resp: serde_json::Value = serde_json::from_slice(&msg.payload)?;
        if resp.get("success").and_then(|v| v.as_bool()) == Some(true) {
            let result = resp
                .get("result")
                .cloned()
                .unwrap_or(serde_json::Value::Null);
            Ok(serde_json::to_vec(&result)?)
        } else {
            let err = resp.get("error").cloned().unwrap_or_default();
            let code = err
                .get("code")
                .and_then(|v| v.as_str())
                .unwrap_or("UNKNOWN");
            let message = err
                .get("message")
                .and_then(|v| v.as_str())
                .unwrap_or("unknown error");
            Err(DinkError::Service {
                service: String::new(),
                method: String::new(),
                code: code.to_string(),
                message: message.to_string(),
            })
        }
    }

    // ------------------------------------------------------------------
    // Request handler (unary)
    // ------------------------------------------------------------------

    fn spawn_request_handler(
        &self,
        service_name: String,
        handler: Arc<dyn ServiceHandler>,
        mut sub: async_nats::Subscriber,
        mut cancel_rx: broadcast::Receiver<()>,
    ) {
        let nc = self.nc.clone();
        let svc_prefix = format!(
            "edge.{}.{}.services.{}.",
            self.app_id, self.edge_id, service_name
        );

        tokio::spawn(async move {
            use futures::StreamExt;
            loop {
                tokio::select! {
                    msg = sub.next() => {
                        let msg = match msg {
                            Some(m) => m,
                            None => break,
                        };
                let subject = msg.subject.as_str();

                // Skip stream subjects — handled by stream handler
                if subject.ends_with(".stream") {
                    continue;
                }

                // Extract method: everything after "services.{name}."
                let method = match subject.strip_prefix(&svc_prefix) {
                    Some(m) => m.to_string(),
                    None => {
                        warn!(subject = %subject, "could not extract method from subject");
                        continue;
                    }
                };

                let reply = match msg.reply {
                    Some(ref r) => r.clone(),
                    None => {
                        debug!(subject = %subject, "no reply subject, skipping");
                        continue;
                    }
                };

                let handler = handler.clone();
                let nc = nc.clone();
                let service_name = service_name.clone();

                tokio::spawn(async move {
                    // Unwrap envelope if present: {"metadata": {...}, "payload": {...}}
                    let request_data = match serde_json::from_slice::<serde_json::Value>(&msg.payload) {
                        Ok(envelope) => {
                            if let Some(payload) = envelope.get("payload") {
                                serde_json::to_vec(payload).unwrap_or_else(|_| msg.payload.to_vec())
                            } else {
                                msg.payload.to_vec()
                            }
                        }
                        Err(_) => msg.payload.to_vec(),
                    };
                    let response = match handler.handle_request(&method, &request_data).await {
                        Ok(data) => {
                            // data is already serialized JSON bytes from the handler
                            let result: serde_json::Value =
                                serde_json::from_slice(&data).unwrap_or(serde_json::Value::Null);
                            json!({ "result": result, "success": true })
                        }
                        Err(e) => {
                            error!(
                                service = %service_name,
                                method = %method,
                                error = %e,
                                "handler error"
                            );
                            json!({
                                "success": false,
                                "error": {
                                    "code": "HANDLER_ERROR",
                                    "message": e.to_string(),
                                }
                            })
                        }
                    };

                    let payload = serde_json::to_vec(&response).unwrap_or_default();
                    if let Err(e) = nc.publish(reply, Bytes::from(payload)).await {
                        error!(error = %e, "failed to publish response");
                    }
                });
                    }
                    _ = cancel_rx.recv() => {
                        debug!("request handler cancelled");
                        break;
                    }
                }
            }
        });
    }

    // ------------------------------------------------------------------
    // Stream handler
    // ------------------------------------------------------------------

    fn spawn_stream_handler(
        &self,
        service_name: String,
        handler: Arc<dyn ServiceHandler>,
        mut sub: async_nats::Subscriber,
        mut cancel_rx: broadcast::Receiver<()>,
    ) {
        let nc = self.nc.clone();
        let app_id = self.app_id.clone();
        let edge_id = self.edge_id.clone();

        tokio::spawn(async move {
            use futures::StreamExt;
            loop {
                tokio::select! {
                    msg = sub.next() => {
                        let msg = match msg {
                            Some(m) => m,
                            None => break,
                        };
                        // Extract method from subject
                        let subject = msg.subject.as_str();
                        let parts: Vec<&str> = subject.split('.').collect();
                        let method = {
                            let svc_idx = parts.iter().position(|p| p == &"services");
                            match svc_idx {
                                Some(i) if i + 2 < parts.len() => parts[i + 2].to_string(),
                                _ => {
                                    warn!(subject = %subject, "cannot extract stream method");
                                    continue;
                                }
                            }
                        };

                        let reply = match msg.reply {
                            Some(ref r) => r.clone(),
                            None => {
                                debug!(subject = %subject, "no reply subject for stream");
                                continue;
                            }
                        };

                        let init: serde_json::Value = match serde_json::from_slice(&msg.payload) {
                            Ok(v) => v,
                            Err(e) => {
                                error!(error = %e, "invalid stream initiation payload");
                                continue;
                            }
                        };

                        let reply_subject = match init.get("reply_subject").and_then(|v| v.as_str()) {
                            Some(s) => s.to_string(),
                            None => {
                                error!("stream initiation missing reply_subject");
                                continue;
                            }
                        };

                        let data = init
                            .get("data")
                            .map(|v| serde_json::to_vec(v).unwrap_or_default())
                            .unwrap_or_default();

                        let cancel_subject = format!(
                            "edge.{}.{}.stream-cancel.{}",
                            app_id, edge_id, uuid::Uuid::new_v4()
                        );

                        let ack = json!({
                            "status": "streaming",
                            "cancelSubject": cancel_subject,
                        });
                        let ack_bytes = serde_json::to_vec(&ack).unwrap_or_default();
                        if let Err(e) = nc.publish(reply, Bytes::from(ack_bytes)).await {
                            error!(error = %e, "failed to send stream ack");
                            continue;
                        }

                        let handler = handler.clone();
                        let nc = nc.clone();
                        let service_name = service_name.clone();

                        tokio::spawn(async move {
                            let (emit_tx, mut emit_rx) = tokio::sync::mpsc::unbounded_channel::<Vec<u8>>();
                            let emit_nc = nc.clone();
                            let emit_subj = reply_subject.clone();

                            // Background task: publish chunks in order, guaranteeing delivery
                            let flush_handle = tokio::spawn(async move {
                                while let Some(payload) = emit_rx.recv().await {
                                    if let Err(e) = emit_nc.publish(emit_subj.clone(), Bytes::from(payload)).await {
                                        error!(error = %e, "stream emit publish failed");
                                        break;
                                    }
                                }
                            });

                            let emit = Box::new(move |payload: Vec<u8>| -> crate::error::Result<()> {
                                emit_tx.send(payload).map_err(|e| crate::error::DinkError::Nats(e.to_string()))
                            });

                            match handler.handle_stream(&method, &data, emit).await {
                                Ok(()) => {
                                    // Wait for all chunks to be flushed before sending done
                                    let _ = flush_handle.await;
                                    let done_subject = format!("{}.done", reply_subject);
                                    let _ = nc.publish(done_subject, Bytes::from_static(b"{}")).await;
                                }
                                Err(e) => {
                                    // Wait for pending chunks even on error
                                    let _ = flush_handle.await;
                                    error!(service = %service_name, method = %method, error = %e, "stream handler error");
                                    let err_subject = format!("{}.error", reply_subject);
                                    let err_payload = json!({
                                        "success": false,
                                        "error": { "code": "STREAM_ERROR", "message": e.to_string() }
                                    });
                                    let _ = nc.publish(
                                        err_subject,
                                        Bytes::from(serde_json::to_vec(&err_payload).unwrap_or_default()),
                                    ).await;
                                }
                            }
                        });
                    }
                    _ = cancel_rx.recv() => {
                        debug!("stream handler cancelled");
                        break;
                    }
                }
            }
        });
    }
}

/// Maps `EdgeConfig.max_reconnects` (i32) to the `Option<usize>` expected by async_nats.
/// - negative (e.g. -1) → `None` (unlimited reconnects)
/// - 0 → `Some(0)` (no reconnects)
/// - positive n → `Some(n)`
pub(crate) fn map_max_reconnects(value: i32) -> Option<usize> {
    if value < 0 {
        None
    } else {
        Some(value as usize)
    }
}

/// Returns a reconnect delay callback that always returns the given fixed duration.
pub(crate) fn fixed_reconnect_delay(
    wait: std::time::Duration,
) -> impl Fn(usize) -> std::time::Duration {
    move |_attempt| wait
}

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

    // ---- map_max_reconnects ----

    #[test]
    fn max_reconnects_negative_means_unlimited() {
        assert_eq!(map_max_reconnects(-1), None);
        assert_eq!(map_max_reconnects(-100), None);
        assert_eq!(map_max_reconnects(i32::MIN), None);
    }

    #[test]
    fn max_reconnects_zero_means_no_reconnects() {
        assert_eq!(map_max_reconnects(0), Some(0));
    }

    #[test]
    fn max_reconnects_positive_maps_directly() {
        assert_eq!(map_max_reconnects(1), Some(1));
        assert_eq!(map_max_reconnects(5), Some(5));
        assert_eq!(map_max_reconnects(1000), Some(1000));
    }

    // ---- fixed_reconnect_delay ----

    #[test]
    fn reconnect_delay_returns_configured_wait() {
        let wait = Duration::from_millis(500);
        let cb = fixed_reconnect_delay(wait);
        assert_eq!(cb(0), wait);
        assert_eq!(cb(1), wait);
        assert_eq!(cb(99), wait);
    }

    #[test]
    fn reconnect_delay_ignores_attempt_number() {
        let wait = Duration::from_secs(3);
        let cb = fixed_reconnect_delay(wait);
        // Every attempt returns the same duration — no backoff.
        for attempt in 0..20 {
            assert_eq!(cb(attempt), wait, "attempt {attempt} differed");
        }
    }

    // ---- EdgeConfig defaults wire correctly ----

    #[test]
    fn default_config_maps_to_unlimited_reconnects() {
        let cfg = EdgeConfig::default();
        assert_eq!(map_max_reconnects(cfg.max_reconnects), None);
    }

    #[test]
    fn default_config_reconnect_wait_is_2s() {
        let cfg = EdgeConfig::default();
        let cb = fixed_reconnect_delay(cfg.reconnect_wait);
        assert_eq!(cb(0), Duration::from_secs(2));
    }
}