//! Session Persistence Backends

use std::collections::HashMap;
use std::sync::Arc;

use tokio::sync::RwLock;
use uuid::Uuid;

use super::archive::verify_restored_session_roundtrip;
use chrono::{DateTime, Utc};

use super::state::{Session, SessionId, SessionMessage, SessionState, SessionType};
use super::types::QueueItem;
use super::{SessionError, SessionResult};
use crate::graph::{GraphEvent, GraphMaterializer, GraphValidator, SessionGraph};

pub(crate) fn validate_session_graph(session: &Session, backend: &str) -> SessionResult<()> {
    let report = GraphValidator::validate(&session.graph);
    if report.is_valid() {
        return Ok(());
    }

    let issues = report
        .issues
        .into_iter()
        .map(|issue| format!("{}: {}", issue.code, issue.message))
        .collect::<Vec<_>>()
        .join("; ");
    Err(SessionError::Storage {
        message: format!(
            "Invalid graph for session {} in {backend} persistence: {issues}",
            session.id
        ),
    })
}

/// Filter criteria for searching sessions.
///
/// All fields are optional; only set fields are applied as AND conditions.
/// Applications build richer search on top of this (e.g., full-text search
/// in summaries, tag-based filtering).
#[derive(Debug, Clone, Default)]
pub struct SessionFilter {
    /// Only sessions created after this time.
    pub created_after: Option<DateTime<Utc>>,
    /// Only sessions created before this time.
    pub created_before: Option<DateTime<Utc>>,
    /// Only sessions of this type (Main or Subagent).
    pub session_type: Option<SessionType>,
    /// Only sessions owned by this tenant.
    pub tenant_id: Option<String>,
    /// Only sessions owned by this principal.
    pub principal_id: Option<String>,
    /// Only sessions in this state.
    pub state: Option<SessionState>,
    /// Maximum number of results.
    pub limit: Option<usize>,
}

impl SessionFilter {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn created_after(mut self, time: DateTime<Utc>) -> Self {
        self.created_after = Some(time);
        self
    }

    pub fn created_before(mut self, time: DateTime<Utc>) -> Self {
        self.created_before = Some(time);
        self
    }

    pub fn tenant(mut self, tenant_id: impl Into<String>) -> Self {
        self.tenant_id = Some(tenant_id.into());
        self
    }

    pub fn principal(mut self, principal_id: impl Into<String>) -> Self {
        self.principal_id = Some(principal_id.into());
        self
    }

    pub fn limit(mut self, limit: usize) -> Self {
        self.limit = Some(limit);
        self
    }

    /// Check if a session matches this filter.
    pub fn matches(&self, session: &Session) -> bool {
        if self
            .created_after
            .is_some_and(|after| session.created_at < after)
        {
            return false;
        }
        if self
            .created_before
            .is_some_and(|before| session.created_at > before)
        {
            return false;
        }
        if let Some(ref st) = self.session_type
            && std::mem::discriminant(&session.session_type) != std::mem::discriminant(st)
        {
            return false;
        }
        if let Some(ref tid) = self.tenant_id
            && session.tenant_id.as_ref() != Some(tid)
        {
            return false;
        }
        if let Some(ref pid) = self.principal_id
            && session.principal_id.as_ref() != Some(pid)
        {
            return false;
        }
        if let Some(ref state) = self.state
            && session.state != *state
        {
            return false;
        }
        true
    }
}

/// A boxed synchronous mutation applied to a session inside [`Persistence::with_session_lock`].
pub type SessionMutationFn = Box<dyn FnOnce(&mut Session) -> SessionResult<()> + Send>;

#[async_trait::async_trait]
pub trait Persistence: Send + Sync {
    fn name(&self) -> &str;

    // Core CRUD
    async fn save(&self, session: &Session) -> SessionResult<()>;
    async fn load(&self, id: &SessionId) -> SessionResult<Option<Session>>;
    async fn delete(&self, id: &SessionId) -> SessionResult<bool>;
    async fn list(&self, tenant_id: Option<&str>) -> SessionResult<Vec<SessionId>>;
    async fn list_children(&self, parent_id: &SessionId) -> SessionResult<Vec<SessionId>> {
        let ids = self.list(None).await?;
        let mut children = Vec::new();
        for id in ids {
            let Some(session) = self.load(&id).await? else {
                continue;
            };
            if session.parent_id == Some(*parent_id) {
                children.push(id);
            }
        }
        Ok(children)
    }

    // Queue
    async fn enqueue(
        &self,
        session_id: &SessionId,
        content: String,
        priority: i32,
    ) -> SessionResult<QueueItem>;
    async fn dequeue(&self, session_id: &SessionId) -> SessionResult<Option<QueueItem>>;
    async fn cancel_queued(&self, item_id: Uuid) -> SessionResult<bool>;
    async fn pending_queue(&self, session_id: &SessionId) -> SessionResult<Vec<QueueItem>>;
    async fn replace_pending_queue(
        &self,
        session_id: &SessionId,
        items: &[QueueItem],
    ) -> SessionResult<()>;

    /// Restore a session snapshot plus its pending queue as a single backend operation.
    ///
    /// Implementations should fail closed if the session already exists and should avoid
    /// exposing partially restored or unverified state to concurrent readers.
    ///
    /// Implementations are expected to validate the persisted round-trip before making the
    /// restored session visible. No default implementation is provided because these guarantees
    /// are backend-specific.
    async fn restore_bundle(
        &self,
        session: &Session,
        pending_queue: &[QueueItem],
    ) -> SessionResult<()>;

    // Cleanup
    async fn cleanup_expired(&self) -> SessionResult<usize>;

    /// Apply a synchronous mutation to a session under an advisory lock.
    ///
    /// The default implementation performs an **unlocked** load-modify-save cycle,
    /// which is NOT safe under concurrent access. Backends that support concurrent
    /// sessions **must** override this with a proper locking mechanism.
    async fn with_session_lock(&self, id: &SessionId, f: SessionMutationFn) -> SessionResult<()> {
        tracing::warn!(
            session_id = %id,
            backend = std::any::type_name::<Self>(),
            "with_session_lock: using default unlocked load-modify-save — \
             override this method for concurrent safety"
        );
        let mut session = self
            .load(id)
            .await?
            .ok_or_else(|| SessionError::NotFound { id: id.to_string() })?;
        f(&mut session)?;
        self.save(&session).await
    }

    async fn append_graph_event(
        &self,
        session_id: &SessionId,
        event: GraphEvent,
    ) -> SessionResult<()> {
        self.with_session_lock(
            session_id,
            Box::new(move |session| {
                let graph_id = session.graph.id;
                let created_at = session.graph.created_at;
                let primary_branch = session.graph.primary_branch;
                session.graph.events.push(event);
                session.graph = GraphMaterializer::from_events_with_primary(
                    &session.graph.events,
                    Some(primary_branch),
                );
                session.graph.id = graph_id;
                session.graph.created_at = created_at;
                Ok(())
            }),
        )
        .await
    }

    async fn load_graph(&self, session_id: &SessionId) -> SessionResult<Option<SessionGraph>> {
        Ok(self
            .load(session_id)
            .await?
            .map(|session| session.to_graph()))
    }

    /// Append a message to an existing session.
    ///
    /// Concurrency contract: implementations may hold a write lock for the duration
    /// of this call. Callers must not hold other persistence locks to avoid deadlocks.
    /// The default implementation performs a load-modify-save cycle; backends should
    /// override this with a more efficient single-lock approach when possible.
    async fn add_message(
        &self,
        session_id: &SessionId,
        message: SessionMessage,
    ) -> SessionResult<()> {
        self.with_session_lock(
            session_id,
            Box::new(move |session| {
                session.add_message(message)?;
                Ok(())
            }),
        )
        .await
    }

    async fn set_state(&self, session_id: &SessionId, state: SessionState) -> SessionResult<()> {
        self.with_session_lock(
            session_id,
            Box::new(move |session| {
                session.set_state(state);
                Ok(())
            }),
        )
        .await
    }

    /// Search sessions matching a filter.
    ///
    /// The default implementation loads all sessions and filters in memory.
    /// Backends with native query support (e.g., PostgreSQL) should override
    /// this with a more efficient implementation.
    async fn search(&self, filter: &SessionFilter) -> SessionResult<Vec<SessionId>> {
        let all_ids = self.list(filter.tenant_id.as_deref()).await?;
        let mut results = Vec::new();

        for id in all_ids {
            if let Some(limit) = filter.limit
                && results.len() >= limit
            {
                break;
            }
            if let Some(session) = self.load(&id).await?
                && filter.matches(&session)
            {
                results.push(id);
            }
        }

        Ok(results)
    }
}

#[derive(Debug, Default)]
pub struct MemoryPersistence {
    sessions: Arc<RwLock<HashMap<String, Session>>>,
    queue: Arc<RwLock<HashMap<String, Vec<QueueItem>>>>,
}

impl MemoryPersistence {
    pub fn new() -> Self {
        Self::default()
    }

    pub async fn count(&self) -> usize {
        self.sessions.read().await.len()
    }

    pub async fn clear(&self) {
        self.sessions.write().await.clear();
        self.queue.write().await.clear();
    }

    pub async fn load_graph(&self, session_id: &SessionId) -> SessionResult<Option<SessionGraph>> {
        <Self as Persistence>::load_graph(self, session_id).await
    }
}

#[async_trait::async_trait]
impl Persistence for MemoryPersistence {
    fn name(&self) -> &str {
        "memory"
    }

    async fn with_session_lock(&self, id: &SessionId, f: SessionMutationFn) -> SessionResult<()> {
        let mut sessions = self.sessions.write().await;
        let session = sessions
            .get_mut(&id.to_string())
            .ok_or_else(|| SessionError::NotFound { id: id.to_string() })?;
        f(session)?;
        validate_session_graph(session, "memory")?;
        Ok(())
    }

    async fn save(&self, session: &Session) -> SessionResult<()> {
        validate_session_graph(session, "memory")?;
        self.sessions
            .write()
            .await
            .insert(session.id.to_string(), session.clone());
        Ok(())
    }

    async fn add_message(
        &self,
        session_id: &SessionId,
        message: SessionMessage,
    ) -> SessionResult<()> {
        let mut sessions = self.sessions.write().await;
        if let Some(session) = sessions.get_mut(&session_id.to_string()) {
            session.add_message(message)?;
            Ok(())
        } else {
            Err(SessionError::NotFound {
                id: session_id.to_string(),
            })
        }
    }

    async fn load(&self, id: &SessionId) -> SessionResult<Option<Session>> {
        Ok(self.sessions.read().await.get(&id.to_string()).cloned())
    }

    async fn delete(&self, id: &SessionId) -> SessionResult<bool> {
        let key = id.to_string();
        let mut sessions = self.sessions.write().await;
        let mut queue = self.queue.write().await;
        queue.remove(&key);
        Ok(sessions.remove(&key).is_some())
    }

    async fn list(&self, tenant_id: Option<&str>) -> SessionResult<Vec<SessionId>> {
        Ok(self
            .sessions
            .read()
            .await
            .values()
            .filter(|s| {
                tenant_id
                    .map(|t| s.tenant_id.as_deref() == Some(t))
                    .unwrap_or(true)
            })
            .map(|s| s.id)
            .collect())
    }

    async fn list_children(&self, parent_id: &SessionId) -> SessionResult<Vec<SessionId>> {
        Ok(self
            .sessions
            .read()
            .await
            .values()
            .filter(|session| session.parent_id == Some(*parent_id))
            .map(|session| session.id)
            .collect())
    }

    async fn enqueue(
        &self,
        session_id: &SessionId,
        content: String,
        priority: i32,
    ) -> SessionResult<QueueItem> {
        let item = QueueItem::enqueue(*session_id, content).priority(priority);
        self.queue
            .write()
            .await
            .entry(session_id.to_string())
            .or_default()
            .push(item.clone());
        Ok(item)
    }

    async fn dequeue(&self, session_id: &SessionId) -> SessionResult<Option<QueueItem>> {
        let mut queue = self.queue.write().await;
        if let Some(items) = queue.get_mut(&session_id.to_string()) {
            items.sort_by(|a, b| b.priority.cmp(&a.priority));
            if let Some(pos) = items
                .iter()
                .position(|i| i.status == super::types::QueueStatus::Pending)
            {
                items[pos].start_processing();
                return Ok(Some(items[pos].clone()));
            }
        }
        Ok(None)
    }

    async fn cancel_queued(&self, item_id: Uuid) -> SessionResult<bool> {
        for items in self.queue.write().await.values_mut() {
            if let Some(item) = items.iter_mut().find(|i| i.id == item_id) {
                item.cancel();
                return Ok(true);
            }
        }
        Ok(false)
    }

    async fn pending_queue(&self, session_id: &SessionId) -> SessionResult<Vec<QueueItem>> {
        Ok(self
            .queue
            .read()
            .await
            .get(&session_id.to_string())
            .map(|items| {
                items
                    .iter()
                    .filter(|i| i.status == super::types::QueueStatus::Pending)
                    .cloned()
                    .collect()
            })
            .unwrap_or_default())
    }

    async fn replace_pending_queue(
        &self,
        session_id: &SessionId,
        items: &[QueueItem],
    ) -> SessionResult<()> {
        let normalized: Vec<QueueItem> = items
            .iter()
            .cloned()
            .map(|mut item| {
                item.session_id = *session_id;
                item.status = super::types::QueueStatus::Pending;
                item.processed_at = None;
                item
            })
            .collect();
        self.queue
            .write()
            .await
            .insert(session_id.to_string(), normalized);
        Ok(())
    }

    async fn restore_bundle(
        &self,
        session: &Session,
        pending_queue: &[QueueItem],
    ) -> SessionResult<()> {
        let normalized_queue: Vec<QueueItem> = pending_queue
            .iter()
            .cloned()
            .map(|mut item| {
                item.session_id = session.id;
                item.status = super::types::QueueStatus::Pending;
                item.processed_at = None;
                item
            })
            .collect();
        verify_restored_session_roundtrip(session, &normalized_queue, session, &normalized_queue)?;

        let mut sessions = self.sessions.write().await;
        let mut queue = self.queue.write().await;

        if sessions.contains_key(&session.id.to_string()) {
            return Err(SessionError::Storage {
                message: format!(
                    "Archive restore refuses to overwrite existing session {}",
                    session.id
                ),
            });
        }

        sessions.insert(session.id.to_string(), session.clone());
        queue.insert(session.id.to_string(), normalized_queue);

        Ok(())
    }

    async fn cleanup_expired(&self) -> SessionResult<usize> {
        // Hold both write locks simultaneously to prevent races where a
        // concurrent operation could observe a session removed from `sessions`
        // but still present in `queue`.
        let mut sessions = self.sessions.write().await;
        let mut queue = self.queue.write().await;

        let expired_keys: Vec<String> = sessions
            .iter()
            .filter(|(_, s)| s.is_expired())
            .map(|(k, _)| k.clone())
            .collect();

        for key in &expired_keys {
            sessions.remove(key);
            queue.remove(key);
        }

        Ok(expired_keys.len())
    }
}

pub struct PersistenceFactory;

impl PersistenceFactory {
    pub fn memory() -> Arc<dyn Persistence> {
        Arc::new(MemoryPersistence::new())
    }

    /// Create a JSONL persistence backend (requires `jsonl` feature).
    #[cfg(feature = "jsonl")]
    pub async fn jsonl(
        config: super::persistence_jsonl::JsonlConfig,
    ) -> SessionResult<Arc<dyn Persistence>> {
        Ok(Arc::new(
            super::persistence_jsonl::JsonlPersistence::new(config).await?,
        ))
    }

    /// Create a JSONL persistence backend with default configuration (requires `jsonl` feature).
    #[cfg(feature = "jsonl")]
    pub async fn jsonl_default() -> SessionResult<Arc<dyn Persistence>> {
        Ok(Arc::new(
            super::persistence_jsonl::JsonlPersistence::default_config().await?,
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::graph::{GraphEvent, GraphEventBody, NodeKind};
    use crate::session::state::SessionConfig;
    use crate::types::ContentBlock;

    #[tokio::test]
    async fn test_save_load() {
        let persistence = MemoryPersistence::new();
        let session = Session::new(SessionConfig::default());
        let id = session.id;

        persistence.save(&session).await.unwrap();
        let loaded = persistence.load(&id).await.unwrap();

        assert!(loaded.is_some());
        assert_eq!(loaded.unwrap().id, id);
    }

    #[tokio::test]
    async fn test_delete() {
        let persistence = MemoryPersistence::new();
        let session = Session::new(SessionConfig::default());
        let id = session.id;

        persistence.save(&session).await.unwrap();
        assert!(persistence.delete(&id).await.unwrap());
        assert!(persistence.load(&id).await.unwrap().is_none());
    }

    #[tokio::test]
    async fn test_list_by_tenant() {
        let persistence = MemoryPersistence::new();

        let mut s1 = Session::new(SessionConfig::default());
        s1.tenant_id = Some("tenant-a".to_string());
        s1.principal_id = Some("user-1".to_string());

        let mut s2 = Session::new(SessionConfig::default());
        s2.tenant_id = Some("tenant-b".to_string());
        s2.principal_id = Some("user-2".to_string());

        persistence.save(&s1).await.unwrap();
        persistence.save(&s2).await.unwrap();

        assert_eq!(persistence.list(None).await.unwrap().len(), 2);
        assert_eq!(persistence.list(Some("tenant-a")).await.unwrap().len(), 1);
    }

    #[tokio::test]
    async fn test_add_message() {
        let persistence = MemoryPersistence::new();
        let session = Session::new(SessionConfig::default());
        let id = session.id;

        persistence.save(&session).await.unwrap();
        persistence
            .add_message(&id, SessionMessage::user(vec![ContentBlock::text("Hello")]))
            .await
            .unwrap();

        let loaded = persistence.load(&id).await.unwrap().unwrap();
        assert_eq!(loaded.current_branch_messages().len(), 1);
    }

    #[tokio::test]
    async fn test_queue_priority() {
        let persistence = MemoryPersistence::new();
        let session = Session::new(SessionConfig::default());
        let id = session.id;

        persistence.save(&session).await.unwrap();
        persistence
            .enqueue(&id, "Low".to_string(), 1)
            .await
            .unwrap();
        persistence
            .enqueue(&id, "High".to_string(), 10)
            .await
            .unwrap();

        let next = persistence.dequeue(&id).await.unwrap().unwrap();
        assert_eq!(next.content, "High");
    }

    #[tokio::test]
    async fn test_cleanup_expired() {
        let persistence = MemoryPersistence::new();
        let config = SessionConfig {
            ttl_secs: Some(0),
            ..Default::default()
        };
        let session = Session::new(config);

        persistence.save(&session).await.unwrap();
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;

        assert_eq!(persistence.cleanup_expired().await.unwrap(), 1);
        assert_eq!(persistence.count().await, 0);
    }

    #[tokio::test]
    async fn test_append_graph_event() {
        let persistence = MemoryPersistence::new();
        let session = Session::new(SessionConfig::default());
        let id = session.id;

        persistence.save(&session).await.unwrap();
        persistence
            .append_graph_event(
                &id,
                GraphEvent::new(GraphEventBody::NodeAppended {
                    node_id: uuid::Uuid::new_v4(),
                    branch_id: session.graph.primary_branch,
                    parent_id: None,
                    kind: NodeKind::User,
                    tags: vec!["test".to_string()],
                    payload: serde_json::json!({"text": "hello"}),
                    provenance: None,
                }),
            )
            .await
            .unwrap();

        let graph = persistence.load_graph(&id).await.unwrap().unwrap();
        assert_eq!(graph.events.len(), 1);
        assert_eq!(graph.branch_nodes(graph.primary_branch).len(), 1);
    }

    #[tokio::test]
    async fn test_graph_first_roundtrip_restores_projection() {
        let persistence = MemoryPersistence::new();
        let mut session = Session::new(SessionConfig::default());
        session
            .add_message(SessionMessage::user(vec![ContentBlock::text("Hello")]))
            .unwrap();
        session.clear_messages();

        persistence.save(&session).await.unwrap();

        let loaded = persistence.load(&session.id).await.unwrap().unwrap();
        assert_eq!(loaded.current_branch_messages().len(), 1);
    }

    #[tokio::test]
    async fn test_graph_first_roundtrip_preserves_bookmarks_and_checkpoints() {
        let persistence = MemoryPersistence::new();
        let mut session = Session::new(SessionConfig::default());
        session
            .add_message(SessionMessage::user(vec![ContentBlock::text("Hello")]))
            .unwrap();
        session
            .graph
            .create_checkpoint(
                session.graph.primary_branch,
                "milestone",
                Some("saved".to_string()),
                vec!["tag".to_string()],
                None,
                None,
            )
            .unwrap();
        session.bookmark_current_head("head", Some("note".to_string()));
        session.clear_messages();

        persistence.save(&session).await.unwrap();

        let loaded = persistence.load(&session.id).await.unwrap().unwrap();
        assert_eq!(loaded.current_branch_messages().len(), 1);
        assert_eq!(loaded.graph.checkpoints.len(), 1);
        assert_eq!(loaded.graph.bookmarks.len(), 1);
    }
}
