Kordophone/kordophoned/src/xpc/agent.rs

use crate::xpc::interface::SERVICE_NAME;
use kordophoned::daemon::settings::Settings;
use kordophoned::daemon::{events::Event, signals::Signal, DaemonResult};
use std::collections::HashMap;
use std::ffi::CString;
use std::os::raw::c_char;
use std::ptr;
use std::sync::Arc;
use std::thread;
use tokio::sync::{mpsc, oneshot, Mutex};
use xpc_connection::{message_to_xpc_object, xpc_object_to_message, Message, MessageError};
use xpc_connection_sys as xpc_sys;

static LOG_TARGET: &str = "xpc";

/// Wrapper for raw XPC connection pointer to declare cross-thread usage.
/// Safety: libxpc connections are reference-counted and may be used to send from other threads.
#[derive(Copy, Clone)]
struct XpcConn(pub xpc_sys::xpc_connection_t);
unsafe impl Send for XpcConn {}
unsafe impl Sync for XpcConn {}

/// XPC IPC agent that forwards daemon events and signals over libxpc.
#[derive(Clone)]
pub struct XpcAgent {
    event_sink: mpsc::Sender<Event>,
    signal_receiver: Arc<Mutex<Option<mpsc::Receiver<Signal>>>>,
}

impl XpcAgent {
    /// Create a new XPC agent with an event sink and signal receiver.
    pub fn new(event_sink: mpsc::Sender<Event>, signal_receiver: mpsc::Receiver<Signal>) -> Self {
        Self { event_sink, signal_receiver: Arc::new(Mutex::new(Some(signal_receiver))) }
    }

    /// Run the XPC agent and host the XPC service. Implements generic dispatch.
    pub async fn run(self) {
        use block::ConcreteBlock;
        use std::ops::Deref;
        use std::sync::Mutex as StdMutex;

        log::info!(target: LOG_TARGET, "XPCAgent running");

        // Construct the Mach service name without a trailing NUL for CString.
        let service_name = SERVICE_NAME.trim_end_matches('\0');
        let mach_port_name = match CString::new(service_name) {
            Ok(c) => c,
            Err(e) => {
                log::error!(target: LOG_TARGET, "Invalid XPC service name: {e}");
                return;
            }
        };

        log::info!(
            target: LOG_TARGET,
            "Waiting for XPC connections on {}",
            service_name
        );

        // Multi-thread runtime to drive async dispatch from XPC event handlers.
        let rt = match tokio::runtime::Runtime::new() {
            Ok(rt) => Arc::new(rt),
            Err(e) => {
                log::error!(target: LOG_TARGET, "Failed to create Tokio runtime: {}", e);
                return;
            }
        };

        // Shared list of connected clients for signal fanout
        let connections: Arc<StdMutex<Vec<XpcConn>>> = Arc::new(StdMutex::new(Vec::new()));
        // Forward daemon signals to all connected clients
        {
            let receiver_arc = self.signal_receiver.clone();
            let conns = connections.clone();
            rt.spawn(async move {
                let mut receiver = receiver_arc
                    .lock()
                    .await
                    .take()
                    .expect("Signal receiver already taken");
                while let Some(signal) = receiver.recv().await {
                    log::info!(target: LOG_TARGET, "Broadcasting signal: {:?}", signal);
                    let msg = signal_to_message(signal);
                    let xobj = unsafe { message_to_xpc_object(msg) };
                    let list = conns.lock().unwrap();
                    log::info!(target: LOG_TARGET, "Active XPC clients: {}", list.len());
                    for c in list.iter() {
                        log::info!(target: LOG_TARGET, "Sending signal to client");
                        unsafe { xpc_sys::xpc_connection_send_message(c.0, xobj) };
                    }
                    unsafe { xpc_sys::xpc_release(xobj) };
                }
            });
        }

        // Create the XPC Mach service listener.
        let service = unsafe {
            xpc_sys::xpc_connection_create_mach_service(
                mach_port_name.as_ptr(),
                ptr::null_mut(),
                xpc_sys::XPC_CONNECTION_MACH_SERVICE_LISTENER as u64,
            )
        };

        // Event handler for the service: accepts new client connections.
        let agent = self.clone();
        let rt_accept = rt.clone();
        let conns_accept = connections.clone();
        let service_handler = ConcreteBlock::new(move |event: xpc_sys::xpc_object_t| {
            unsafe {
                // Treat incoming events as connections; ignore others
                // We detect connections by trying to set a per-connection handler.
                    let client = event as xpc_sys::xpc_connection_t;
                    log::info!(target: LOG_TARGET, "New XPC connection accepted");
                    // Do not register for signals until the client explicitly subscribes

                    // Per-connection handler
                    let agent_conn = agent.clone();
                    let rt_conn = rt_accept.clone();
                    let conns_for_handler = conns_accept.clone();
                    let conn_handler = ConcreteBlock::new(move |msg: xpc_sys::xpc_object_t| {
                        unsafe {
                            // Convert to higher-level Message for type matching
                            match xpc_object_to_message(msg) {
                                Message::Dictionary(map) => {
                                    // Trace inbound method
                                    let method = dict_get_str(&map, "method").or_else(|| dict_get_str(&map, "type")).unwrap_or_else(|| "<unknown>".to_string());
                                    log::info!(target: LOG_TARGET, "XPC request received: {}", method);
                                    let response = rt_conn.block_on(dispatch(&agent_conn, &conns_for_handler, client, &map));
                                    let reply = xpc_sys::xpc_dictionary_create_reply(msg);
                                    if !reply.is_null() {
                                        let payload = message_to_xpc_object(response);
                                        let apply_block = ConcreteBlock::new(move |key: *const c_char, value: xpc_sys::xpc_object_t| {
                                            xpc_sys::xpc_dictionary_set_value(reply, key, value);
                                        })
                                        .copy();
                                        xpc_sys::xpc_dictionary_apply(payload, apply_block.deref() as *const _ as *mut _);
                                        xpc_sys::xpc_connection_send_message(client, reply);
                                        xpc_sys::xpc_release(payload);
                                        xpc_sys::xpc_release(reply);
                                        log::info!(target: LOG_TARGET, "XPC reply sent for method: {}", method);
                                    } else {
                                        log::warn!(target: LOG_TARGET, "No reply port for method: {}", method);
                                    }
                                }
                                Message::Error(e) => {
                                    match e {
                                        MessageError::ConnectionInvalid => {
                                            // Normal for one-shot RPC connections; keep logs quiet
                                            let mut list = conns_for_handler.lock().unwrap();
                                            let before = list.len();
                                            list.retain(|c| c.0 != client);
                                            let after = list.len();
                                            if after < before {
                                                log::info!(target: LOG_TARGET, "Removed closed XPC client from subscribers ({} -> {})", before, after);
                                            } else {
                                                log::debug!(target: LOG_TARGET, "XPC connection closed (no subscription)");
                                            }
                                        }
                                        other => {
                                            log::warn!(target: LOG_TARGET, "XPC error event: {:?}", other);
                                        }
                                    }
                                }
                                _ => {}
                            }
                        }
                    })
                    .copy();

                    xpc_sys::xpc_connection_set_event_handler(
                        client,
                        conn_handler.deref() as *const _ as *mut _,
                    );
                    xpc_sys::xpc_connection_resume(client);
                }
            }
        )
        .copy();

        unsafe {
            xpc_sys::xpc_connection_set_event_handler(
                service,
                service_handler.deref() as *const _ as *mut _,
            );
            xpc_sys::xpc_connection_resume(service);
        }

        // Keep this future alive forever.
        futures_util::future::pending::<()>().await;
    }

    /// Send an event to the daemon and await its reply.
    pub async fn send_event<T>(
        &self,
        make_event: impl FnOnce(kordophoned::daemon::events::Reply<T>) -> Event,
    ) -> DaemonResult<T> {
        let (tx, rx) = oneshot::channel();
        self.event_sink
            .send(make_event(tx))
            .await
            .map_err(|_| "Failed to send event")?;
        rx.await.map_err(|_| "Failed to receive reply".into())
    }
}

fn cstr(s: &str) -> CString {
    CString::new(s).unwrap_or_else(|_| CString::new("").unwrap())
}

fn get_string_field(map: &HashMap<CString, Message>, key: &str) -> Option<String> {
    let k = CString::new(key).ok()?;
    map.get(&k).and_then(|v| match v {
        Message::String(s) => Some(s.to_string_lossy().into_owned()),
        _ => None,
    })
}

fn get_dictionary_field<'a>(
    map: &'a HashMap<CString, Message>,
    key: &str,
) -> Option<&'a HashMap<CString, Message>> {
    let k = CString::new(key).ok()?;
    map.get(&k).and_then(|v| match v {
        Message::Dictionary(d) => Some(d),
        _ => None,
    })
}

fn make_error_reply(code: &str, message: &str) -> Message {
    log::error!(target: LOG_TARGET, "XPC error: {code}: {message}");
    
    let mut reply: HashMap<CString, Message> = HashMap::new();
    reply.insert(cstr("type"), Message::String(cstr("Error")));
    reply.insert(cstr("error"), Message::String(cstr(code)));
    reply.insert(cstr("message"), Message::String(cstr(message)));

    Message::Dictionary(reply)
}

type XpcMap = HashMap<CString, Message>;

fn dict_get_str(map: &XpcMap, key: &str) -> Option<String> {
    let k = CString::new(key).ok()?;
    match map.get(&k) {
        Some(Message::String(v)) => Some(v.to_string_lossy().into_owned()),
        _ => None,
    }
}

fn dict_get_i64_from_str(map: &XpcMap, key: &str) -> Option<i64> {
    dict_get_str(map, key).and_then(|s| s.parse::<i64>().ok())
}

fn dict_put_str(map: &mut XpcMap, key: &str, value: impl AsRef<str>) {
    map.insert(cstr(key), Message::String(cstr(value.as_ref())));
}

fn dict_put_i64_as_str(map: &mut XpcMap, key: &str, value: i64) {
    dict_put_str(map, key, value.to_string());
}

fn array_from_strs(values: impl IntoIterator<Item = String>) -> Message {
    let arr = values
        .into_iter()
        .map(|s| Message::String(cstr(&s)))
        .collect();
    Message::Array(arr)
}

fn make_ok_reply() -> Message {
    let mut reply: XpcMap = HashMap::new();
    dict_put_str(&mut reply, "type", "Ok");
    Message::Dictionary(reply)
}

/// Attach an optional request_id to a dictionary reply message.
fn attach_request_id(mut message: Message, request_id: Option<String>) -> Message {
    if let (Some(id), Message::Dictionary(ref mut m)) = (request_id, &mut message) {
        dict_put_str(m, "request_id", &id);
    }
    message
}

async fn dispatch(
    agent: &XpcAgent,
    subscribers: &std::sync::Mutex<Vec<XpcConn>>, 
    current_client: xpc_sys::xpc_connection_t,
    root: &HashMap<CString, Message>,
) -> Message {
    // Standardized request: { method: String, arguments: Dictionary?, request_id: String? }
    let request_id = dict_get_str(root, "request_id");

    let method = match dict_get_str(root, "method").or_else(|| dict_get_str(root, "type")) {
        Some(m) => m,
        None => return attach_request_id(make_error_reply("InvalidRequest", "Missing method/type"), request_id),
    };

    let _arguments = get_dictionary_field(root, "arguments");

    let mut response = match method.as_str() {
        // Example implemented method: GetVersion
        "GetVersion" => match agent.send_event(Event::GetVersion).await {
            Ok(version) => {
                let mut reply: XpcMap = HashMap::new();
                dict_put_str(&mut reply, "type", "GetVersionResponse");
                dict_put_str(&mut reply, "version", &version);
                Message::Dictionary(reply)
            }
            Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
        },

        "GetConversations" => {
            // Defaults
            let mut limit: i32 = 100;
            let mut offset: i32 = 0;

            if let Some(args) = get_dictionary_field(root, "arguments") {
                if let Some(v) = dict_get_i64_from_str(args, "limit") {
                    limit = v as i32;
                }
                if let Some(v) = dict_get_i64_from_str(args, "offset") {
                    offset = v as i32;
                }
            }

            match agent
                .send_event(|r| Event::GetAllConversations(limit, offset, r))
                .await
            {
                Ok(conversations) => {
                    // Build array of conversation dictionaries
                    let mut items: Vec<Message> = Vec::with_capacity(conversations.len());
                    for conv in conversations {
                        let mut m: XpcMap = HashMap::new();
                        dict_put_str(&mut m, "guid", &conv.guid);
                        dict_put_str(
                            &mut m,
                            "display_name",
                            &conv.display_name.unwrap_or_default(),
                        );
                        dict_put_i64_as_str(&mut m, "unread_count", conv.unread_count as i64);
                        dict_put_str(
                            &mut m,
                            "last_message_preview",
                            &conv.last_message_preview.unwrap_or_default(),
                        );

                        // participants -> array of strings
                        let participant_names: Vec<String> = conv
                            .participants
                            .into_iter()
                            .map(|p| p.display_name())
                            .collect();
                        m.insert(cstr("participants"), array_from_strs(participant_names));

                        // date as unix timestamp (i64)
                        dict_put_i64_as_str(&mut m, "date", conv.date.and_utc().timestamp());

                        items.push(Message::Dictionary(m));
                    }

                    let mut reply: XpcMap = HashMap::new();
                    dict_put_str(&mut reply, "type", "GetConversationsResponse");
                    reply.insert(cstr("conversations"), Message::Array(items));
                    Message::Dictionary(reply)
                }
                Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
            }
        }

        "SyncConversationList" => match agent.send_event(Event::SyncConversationList).await {
            Ok(()) => make_ok_reply(),
            Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
        },

        "SyncAllConversations" => match agent.send_event(Event::SyncAllConversations).await {
            Ok(()) => make_ok_reply(),
            Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
        },

        "SyncConversation" => {
            let conversation_id = match get_dictionary_field(root, "arguments")
                .and_then(|m| dict_get_str(m, "conversation_id"))
            {
                Some(id) => id,
                None => return make_error_reply("InvalidRequest", "Missing conversation_id"),
            };
            match agent
                .send_event(|r| Event::SyncConversation(conversation_id, r))
                .await
            {
                Ok(()) => make_ok_reply(),
                Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
            }
        }

        "MarkConversationAsRead" => {
            let conversation_id = match get_dictionary_field(root, "arguments")
                .and_then(|m| dict_get_str(m, "conversation_id"))
            {
                Some(id) => id,
                None => return make_error_reply("InvalidRequest", "Missing conversation_id"),
            };
            match agent
                .send_event(|r| Event::MarkConversationAsRead(conversation_id, r))
                .await
            {
                Ok(()) => make_ok_reply(),
                Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
            }
        }

        "GetMessages" => {
            let args = match get_dictionary_field(root, "arguments") {
                Some(a) => a,
                None => return make_error_reply("InvalidRequest", "Missing arguments"),
            };
            let conversation_id = match dict_get_str(args, "conversation_id") {
                Some(id) => id,
                None => return make_error_reply("InvalidRequest", "Missing conversation_id"),
            };
            let last_message_id = dict_get_str(args, "last_message_id");
            match agent
                .send_event(|r| Event::GetMessages(conversation_id, last_message_id, r))
                .await
            {
                Ok(messages) => {
                    let mut items: Vec<Message> = Vec::with_capacity(messages.len());
                    for msg in messages {
                        let mut m: XpcMap = HashMap::new();
                        dict_put_str(&mut m, "id", &msg.id);
                        dict_put_str(&mut m, "text", &msg.text.replace('\u{FFFC}', ""));
                        dict_put_i64_as_str(&mut m, "date", msg.date.and_utc().timestamp());
                        dict_put_str(&mut m, "sender", &msg.sender.display_name());
                        items.push(Message::Dictionary(m));
                    }
                    let mut reply: XpcMap = HashMap::new();
                    dict_put_str(&mut reply, "type", "GetMessagesResponse");
                    reply.insert(cstr("messages"), Message::Array(items));
                    Message::Dictionary(reply)
                }
                Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
            }
        }

        "DeleteAllConversations" => match agent.send_event(Event::DeleteAllConversations).await {
            Ok(()) => make_ok_reply(),
            Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
        },

        "SendMessage" => {
            let args = match get_dictionary_field(root, "arguments") {
                Some(a) => a,
                None => return make_error_reply("InvalidRequest", "Missing arguments"),
            };
            let conversation_id = match dict_get_str(args, "conversation_id") {
                Some(v) => v,
                None => return make_error_reply("InvalidRequest", "Missing conversation_id"),
            };
            let text = dict_get_str(args, "text").unwrap_or_default();
            let attachment_guids: Vec<String> = match args.get(&cstr("attachment_guids")) {
                Some(Message::Array(arr)) => arr
                    .iter()
                    .filter_map(|m| match m {
                        Message::String(s) => Some(s.to_string_lossy().into_owned()),
                        _ => None,
                    })
                    .collect(),
                _ => Vec::new(),
            };
            match agent
                .send_event(|r| Event::SendMessage(conversation_id, text, attachment_guids, r))
                .await
            {
                Ok(uuid) => {
                    let mut reply: XpcMap = HashMap::new();
                    dict_put_str(&mut reply, "type", "SendMessageResponse");
                    dict_put_str(&mut reply, "uuid", &uuid.to_string());
                    Message::Dictionary(reply)
                }
                Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
            }
        }

        "GetAttachmentInfo" => {
            let args = match get_dictionary_field(root, "arguments") {
                Some(a) => a,
                None => return make_error_reply("InvalidRequest", "Missing arguments"),
            };
            let attachment_id = match dict_get_str(args, "attachment_id") {
                Some(v) => v,
                None => return make_error_reply("InvalidRequest", "Missing attachment_id"),
            };
            match agent
                .send_event(|r| Event::GetAttachment(attachment_id, r))
                .await
            {
                Ok(attachment) => {
                    let mut reply: XpcMap = HashMap::new();
                    dict_put_str(&mut reply, "type", "GetAttachmentInfoResponse");
                    dict_put_str(
                        &mut reply,
                        "path",
                        &attachment.get_path_for_preview(false).to_string_lossy(),
                    );
                    dict_put_str(
                        &mut reply,
                        "preview_path",
                        &attachment.get_path_for_preview(true).to_string_lossy(),
                    );
                    dict_put_str(
                        &mut reply,
                        "downloaded",
                        &attachment.is_downloaded(false).to_string(),
                    );
                    dict_put_str(
                        &mut reply,
                        "preview_downloaded",
                        &attachment.is_downloaded(true).to_string(),
                    );
                    Message::Dictionary(reply)
                }
                Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
            }
        }

        "DownloadAttachment" => {
            let args = match get_dictionary_field(root, "arguments") {
                Some(a) => a,
                None => return make_error_reply("InvalidRequest", "Missing arguments"),
            };
            let attachment_id = match dict_get_str(args, "attachment_id") {
                Some(v) => v,
                None => return make_error_reply("InvalidRequest", "Missing attachment_id"),
            };
            let preview = dict_get_str(args, "preview")
                .map(|s| s == "true")
                .unwrap_or(false);
            match agent
                .send_event(|r| Event::DownloadAttachment(attachment_id, preview, r))
                .await
            {
                Ok(()) => make_ok_reply(),
                Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
            }
        }

        "UploadAttachment" => {
            use std::path::PathBuf;
            let args = match get_dictionary_field(root, "arguments") {
                Some(a) => a,
                None => return make_error_reply("InvalidRequest", "Missing arguments"),
            };
            let path = match dict_get_str(args, "path") {
                Some(v) => v,
                None => return make_error_reply("InvalidRequest", "Missing path"),
            };
            match agent
                .send_event(|r| Event::UploadAttachment(PathBuf::from(path), r))
                .await
            {
                Ok(upload_guid) => {
                    let mut reply: XpcMap = HashMap::new();
                    dict_put_str(&mut reply, "type", "UploadAttachmentResponse");
                    dict_put_str(&mut reply, "upload_guid", &upload_guid);
                    Message::Dictionary(reply)
                }
                Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
            }
        }

        "GetAllSettings" => match agent.send_event(Event::GetAllSettings).await {
            Ok(settings) => {
                let mut reply: XpcMap = HashMap::new();
                dict_put_str(&mut reply, "type", "GetAllSettingsResponse");
                dict_put_str(
                    &mut reply,
                    "server_url",
                    &settings.server_url.unwrap_or_default(),
                );
                dict_put_str(
                    &mut reply,
                    "username",
                    &settings.username.unwrap_or_default(),
                );
                Message::Dictionary(reply)
            }
            Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
        },

        "UpdateSettings" => {
            let args = match get_dictionary_field(root, "arguments") {
                Some(a) => a,
                None => return make_error_reply("InvalidRequest", "Missing arguments"),
            };
            let server_url = dict_get_str(args, "server_url");
            let username = dict_get_str(args, "username");
            let settings = Settings {
                server_url,
                username,
                token: None,
            };
            match agent
                .send_event(|r| Event::UpdateSettings(settings, r))
                .await
            {
                Ok(()) => make_ok_reply(),
                Err(e) => make_error_reply("DaemonError", &format!("{}", e)),
            }
        }

        // Subscribe and return immediately
        "SubscribeSignals" => {
            let mut list = subscribers.lock().unwrap();
            // Avoid duplicates
            if !list.iter().any(|c| c.0 == current_client) {
                list.push(XpcConn(current_client));
                log::info!(target: LOG_TARGET, "Client subscribed to signals (total subscribers: {})", list.len());
            }
            make_ok_reply()
        },

        // Unknown method fallback
        other => make_error_reply("UnknownMethod", other),
    };

    // Echo request_id back (if present) so clients can correlate replies
    response = attach_request_id(response, request_id);
    response
}

fn signal_to_message(signal: Signal) -> Message {
    let mut root: XpcMap = HashMap::new();
    let mut args: XpcMap = HashMap::new();
    match signal {
        Signal::ConversationsUpdated => {
            dict_put_str(&mut root, "name", "ConversationsUpdated");
        }
        Signal::MessagesUpdated(conversation_id) => {
            dict_put_str(&mut root, "name", "MessagesUpdated");
            dict_put_str(&mut args, "conversation_id", &conversation_id);
        }
        Signal::AttachmentDownloaded(attachment_id) => {
            dict_put_str(&mut root, "name", "AttachmentDownloadCompleted");
            dict_put_str(&mut args, "attachment_id", &attachment_id);
        }
        Signal::AttachmentUploaded(upload_guid, attachment_guid) => {
            dict_put_str(&mut root, "name", "AttachmentUploadCompleted");
            dict_put_str(&mut args, "upload_guid", &upload_guid);
            dict_put_str(&mut args, "attachment_guid", &attachment_guid);
        }
        Signal::UpdateStreamReconnected => {
            dict_put_str(&mut root, "name", "UpdateStreamReconnected");
        }
    }
    if !args.is_empty() {
        root.insert(cstr("arguments"), Message::Dictionary(args));
    }
    Message::Dictionary(root)
}

// legacy async client handler removed in reply-port implementation