//! 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 std::time::Duration;

use bytes::Bytes;
use tokio::time::Instant;

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,
    /// Unique correlation ID for request tracing (UUID v7, monotonic).
    pub trace_id: String,
    /// The effective deadline for this request. Handlers can check this to bail
    /// early on long-running operations. `None` means no deadline.
    pub deadline: Option<Instant>,
}

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

/// Per-method registration entry.
struct MethodEntry<S> {
    handler: HandlerFn<S>,
    name: &'static str,
    /// Optional per-method timeout override. `None` means use the server default.
    timeout: Option<Duration>,
}

/// Registry mapping method IDs to handler functions.
pub struct MethodRegistry<S> {
    handlers: HashMap<u16, MethodEntry<S>>,
    /// Default timeout applied to methods that don't specify their own.
    default_timeout: Option<Duration>,
}

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

    /// Set the default timeout for all methods that don't have a per-method override.
    pub fn set_default_timeout(&mut self, timeout: Duration) {
        self.default_timeout = Some(timeout);
    }

    /// 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, MethodEntry { handler, name, timeout: None });
    }

    /// Register a handler with a per-method timeout override.
    pub fn register_with_timeout<F, Fut>(
        &mut self,
        method_id: u16,
        name: &'static str,
        timeout: Duration,
        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, MethodEntry { handler, name, timeout: Some(timeout) });
    }

    /// Look up a handler, name, and effective timeout by method ID.
    pub fn get(&self, method_id: u16) -> Option<(&HandlerFn<S>, &'static str, Option<Duration>)> {
        self.handlers.get(&method_id).map(|e| {
            (&e.handler, e.name, e.timeout.or(self.default_timeout))
        })
    }

    /// 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, entry)| (id, entry.name))
    }
}

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

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

    #[test]
    fn registry_default_timeout_applies_to_methods() {
        let mut reg = MethodRegistry::<()>::new();
        reg.set_default_timeout(Duration::from_secs(30));
        reg.register(1, "Test", |_state: Arc<()>, _ctx| async { HandlerResult::ok(Bytes::new()) });

        let (_, name, timeout) = reg.get(1).expect("registered method");
        assert_eq!(name, "Test");
        assert_eq!(timeout, Some(Duration::from_secs(30)));
    }

    #[test]
    fn registry_per_method_timeout_overrides_default() {
        let mut reg = MethodRegistry::<()>::new();
        reg.set_default_timeout(Duration::from_secs(30));
        reg.register_with_timeout(
            1,
            "Slow",
            Duration::from_secs(120),
            |_state: Arc<()>, _ctx| async { HandlerResult::ok(Bytes::new()) },
        );

        let (_, _, timeout) = reg.get(1).expect("registered method");
        assert_eq!(timeout, Some(Duration::from_secs(120)));
    }

    #[test]
    fn registry_no_default_timeout_returns_none() {
        let mut reg = MethodRegistry::<()>::new();
        reg.register(1, "NoTimeout", |_state: Arc<()>, _ctx| async {
            HandlerResult::ok(Bytes::new())
        });

        let (_, _, timeout) = reg.get(1).expect("registered method");
        assert_eq!(timeout, None);
    }

    #[test]
    fn request_context_deadline_is_accessible() {
        let ctx = RequestContext {
            identity_key: None,
            session_token: None,
            payload: Bytes::new(),
            trace_id: String::new(),
            deadline: Some(Instant::now() + Duration::from_secs(10)),
        };
        assert!(ctx.deadline.is_some());

        let ctx_no_deadline = RequestContext {
            identity_key: None,
            session_token: None,
            payload: Bytes::new(),
            trace_id: String::new(),
            deadline: None,
        };
        assert!(ctx_no_deadline.deadline.is_none());
    }
}
