//! Method registry — maps method IDs to handler functions.

use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use bytes::Bytes;

use crate::error::RpcStatus;

/// The result of handling an RPC request.
pub struct HandlerResult {
    pub status: RpcStatus,
    pub payload: Bytes,
}

impl HandlerResult {
    /// Shorthand for a successful response.
    pub fn ok(payload: Bytes) -> Self {
        Self {
            status: RpcStatus::Ok,
            payload,
        }
    }

    /// Shorthand for an error response.
    pub fn err(status: RpcStatus, message: &str) -> Self {
        Self {
            status,
            payload: Bytes::copy_from_slice(message.as_bytes()),
        }
    }
}

/// Context passed to every RPC handler.
pub struct RequestContext {
    /// The authenticated identity key of the caller, if any.
    pub identity_key: Option<Vec<u8>>,
    /// The session token, if provided.
    pub session_token: Option<Vec<u8>>,
    /// The raw request payload (protobuf-encoded).
    pub payload: Bytes,
}

/// Type-erased async handler function.
pub type HandlerFn<S> = Arc<
    dyn Fn(Arc<S>, RequestContext) -> Pin<Box<dyn Future<Output = HandlerResult> + Send>>
        + Send
        + Sync,
>;

/// Registry mapping method IDs to handler functions.
pub struct MethodRegistry<S> {
    handlers: HashMap<u16, (HandlerFn<S>, &'static str)>,
}

impl<S: Send + Sync + 'static> MethodRegistry<S> {
    pub fn new() -> Self {
        Self {
            handlers: HashMap::new(),
        }
    }

    /// Register a handler for a method ID.
    pub fn register<F, Fut>(&mut self, method_id: u16, name: &'static str, handler: F)
    where
        F: Fn(Arc<S>, RequestContext) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = HandlerResult> + Send + 'static,
    {
        let handler = Arc::new(move |state: Arc<S>, ctx: RequestContext| {
            Box::pin(handler(state, ctx)) as Pin<Box<dyn Future<Output = HandlerResult> + Send>>
        });
        self.handlers.insert(method_id, (handler, name));
    }

    /// Look up a handler by method ID.
    pub fn get(&self, method_id: u16) -> Option<&(HandlerFn<S>, &'static str)> {
        self.handlers.get(&method_id)
    }

    /// Return the number of registered methods.
    pub fn len(&self) -> usize {
        self.handlers.len()
    }

    /// Whether the registry is empty.
    pub fn is_empty(&self) -> bool {
        self.handlers.is_empty()
    }

    /// Iterate over all registered (method_id, name) pairs.
    pub fn methods(&self) -> impl Iterator<Item = (u16, &'static str)> + '_ {
        self.handlers.iter().map(|(&id, (_, name))| (id, *name))
    }
}

impl<S: Send + Sync + 'static> Default for MethodRegistry<S> {
    fn default() -> Self {
        Self::new()
    }
}