use anyhow::Result;
use async_trait::async_trait;
use bytes::Bytes;
use std::collections::VecDeque;
use tokio::sync::Mutex;

/// A queue driver: push jobs in, pop them out.
#[async_trait]
pub trait Driver: Send + Sync + 'static {
    /// Block-pop: wait for a job. Returns `None` if shutdown.
    async fn pop(&self, queue: &str) -> Result<Option<Bytes>>;
    /// Push a job (for retry or testing).
    async fn push(&self, queue: &str, payload: Bytes) -> Result<()>;
    /// Current queue depth (for metrics).
    async fn depth(&self, queue: &str) -> Result<u64>;
}

/// In-memory driver (testing, no persistence).
pub struct MemoryDriver {
    queues: Mutex<std::collections::HashMap<String, VecDeque<Bytes>>>,
    notify: tokio::sync::Notify,
}

impl MemoryDriver {
    pub fn new() -> std::sync::Arc<Self> {
        std::sync::Arc::new(Self {
            queues: Mutex::new(Default::default()),
            notify: tokio::sync::Notify::new(),
        })
    }
}

#[async_trait]
impl Driver for MemoryDriver {
    async fn pop(&self, queue: &str) -> Result<Option<Bytes>> {
        loop {
            {
                let mut qs = self.queues.lock().await;
                if let Some(item) = qs.entry(queue.to_string()).or_default().pop_front() {
                    return Ok(Some(item));
                }
            }
            self.notify.notified().await;
        }
    }

    async fn push(&self, queue: &str, payload: Bytes) -> Result<()> {
        let mut qs = self.queues.lock().await;
        qs.entry(queue.to_string()).or_default().push_back(payload);
        self.notify.notify_one();
        Ok(())
    }

    async fn depth(&self, queue: &str) -> Result<u64> {
        Ok(self
            .queues
            .lock()
            .await
            .get(queue)
            .map(|q| q.len() as u64)
            .unwrap_or(0))
    }
}
