use anyhow::{anyhow, Context, Result};
use async_trait::async_trait;
use redb::{Database, ReadableDatabase, ReadableTable, TableDefinition};
use rmqtt::message::MessageManager;
use rmqtt::types::{ClientId, From as MsgFrom, MsgID, Publish, SharedGroup, TopicFilter, TopicName};
use rmqtt::Result as RmqttResult;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet, VecDeque};
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

const COUNTER_MASK: u64 = 0xFFFF_FFFF_FFFF;
const MESSAGES: TableDefinition<u64, &[u8]> = TableDefinition::new("messages");

fn enc<T: Serialize>(v: &T) -> Result<Vec<u8>> {
    bincode::serde::encode_to_vec(v, bincode::config::standard()).map_err(|e| anyhow!("bincode encode: {e}"))
}

fn dec<T: serde::de::DeserializeOwned>(bytes: &[u8]) -> Result<T> {
    bincode::serde::decode_from_slice(bytes, bincode::config::standard())
        .map(|(v, _)| v)
        .map_err(|e| anyhow!("bincode decode: {e}"))
}

#[derive(Clone, Serialize, Deserialize)]
struct StoredMsg {
    from: MsgFrom,
    publish: Publish,
    expiry_at: i64,
    delivered_to: HashSet<String>,
}

fn approx_size(m: &StoredMsg) -> usize {
    m.publish.payload.len() + m.publish.topic.len() + 384
}

struct MemState {
    msgs: HashMap<u64, StoredMsg>,
    order: VecDeque<u64>,
    by_topic: HashMap<TopicName, HashSet<u64>>,
    bytes: usize,
    dirty: HashSet<u64>,
    removed: HashSet<u64>,
}

impl MemState {
    fn evict(&mut self, now: i64, max_msgs: usize, max_bytes: usize) {
        while let Some(&id) = self.order.front() {
            let over_cap = (max_msgs > 0 && self.msgs.len() > max_msgs)
                || (max_bytes > 0 && self.bytes > max_bytes);
            let expired = self.msgs.get(&id).map(|m| m.expiry_at <= now).unwrap_or(true);
            if !over_cap && !expired {
                break;
            }
            self.order.pop_front();
            if let Some(m) = self.msgs.remove(&id) {
                self.bytes = self.bytes.saturating_sub(approx_size(&m));
                self.dirty.remove(&id);
                self.removed.insert(id);
                if let Some(ids) = self.by_topic.get_mut(&m.publish.topic) {
                    ids.remove(&id);
                    if ids.is_empty() {
                        self.by_topic.remove(&m.publish.topic);
                    }
                }
            }
        }
    }
}

#[derive(Clone)]
pub struct RedbMessageStore {
    db: Arc<Database>,
    node_id: u64,
    counter: Arc<AtomicU64>,
    state: Arc<Mutex<MemState>>,
    max_msgs: usize,
    max_bytes: usize,
    tx: tokio::sync::mpsc::UnboundedSender<(u64, StoredMsg)>,
}

impl RedbMessageStore {
    pub fn open(path: &Path, node_id: u64) -> Result<Self> {
        if let Some(dir) = path.parent() {
            if !dir.as_os_str().is_empty() {
                std::fs::create_dir_all(dir)?;
            }
        }
        let db = Database::builder()
            .set_cache_size(crate::db::redb_cache_bytes())
            .create(path)
            .context("opening redb message store")?;
        let wtx = db.begin_write()?;
        {
            wtx.open_table(MESSAGES)?;
        }
        wtx.commit()?;

        let now = now_millis();
        let mut msgs = HashMap::new();
        let mut start = 0u64;
        {
            let rtx = db.begin_read()?;
            let t = rtx.open_table(MESSAGES)?;
            for item in t.iter()? {
                let (k, v) = item?;
                let id = k.value();
                if (id >> 48) == node_id {
                    let c = (id & COUNTER_MASK) + 1;
                    if c > start {
                        start = c;
                    }
                }
                if let Ok(msg) = dec::<StoredMsg>(v.value()) {
                    if msg.expiry_at > now {
                        msgs.insert(id, msg);
                    }
                }
            }
        }
        let mut entries: Vec<(i64, u64)> = msgs.iter().map(|(id, m)| (m.expiry_at, *id)).collect();
        entries.sort_unstable();
        let order: VecDeque<u64> = entries.into_iter().map(|(_, id)| id).collect();
        let bytes: usize = msgs.values().map(approx_size).sum();
        let mut by_topic: HashMap<TopicName, HashSet<u64>> = HashMap::new();
        for (id, m) in &msgs {
            by_topic.entry(m.publish.topic.clone()).or_default().insert(*id);
        }

        let max_msgs = std::env::var("MSG_MAX_MSGS").ok().and_then(|s| s.parse().ok()).unwrap_or(0);
        let max_bytes = std::env::var("MSG_MAX_BYTES")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(256 * 1024 * 1024);
        log::info!("message store memory cap: max_msgs={max_msgs} max_bytes={max_bytes} (0 = unbounded)");

        let state = Arc::new(Mutex::new(MemState {
            msgs,
            order,
            by_topic,
            bytes,
            dirty: HashSet::new(),
            removed: HashSet::new(),
        }));
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<(u64, StoredMsg)>();
        {
            let state = state.clone();
            tokio::spawn(async move {
                while let Some((key, msg)) = rx.recv().await {
                    let now = now_millis();
                    let size = approx_size(&msg);
                    let topic = msg.publish.topic.clone();
                    if let Ok(mut st) = state.lock() {
                        st.removed.remove(&key);
                        match st.msgs.insert(key, msg) {
                            Some(old) => {
                                st.bytes = st.bytes.saturating_sub(approx_size(&old)).saturating_add(size);
                            }
                            None => {
                                st.order.push_back(key);
                                st.bytes = st.bytes.saturating_add(size);
                            }
                        }
                        st.by_topic.entry(topic).or_default().insert(key);
                        st.dirty.insert(key);
                        st.evict(now, max_msgs, max_bytes);
                    }
                }
            });
        }
        Ok(Self {
            db: Arc::new(db),
            node_id,
            counter: Arc::new(AtomicU64::new(start)),
            state,
            max_msgs,
            max_bytes,
            tx,
        })
    }

    pub async fn flush(&self) {
        let now = now_millis();
        let (batch, removed): (Vec<(u64, Vec<u8>)>, Vec<u64>) = {
            let mut st = match self.state.lock() {
                Ok(g) => g,
                Err(_) => return,
            };
            st.evict(now, self.max_msgs, self.max_bytes);
            let dirty: Vec<u64> = st.dirty.drain().collect();
            let mut batch = Vec::with_capacity(dirty.len());
            for id in dirty {
                if let Some(m) = st.msgs.get(&id) {
                    match enc(m) {
                        Ok(bytes) => batch.push((id, bytes)),
                        Err(e) => log::warn!("message store encode failed: {e}"),
                    }
                }
            }
            let removed: Vec<u64> = st.removed.drain().collect();
            (batch, removed)
        };
        if batch.is_empty() && removed.is_empty() {
            return;
        }
        let db = self.db.clone();
        let res = tokio::task::spawn_blocking(move || -> Result<()> {
            let wtx = db.begin_write()?;
            {
                let mut t = wtx.open_table(MESSAGES)?;
                for id in removed {
                    t.remove(id)?;
                }
                for (id, bytes) in batch {
                    t.insert(id, bytes.as_slice())?;
                }
            }
            wtx.commit()?;
            Ok(())
        })
        .await;
        if let Ok(Err(e)) = res {
            log::warn!("message store flush failed: {e}");
        }
    }
}

fn now_millis() -> i64 {
    SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_millis() as i64).unwrap_or(0)
}

fn topic_matches(topic: &str, filter: &str) -> bool {
    let mut t = topic.split('/');
    let mut f = filter.split('/');
    loop {
        match (f.next(), t.next()) {
            (Some("#"), _) => return true,
            (Some("+"), Some(_)) => continue,
            (Some(fseg), Some(tseg)) if fseg == tseg => continue,
            (None, None) => return true,
            _ => return false,
        }
    }
}

fn has_wildcard(filter: &str) -> bool {
    filter.bytes().any(|b| b == b'+' || b == b'#')
}

#[async_trait]
impl MessageManager for RedbMessageStore {
    fn enable(&self) -> bool {
        true
    }

    fn next_msg_id(&self) -> MsgID {
        let c = self.counter.fetch_add(1, Ordering::Relaxed) & COUNTER_MASK;
        ((self.node_id << 48) | c) as MsgID
    }

    async fn store(
        &self,
        msg_id: MsgID,
        from: MsgFrom,
        p: Publish,
        expiry_interval: Duration,
        sub_client_ids: Option<Vec<(ClientId, Option<(TopicFilter, SharedGroup)>)>>,
    ) -> RmqttResult<()> {
        let now = now_millis();
        let expiry_at = now.saturating_add(expiry_interval.as_millis() as i64);
        let delivered_to: HashSet<String> = sub_client_ids
            .unwrap_or_default()
            .into_iter()
            .map(|(cid, _)| cid.to_string())
            .collect();
        let key = msg_id as u64;
        let msg = StoredMsg { from, publish: p, expiry_at, delivered_to };
        let _ = self.tx.send((key, msg));
        Ok(())
    }

    async fn get(
        &self,
        client_id: &str,
        topic_filter: &str,
        _group: Option<&SharedGroup>,
    ) -> RmqttResult<Vec<(MsgID, MsgFrom, Publish)>> {
        let now = now_millis();
        let mut out = Vec::new();
        if let Ok(mut st) = self.state.lock() {
            st.evict(now, self.max_msgs, self.max_bytes);
            let candidates: Vec<u64> = if has_wildcard(topic_filter) {
                st.by_topic
                    .iter()
                    .filter_map(|(t, ids)| topic_matches(t.as_ref(), topic_filter).then_some(ids))
                    .flatten()
                    .copied()
                    .collect()
            } else {
                let key = TopicName::from(topic_filter.to_owned());
                st.by_topic.get(&key).map(|ids| ids.iter().copied().collect()).unwrap_or_default()
            };
            let mut touched = Vec::new();
            for id in candidates {
                if let Some(msg) = st.msgs.get(&id) {
                    if msg.expiry_at <= now || msg.delivered_to.contains(client_id) {
                        continue;
                    }
                    out.push((id as MsgID, msg.from.clone(), msg.publish.clone()));
                    touched.push(id);
                }
            }
            for id in touched {
                if let Some(m) = st.msgs.get_mut(&id) {
                    m.delivered_to.insert(client_id.to_string());
                }
                st.dirty.insert(id);
            }
        }
        Ok(out)
    }

    async fn count(&self) -> isize {
        match self.state.lock() {
            Ok(st) => st.msgs.len() as isize,
            Err(_) => -1,
        }
    }

    fn should_merge_on_get(&self) -> bool {
        true
    }
}
