use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use futures::future::BoxFuture;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;

#[derive(Debug, Clone)]
pub struct Mailbox<M> {
    tx: mpsc::UnboundedSender<M>,
}

impl<M> Mailbox<M> {
    pub fn send(&self, msg: M) -> bool {
        self.tx.send(msg).is_ok()
    }

    pub fn is_closed(&self) -> bool {
        self.tx.is_closed()
    }
}

type SpawnFn<M> = dyn Fn(Arc<ActorRegistry<M>>, Uuid, mpsc::UnboundedReceiver<M>) -> BoxFuture<'static, ()>
    + Send
    + Sync;

pub struct ActorRegistry<M> {
    name: &'static str,
    actors: Mutex<HashMap<Uuid, mpsc::UnboundedSender<M>>>,
    spawn: Arc<SpawnFn<M>>,
    shutdown: CancellationToken,
}

impl<M> std::fmt::Debug for ActorRegistry<M> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let live = self.actors.lock().unwrap().len();
        f.debug_struct("ActorRegistry")
            .field("name", &self.name)
            .field("live_actors", &live)
            .finish()
    }
}

impl<M> ActorRegistry<M>
where
    M: Send + 'static,
{
    pub fn new<F, Fut>(name: &'static str, spawn: F, shutdown: CancellationToken) -> Arc<Self>
    where
        F: Fn(Arc<Self>, Uuid, mpsc::UnboundedReceiver<M>) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ()> + Send + 'static,
    {
        let spawn: Arc<SpawnFn<M>> = Arc::new(move |reg, key, rx| Box::pin(spawn(reg, key, rx)));
        Arc::new(Self {
            name,
            actors: Mutex::new(HashMap::new()),
            spawn,
            shutdown,
        })
    }

    pub fn shutdown_token(&self) -> CancellationToken {
        self.shutdown.clone()
    }

    pub fn live_count(&self) -> usize {
        self.actors.lock().unwrap().len()
    }

    pub fn dispatch(self: &Arc<Self>, key: Uuid, mut msg: M) {
        {
            let guard = self.actors.lock().unwrap();
            if let Some(tx) = guard.get(&key) {
                match tx.send(msg) {
                    Ok(()) => return,
                    Err(e) => msg = e.0,
                }
            }
        }
        let (tx, rx) = mpsc::unbounded_channel();
        let mailbox = Mailbox { tx: tx.clone() };
        {
            let mut guard = self.actors.lock().unwrap();
            if let Some(existing) = guard.get(&key) {
                match existing.send(msg) {
                    Ok(()) => return,
                    Err(e) => msg = e.0,
                }
            }
            guard.insert(key, tx);
        }
        let registry = self.clone();
        let spawn = self.spawn.clone();
        tokio::spawn(async move {
            let fut = (spawn)(registry.clone(), key, rx);
            fut.await;
            registry.actors.lock().unwrap().remove(&key);
        });
        let _ = mailbox.send(msg);
        crate::metrics::actor_spawned(self.name);
    }
}

pub async fn run_actor<M>(
    mut rx: mpsc::UnboundedReceiver<M>,
    shutdown: CancellationToken,
    mut apply: impl FnMut(M) -> BoxFuture<'static, ()>,
) {
    loop {
        tokio::select! {
            _ = shutdown.cancelled() => break,
            maybe = rx.recv() => {
                match maybe {
                    Some(msg) => apply(msg).await,
                    None => break,
                }
            }
        }
    }
}

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

    #[tokio::test]
    async fn dispatch_spawns_and_orders() {
        let shutdown = CancellationToken::new();
        let seen = Arc::new(Mutex::new(Vec::new()));
        let reg = ActorRegistry::new(
            "test",
            {
                let seen = seen.clone();
                move |reg: Arc<ActorRegistry<u32>>, _key, rx| {
                    let seen = seen.clone();
                    async move {
                        run_actor(rx, reg.shutdown_token(), move |msg| {
                            let seen = seen.clone();
                            Box::pin(async move {
                                seen.lock().unwrap().push(msg);
                            })
                        })
                        .await;
                    }
                }
            },
            shutdown,
        );
        let key = Uuid::now_v7();
        reg.dispatch(key, 1);
        reg.dispatch(key, 2);
        reg.dispatch(key, 3);
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        let got = seen.lock().unwrap().clone();
        assert_eq!(got, vec![1, 2, 3]);
    }
}