//! In-memory `ArchiveBackend` for coordinator unit tests.
//!
//! Fully supports every sub-trait. State lives in `Mutex<HashMap<…>>`.

#![cfg(any(test, feature = "test-utils"))]

use std::collections::HashMap;
use std::sync::Mutex;

use async_trait::async_trait;
use uuid::Uuid;

use crate::backend::{
    ArchiveBackend, ArchiveCapability, CapabilitySet, ConnectorRegistryBackend,
    DagBackend, HealthStatus, MetaEventsBackend, SessionMappingBackend,
};
use crate::error::{ArchivistError, Result};
use crate::types::{
    ConnectorRecord, DagEdge, MessageCursor, MessagePage, MessageRecord,
    MetaEventRecord, SessionListQuery, SessionMapping, SessionMetadata, SessionPage,
};

pub struct MockBackend {
    capabilities: CapabilitySet,
    sessions: Mutex<HashMap<Uuid, SessionMetadata>>,
    messages: Mutex<HashMap<Uuid, Vec<MessageRecord>>>,
    connectors: Mutex<HashMap<Uuid, ConnectorRecord>>,
    mappings: Mutex<HashMap<(Uuid, String), Uuid>>,
    meta_events: Mutex<HashMap<Uuid, Vec<MetaEventRecord>>>,
    dag_edges: Mutex<Vec<DagEdge>>,
    fail_next_writes: std::sync::atomic::AtomicUsize,
    fail_next_reads: std::sync::atomic::AtomicUsize,
    permanent_error: std::sync::Mutex<Option<String>>,
    append_calls: std::sync::Mutex<std::collections::HashMap<Uuid, usize>>,
    per_op_delay: std::sync::Mutex<std::time::Duration>,
}

impl MockBackend {
    pub fn new() -> Self {
        let mut capabilities = CapabilitySet::new();
        capabilities.insert(ArchiveCapability::Dag);
        capabilities.insert(ArchiveCapability::MetaEvents);
        capabilities.insert(ArchiveCapability::ConnectorRegistry);
        capabilities.insert(ArchiveCapability::SessionMapping);
        Self {
            capabilities,
            sessions: Mutex::new(HashMap::new()),
            messages: Mutex::new(HashMap::new()),
            connectors: Mutex::new(HashMap::new()),
            mappings: Mutex::new(HashMap::new()),
            meta_events: Mutex::new(HashMap::new()),
            dag_edges: Mutex::new(Vec::new()),
            fail_next_writes: std::sync::atomic::AtomicUsize::new(0),
            fail_next_reads: std::sync::atomic::AtomicUsize::new(0),
            permanent_error: std::sync::Mutex::new(None),
            append_calls: std::sync::Mutex::new(std::collections::HashMap::new()),
            per_op_delay: std::sync::Mutex::new(std::time::Duration::ZERO),
        }
    }
}

impl MockBackend {
    /// Build a mock with the exact capability set provided. All other state
    /// starts empty (same as `new()`).
    pub fn with_capabilities(capabilities: CapabilitySet) -> Self {
        let mut m = Self::new();
        m.capabilities = capabilities;
        m
    }

    /// Test helper: does this mock have any meta events for the given session?
    pub fn has_meta_events(&self, scroll_id: uuid::Uuid) -> bool {
        self.meta_events
            .lock()
            .unwrap()
            .get(&scroll_id)
            .map(|v| !v.is_empty())
            .unwrap_or(false)
    }

    /// Queue up `count` injected write failures. The next `count` calls to
    /// any mutating API return `ArchivistError::Other("injected write failure")`
    /// before touching state.
    pub fn inject_write_failures(&self, count: usize) {
        self.fail_next_writes
            .store(count, std::sync::atomic::Ordering::SeqCst);
    }

    /// Queue up `count` injected read failures for per-scroll_id reads.
    pub fn inject_read_failures(&self, count: usize) {
        self.fail_next_reads
            .store(count, std::sync::atomic::Ordering::SeqCst);
    }

    /// Simulate a permanently broken backend.
    pub fn break_permanently(&self, reason: impl Into<String>) {
        *self.permanent_error.lock().unwrap() = Some(reason.into());
    }

    pub fn clear_failures(&self) {
        self.fail_next_writes
            .store(0, std::sync::atomic::Ordering::SeqCst);
        self.fail_next_reads
            .store(0, std::sync::atomic::Ordering::SeqCst);
        *self.permanent_error.lock().unwrap() = None;
    }

    /// Test helper: how many `MessageRecord`s this mock has for the given session.
    pub fn appended_count(&self, scroll_id: uuid::Uuid) -> usize {
        self.messages
            .lock()
            .unwrap()
            .get(&scroll_id)
            .map(|v| v.len())
            .unwrap_or(0)
    }

    /// Test helper: how many times `append_messages` was invoked for the
    /// given session (regardless of message count per invocation).
    pub fn append_call_count(&self, scroll_id: uuid::Uuid) -> usize {
        self.append_calls
            .lock()
            .unwrap()
            .get(&scroll_id)
            .copied()
            .unwrap_or(0)
    }

    /// Test helper: artificially slow every mutating backend operation by
    /// sleeping `d` before it touches state. Used to simulate a slow backend
    /// for backpressure tests.
    pub fn set_per_op_delay(&self, d: std::time::Duration) {
        *self.per_op_delay.lock().unwrap() = d;
    }

    async fn maybe_delay(&self) {
        let d = *self.per_op_delay.lock().unwrap();
        if !d.is_zero() {
            tokio::time::sleep(d).await;
        }
    }

    pub(crate) fn check_write_failure(&self) -> Result<()> {
        if let Some(reason) = self.permanent_error.lock().unwrap().clone() {
            return Err(ArchivistError::Other(reason));
        }
        let prev = self
            .fail_next_writes
            .fetch_update(
                std::sync::atomic::Ordering::SeqCst,
                std::sync::atomic::Ordering::SeqCst,
                |n| if n > 0 { Some(n - 1) } else { None },
            )
            .ok();
        if prev.is_some() {
            return Err(ArchivistError::Other("injected write failure".into()));
        }
        Ok(())
    }

    pub(crate) fn check_read_failure(&self) -> Result<()> {
        if let Some(reason) = self.permanent_error.lock().unwrap().clone() {
            return Err(ArchivistError::Other(reason));
        }
        let prev = self
            .fail_next_reads
            .fetch_update(
                std::sync::atomic::Ordering::SeqCst,
                std::sync::atomic::Ordering::SeqCst,
                |n| if n > 0 { Some(n - 1) } else { None },
            )
            .ok();
        if prev.is_some() {
            return Err(ArchivistError::Other("injected read failure".into()));
        }
        Ok(())
    }
}

impl Default for MockBackend {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl ArchiveBackend for MockBackend {
    fn capabilities(&self) -> &CapabilitySet {
        &self.capabilities
    }
    async fn health_check(&self) -> HealthStatus {
        HealthStatus::Healthy
    }

    async fn put_session(&self, meta: SessionMetadata) -> Result<()> {
        self.check_write_failure()?;
        self.maybe_delay().await;
        self.sessions.lock().unwrap().insert(meta.scroll_id, meta);
        Ok(())
    }
    async fn get_session(&self, scroll_id: Uuid) -> Result<Option<SessionMetadata>> {
        self.check_read_failure()?;
        Ok(self.sessions.lock().unwrap().get(&scroll_id).cloned())
    }
    async fn list_sessions_paged(&self, query: SessionListQuery) -> Result<SessionPage> {
        let mut items: Vec<SessionMetadata> =
            self.sessions.lock().unwrap().values().cloned().collect();
        if !query.connector_uids.is_empty() {
            items.retain(|s| query.connector_uids.contains(&s.connector_uid));
        }
        items.sort_by(|a, b| {
            b.updated_at
                .cmp(&a.updated_at)
                .then(b.scroll_id.cmp(&a.scroll_id))
        });
        let limit = query.limit.min(crate::types::MAX_PAGE_LIMIT).max(1);
        let total_count = items.len();
        let items: Vec<_> = items.into_iter().take(limit).collect();
        Ok(SessionPage {
            items,
            next_cursor: None,
            total_count: Some(total_count),
        })
    }
    async fn delete_session(&self, scroll_id: Uuid) -> Result<()> {
        self.check_write_failure()?;
        self.maybe_delay().await;
        if self.sessions.lock().unwrap().remove(&scroll_id).is_none() {
            return Err(ArchivistError::SessionUnknown(scroll_id));
        }
        self.messages.lock().unwrap().remove(&scroll_id);
        Ok(())
    }

    async fn append_messages(
        &self,
        scroll_id: Uuid,
        msgs: Vec<MessageRecord>,
    ) -> Result<()> {
        self.check_write_failure()?;
        self.maybe_delay().await;
        *self
            .append_calls
            .lock()
            .unwrap()
            .entry(scroll_id)
            .or_insert(0) += 1;
        self.messages
            .lock()
            .unwrap()
            .entry(scroll_id)
            .or_default()
            .extend(msgs);
        Ok(())
    }
    async fn get_messages_paged(
        &self,
        scroll_id: Uuid,
        cursor: Option<MessageCursor>,
        limit: usize,
    ) -> Result<MessagePage> {
        self.check_read_failure()?;
        let mut all = self
            .messages
            .lock()
            .unwrap()
            .get(&scroll_id)
            .cloned()
            .unwrap_or_default();
        all.sort_by(|a, b| a.ts.cmp(&b.ts).then(a.message_id.cmp(&b.message_id)));
        if let Some(c) = cursor.as_ref() {
            all.retain(|m| (m.ts, m.message_id) > (c.ts, c.message_id));
        }
        let total = all.len();
        let taken: Vec<_> = all.into_iter().take(limit.max(1)).collect();
        let next_cursor = if total > taken.len() {
            taken.last().map(|m| MessageCursor {
                ts: m.ts,
                message_id: m.message_id,
            })
        } else {
            None
        };
        Ok(MessagePage {
            items: taken,
            next_cursor,
        })
    }
    async fn count_messages(&self, scroll_id: Uuid) -> Result<usize> {
        self.check_read_failure()?;
        Ok(self
            .messages
            .lock()
            .unwrap()
            .get(&scroll_id)
            .map(|v| v.len())
            .unwrap_or(0))
    }
    async fn clear_session_messages(&self, scroll_id: Uuid) -> Result<()> {
        self.check_write_failure()?;
        self.maybe_delay().await;
        self.messages.lock().unwrap().remove(&scroll_id);
        Ok(())
    }

    fn as_dag(&self) -> Option<&dyn DagBackend> {
        if self.capabilities.contains(&ArchiveCapability::Dag) {
            Some(self)
        } else {
            None
        }
    }
    fn as_meta_events(&self) -> Option<&dyn MetaEventsBackend> {
        if self.capabilities.contains(&ArchiveCapability::MetaEvents) {
            Some(self)
        } else {
            None
        }
    }
    fn as_connector_registry(&self) -> Option<&dyn ConnectorRegistryBackend> {
        if self
            .capabilities
            .contains(&ArchiveCapability::ConnectorRegistry)
        {
            Some(self)
        } else {
            None
        }
    }
    fn as_session_mapping(&self) -> Option<&dyn SessionMappingBackend> {
        if self
            .capabilities
            .contains(&ArchiveCapability::SessionMapping)
        {
            Some(self)
        } else {
            None
        }
    }
}

#[async_trait]
impl ConnectorRegistryBackend for MockBackend {
    async fn put_connector(&self, record: ConnectorRecord) -> Result<()> {
        self.check_write_failure()?;
        self.maybe_delay().await;
        self.connectors
            .lock()
            .unwrap()
            .insert(record.connector_uid, record);
        Ok(())
    }
    async fn get_connector(&self, connector_uid: Uuid) -> Result<Option<ConnectorRecord>> {
        Ok(self
            .connectors
            .lock()
            .unwrap()
            .get(&connector_uid)
            .cloned())
    }
    async fn list_connectors(&self) -> Result<Vec<ConnectorRecord>> {
        Ok(self.connectors.lock().unwrap().values().cloned().collect())
    }
    async fn resolve_connector_uid(&self, client_native_id: &str) -> Result<Option<Uuid>> {
        Ok(self
            .connectors
            .lock()
            .unwrap()
            .values()
            .find(|c| c.client_native_id == client_native_id)
            .map(|c| c.connector_uid))
    }
    async fn update_connector_fingerprint(
        &self,
        connector_uid: Uuid,
        fingerprint: String,
    ) -> Result<()> {
        self.check_write_failure()?;
        self.maybe_delay().await;
        if let Some(r) = self.connectors.lock().unwrap().get_mut(&connector_uid) {
            r.fingerprint = Some(fingerprint);
            Ok(())
        } else {
            Err(ArchivistError::ConnectorUnknown(connector_uid))
        }
    }
}

#[async_trait]
impl SessionMappingBackend for MockBackend {
    async fn put_mapping(
        &self,
        connector_uid: Uuid,
        native_session_id: &str,
        scroll_id: Uuid,
    ) -> Result<()> {
        self.check_write_failure()?;
        self.maybe_delay().await;
        self.mappings
            .lock()
            .unwrap()
            .insert((connector_uid, native_session_id.to_string()), scroll_id);
        Ok(())
    }
    async fn get_mapping(
        &self,
        connector_uid: Uuid,
        native_session_id: &str,
    ) -> Result<Option<Uuid>> {
        Ok(self
            .mappings
            .lock()
            .unwrap()
            .get(&(connector_uid, native_session_id.to_string()))
            .copied())
    }
    async fn list_mappings_for_connector(
        &self,
        connector_uid: Uuid,
    ) -> Result<Vec<SessionMapping>> {
        Ok(self
            .mappings
            .lock()
            .unwrap()
            .iter()
            .filter(|((c, _), _)| *c == connector_uid)
            .map(|((c, n), s)| SessionMapping {
                version: 1,
                connector_uid: *c,
                native_session_id: n.clone(),
                scroll_id: *s,
                created_at: chrono::Utc::now(),
                alias_of: None,
            })
            .collect())
    }
    async fn find_owner(&self, native_session_id: &str) -> Result<Option<(Uuid, Uuid)>> {
        Ok(self
            .mappings
            .lock()
            .unwrap()
            .iter()
            .find(|((_, n), _)| n == native_session_id)
            .map(|((c, _), s)| (*c, *s)))
    }

    async fn rewrite_connector_mappings(
        &self,
        connector_uid: Uuid,
        mappings: Vec<SessionMapping>,
    ) -> Result<()> {
        self.check_write_failure()?;
        self.maybe_delay().await;
        let mut map = self.mappings.lock().unwrap();
        map.retain(|(c, _), _| *c != connector_uid);
        for m in mappings {
            map.insert((connector_uid, m.native_session_id), m.scroll_id);
        }
        Ok(())
    }
}

#[async_trait]
impl DagBackend for MockBackend {
    async fn append_dag_edge(&self, edge: DagEdge) -> Result<()> {
        self.check_write_failure()?;
        self.maybe_delay().await;
        self.dag_edges.lock().unwrap().push(edge);
        Ok(())
    }
    async fn get_children(&self, parent: Uuid) -> Result<Vec<SessionMetadata>> {
        self.check_read_failure()?;
        let edges = self.dag_edges.lock().unwrap();
        let sessions = self.sessions.lock().unwrap();
        Ok(edges
            .iter()
            .filter(|e| e.parent == parent)
            .filter_map(|e| sessions.get(&e.child).cloned())
            .collect())
    }
    async fn get_dag_edges(&self, root: Uuid) -> Result<Vec<DagEdge>> {
        self.check_read_failure()?;
        Ok(self
            .dag_edges
            .lock()
            .unwrap()
            .iter()
            .filter(|e| e.parent == root)
            .cloned()
            .collect())
    }
}

#[async_trait]
impl MetaEventsBackend for MockBackend {
    async fn append_meta_events(
        &self,
        scroll_id: Uuid,
        events: Vec<MetaEventRecord>,
    ) -> Result<()> {
        self.check_write_failure()?;
        self.maybe_delay().await;
        self.meta_events
            .lock()
            .unwrap()
            .entry(scroll_id)
            .or_default()
            .extend(events);
        Ok(())
    }
    async fn get_meta_events(&self, scroll_id: Uuid) -> Result<Vec<MetaEventRecord>> {
        self.check_read_failure()?;
        Ok(self
            .meta_events
            .lock()
            .unwrap()
            .get(&scroll_id)
            .cloned()
            .unwrap_or_default())
    }
    async fn update_meta_session_status(
        &self,
        scroll_id: Uuid,
        is_connected: bool,
        current_session_id: Option<Uuid>,
    ) -> Result<()> {
        self.check_write_failure()?;
        self.maybe_delay().await;
        if let Some(s) = self.sessions.lock().unwrap().get_mut(&scroll_id) {
            s.is_connected = Some(is_connected);
            s.current_session_id = current_session_id;
            Ok(())
        } else {
            Err(ArchivistError::SessionUnknown(scroll_id))
        }
    }
    async fn list_meta_sessions(&self) -> Result<Vec<SessionMetadata>> {
        Ok(self
            .sessions
            .lock()
            .unwrap()
            .values()
            .filter(|s| matches!(s.kind, crate::types::SessionKind::AcpConnection))
            .cloned()
            .collect())
    }
    async fn find_meta_session_by_client(
        &self,
        client_id: &str,
    ) -> Result<Option<SessionMetadata>> {
        Ok(self
            .sessions
            .lock()
            .unwrap()
            .values()
            .find(|s| s.acp_client_id.as_deref() == Some(client_id))
            .cloned())
    }
}

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

    #[tokio::test]
    async fn injected_write_failure_returns_error_then_recovers() {
        let m = MockBackend::new();
        m.inject_write_failures(2);
        let scroll = uuid::Uuid::nil();
        assert!(m.append_messages(scroll, vec![]).await.is_err());
        assert!(m.append_messages(scroll, vec![]).await.is_err());
        assert!(m.append_messages(scroll, vec![]).await.is_ok()); // back to normal
    }
}
