//! Session registry - manages session actors.
//!
//! Lives on the server layer. Not an actor — just a plain data structure
//! protected by a mutex (acceptable: only accessed for routing, not during execution).

use crate::agent::agent_config::AgentConfig;
use crate::agent::core::{AgentMode, SessionRuntime};
use crate::agent::remote::SessionActorRef;
#[cfg(feature = "remote")]
use crate::agent::remote::event_relay::RemoteSessionDisconnect;
#[cfg(feature = "remote")]
use crate::agent::remote::scope::{MeshScopeId, scoped_event_relay, scoped_session};
use crate::agent::session_actor::SessionActor;
use crate::error::AgentError;
use crate::profiles::ProfileMetadata;
use agent_client_protocol::schema::{
    Error, ListSessionsRequest, ListSessionsResponse, McpServer, SessionConfigOption,
    SessionConfigOptionCategory, SessionConfigSelectOption, SessionInfo, SessionMode,
    SessionModeState,
};
use kameo::actor::{ActorRef, Spawn};
#[cfg(feature = "remote")]
use querymt_remote::MeshRuntimeHandle;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
#[cfg(feature = "remote")]
use tokio::sync::mpsc;

#[cfg(feature = "remote")]
pub(crate) fn select_relay_scope(
    active_scopes: &[MeshScopeId],
    preferred_scope: Option<MeshScopeId>,
) -> MeshScopeId {
    preferred_scope
        .filter(|scope| active_scopes.contains(scope))
        .or_else(|| active_scopes.first().cloned())
        .unwrap_or_else(MeshScopeId::lan_default)
}

pub fn all_session_modes() -> Vec<SessionMode> {
    vec![
        SessionMode::new("build", "Build").description("Full read/write mode"),
        SessionMode::new("plan", "Plan").description("Read-only planning mode"),
        SessionMode::new("review", "Review").description("Read-only review mode"),
    ]
}

pub fn mode_state(mode: AgentMode) -> SessionModeState {
    SessionModeState::new(mode.as_str(), all_session_modes())
}

/// Build the full set of session configuration options for the given mode and reasoning effort.
///
/// This is the single source of truth for config option shape — used by session creation,
/// set_session_config_option responses, and config_option_update notifications.
pub fn config_options(
    mode: AgentMode,
    reasoning_effort: Option<querymt::chat::ReasoningEffort>,
) -> Vec<SessionConfigOption> {
    config_options_with_profiles(mode, reasoning_effort, None, &[])
}

pub fn config_options_with_profiles(
    mode: AgentMode,
    reasoning_effort: Option<querymt::chat::ReasoningEffort>,
    current_profile_id: Option<&str>,
    profiles: &[ProfileMetadata],
) -> Vec<SessionConfigOption> {
    let effort_value = reasoning_effort
        .map(|e| e.to_string())
        .unwrap_or_else(|| "auto".to_string());
    let mut options = Vec::new();

    if !profiles.is_empty() {
        let profile_value = current_profile_id
            .map(str::to_string)
            .unwrap_or_else(|| profiles[0].id.clone());
        let profile_options: Vec<SessionConfigSelectOption> = profiles
            .iter()
            .map(|profile| {
                let mut option =
                    SessionConfigSelectOption::new(profile.id.clone(), profile.name.clone());
                if let Some(description) = profile.description.as_deref() {
                    option = option.description(description);
                }
                option
            })
            .collect();
        options.push(
            SessionConfigOption::select("profile", "Profile", profile_value, profile_options)
                .description("Controls the active profile for this session"),
        );
    }

    options.extend([
        SessionConfigOption::select(
            "mode",
            "Session Mode",
            mode.as_str(),
            vec![
                SessionConfigSelectOption::new("build", "Build")
                    .description("Full read/write mode"),
                SessionConfigSelectOption::new("plan", "Plan")
                    .description("Read-only planning mode"),
                SessionConfigSelectOption::new("review", "Review")
                    .description("Read-only review mode"),
            ],
        )
        .description("Controls how the agent operates for this session")
        .category(SessionConfigOptionCategory::Mode),
        SessionConfigOption::select(
            "reasoning_effort",
            "Reasoning Effort",
            effort_value,
            vec![
                SessionConfigSelectOption::new("auto", "Auto")
                    .description("Use model-specific defaults"),
                SessionConfigSelectOption::new("low", "Low")
                    .description("Minimal thinking, fastest responses"),
                SessionConfigSelectOption::new("medium", "Medium").description("Balanced thinking"),
                SessionConfigSelectOption::new("high", "High").description("Thorough thinking"),
                SessionConfigSelectOption::new("max", "Max")
                    .description("Deepest thinking, highest budget"),
            ],
        )
        .description("Controls reasoning depth for this session")
        .category(SessionConfigOptionCategory::ThoughtLevel),
    ]);

    options
}

/// Controls how a materialized session actor is integrated with remote infra.
pub struct SessionMaterializationOptions {
    pub attach_mesh_handle: bool,
    pub register_in_dht: bool,
}

/// Manages session actors. Lives on the server layer.
pub struct SessionMaterialization {
    pub actor_ref: ActorRef<SessionActor>,
    pub runtime: Arc<SessionRuntime>,
}

#[cfg(feature = "remote")]
#[derive(Clone)]
struct RemoteRelayRegistration {
    relay_actor_id: u64,
    relay_dht_name: String,
    relay_ref: ActorRef<crate::agent::remote::EventRelayActor>,
    linked: bool,
}

#[cfg(feature = "remote")]
#[derive(Clone, Copy)]
struct RemoteDetachOptions {
    preserve_bookmark: bool,
    notify_remote: bool,
}

pub struct SessionRegistry {
    pub config: Arc<AgentConfig>,
    sessions: HashMap<String, SessionActorRef>,
    local_actor_refs: HashMap<String, ActorRef<SessionActor>>,
    /// Tracks the per-session relay actor so attach/detach can manage
    /// event forwarding and remote actor links together.
    #[cfg(feature = "remote")]
    relay_actor_ids: HashMap<String, RemoteRelayRegistration>,
    /// Mesh handle for cleaning up re-registration closures when sessions
    /// are removed (Phase 4 of Bug 1 fix).
    #[cfg(feature = "remote")]
    mesh: Option<crate::agent::remote::MeshHandle>,
    #[cfg(feature = "remote")]
    remote_disconnect_tx: Option<mpsc::UnboundedSender<RemoteSessionDisconnect>>,
    /// Client bridge for workspace queries and notifications. Set once the ACP
    /// connection is established. Propagated to new session actors via `SetBridge`
    /// so that tools like `language_query` can access the client's language server.
    bridge: Option<crate::acp::client_bridge::ClientBridgeSender>,
}

impl SessionRegistry {
    pub fn new(config: Arc<AgentConfig>) -> Self {
        Self {
            config,
            sessions: HashMap::new(),
            local_actor_refs: HashMap::new(),
            #[cfg(feature = "remote")]
            relay_actor_ids: HashMap::new(),
            #[cfg(feature = "remote")]
            mesh: None,
            #[cfg(feature = "remote")]
            remote_disconnect_tx: None,
            bridge: None,
        }
    }

    /// Set the client bridge for ACP communication.
    ///
    /// When set, newly created sessions will receive the bridge via `SetBridge`
    /// so that session actors can send notifications and tools like
    /// `language_query` can access the client's language server.
    pub fn set_bridge(&mut self, bridge: crate::acp::client_bridge::ClientBridgeSender) {
        self.bridge = Some(bridge);
    }

    #[cfg(feature = "remote")]
    pub(crate) fn set_remote_disconnect_tx(
        &mut self,
        tx: mpsc::UnboundedSender<RemoteSessionDisconnect>,
    ) {
        self.remote_disconnect_tx = Some(tx);
    }

    /// Set the mesh handle so that `remove()` and `detach_remote_session()`
    /// can deregister actors from the re-registration map.
    ///
    /// When a mesh handle is provided, all existing **local** sessions in the
    /// registry are registered in the DHT so that remote peers can discover
    /// and attach to them.
    #[cfg(feature = "remote")]
    pub fn set_mesh(&mut self, mesh: Option<crate::agent::remote::MeshHandle>) {
        // Register all existing local sessions in each active scope so remote peers can attach.
        if let Some(ref mesh) = mesh {
            let runtime = MeshRuntimeHandle::from(mesh.clone());
            let scopes = runtime.active_scopes();
            for (session_id, actor_ref) in &self.local_actor_refs {
                for scope in &scopes {
                    let dht_name = scoped_session(scope, session_id);
                    let runtime = runtime.clone();
                    let actor_ref = actor_ref.clone();
                    tokio::spawn(async move {
                        runtime.register_actor(actor_ref, dht_name).await;
                    });
                }
            }
        }
        self.mesh = mesh;
    }

    /// Merge MCP servers from the agent config with any client-supplied servers.
    ///
    /// Config servers act as defaults; client-supplied servers with the same name
    /// take precedence (client wins).
    fn merged_mcp_servers(&self, req_servers: &[McpServer]) -> Vec<McpServer> {
        // Start with config servers converted to ACP format.
        let mut merged: Vec<McpServer> =
            self.config.mcp_servers.iter().map(|s| s.to_acp()).collect();

        // For each client-supplied server, replace any config server with the same
        // name or append it if not present.
        for req_server in req_servers {
            let req_name = match req_server {
                McpServer::Stdio(s) => s.name.as_str(),
                McpServer::Http(s) => s.name.as_str(),
                _ => continue,
            };
            if let Some(pos) = merged.iter().position(|s| match s {
                McpServer::Stdio(cs) => cs.name == req_name,
                McpServer::Http(cs) => cs.name == req_name,
                _ => false,
            }) {
                merged[pos] = req_server.clone();
            } else {
                merged.push(req_server.clone());
            }
        }

        merged
    }

    /// Get a reference to the session actor for routing.
    pub fn get(&self, session_id: &str) -> Option<&SessionActorRef> {
        self.sessions.get(session_id)
    }

    /// Insert a pre-spawned session actor into the registry.
    ///
    /// Accepts anything that converts into a `SessionActorRef`, including
    /// a bare `ActorRef<SessionActor>` (via the `From` impl).
    pub fn insert(&mut self, session_id: String, actor_ref: impl Into<SessionActorRef>) {
        self.sessions.insert(session_id.clone(), actor_ref.into());
        self.local_actor_refs.remove(&session_id);
    }

    /// Register a prepared session actor into the registry (fast, map-only operation).
    ///
    /// This is the "Register" phase of the 3-phase materialization pattern.
    /// It performs ONLY in-memory HashMap operations and should complete in microseconds.
    ///
    /// This does NOT:
    /// - Query the database
    /// - Initialize MCP servers
    /// - Spawn actors (already done in Prepare phase)
    /// - Register with DHT (done in Finalize phase)
    /// - Emit events (done in Finalize phase)
    /// - Set bridge (done in Finalize phase to avoid holding lock during await)
    ///
    /// # Arguments
    ///
    /// * `prepared` - The prepared session from `SessionMaterializer::prepare_*()` methods
    ///
    /// # Returns
    ///
    /// The `SessionActorRef` for routing to the session actor.
    pub async fn register_prepared_session(
        &mut self,
        prepared: &crate::agent::session_materializer::PreparedSession,
    ) -> SessionActorRef {
        let session_ref = SessionActorRef::from(prepared.actor_ref.clone());

        // Insert into in-memory maps (microseconds)
        // This is the ONLY operation that should happen under the lock.
        // Bridge setup is deferred to finalize_session to avoid holding
        // the lock during an async actor call.
        self.sessions
            .insert(prepared.session_id.clone(), session_ref.clone());
        self.local_actor_refs
            .insert(prepared.session_id.clone(), prepared.actor_ref.clone());

        session_ref
    }

    pub fn local_actor_ref(&self, session_id: &str) -> Option<&ActorRef<SessionActor>> {
        self.local_actor_refs.get(session_id)
    }

    pub async fn materialize_session_actor(
        &mut self,
        session_id: String,
        cwd: Option<PathBuf>,
        mcp_servers: &[McpServer],
        initialize_fork: bool,
        options: &mut SessionMaterializationOptions,
    ) -> Result<SessionMaterialization, Error> {
        let merged_mcp = self.merged_mcp_servers(mcp_servers);
        let tool_state = crate::agent::core::McpToolState::empty();
        let mcp_services = crate::agent::protocol::build_mcp_state(
            &merged_mcp,
            self.config.pending_elicitations(),
            self.config.event_sink.clone(),
            session_id.clone(),
            &crate::agent::mcp::agent_implementation(),
            tool_state.clone(),
        )
        .await?;

        if initialize_fork {
            crate::session::runtime::SessionForkHelper::initialize_fork(
                self.config.provider.history_store(),
                &session_id,
            )
            .await
            .map_err(|e| Error::internal_error().data(e.to_string()))?;
        }

        let runtime = SessionRuntime::new(cwd.clone(), mcp_services, tool_state);
        #[cfg(feature = "remote")]
        let actor = SessionActor::new(self.config.clone(), session_id.clone(), runtime.clone())
            .with_mesh(if options.attach_mesh_handle {
                self.mesh.clone()
            } else {
                None
            });
        #[cfg(not(feature = "remote"))]
        let actor = { SessionActor::new(self.config.clone(), session_id.clone(), runtime.clone()) };
        let actor_ref = SessionActor::spawn(actor);
        let session_ref = SessionActorRef::from(actor_ref.clone());

        if let Some(ref bridge) = self.bridge
            && let Err(e) = session_ref.set_bridge(bridge.clone()).await
        {
            log::warn!(
                "Session {}: failed to set bridge on session actor: {}",
                session_id,
                e
            );
        }

        self.sessions
            .insert(session_id.clone(), session_ref.clone());
        self.local_actor_refs
            .insert(session_id.clone(), actor_ref.clone());

        #[cfg(feature = "remote")]
        if options.register_in_dht
            && let Some(ref mesh) = self.mesh
        {
            let runtime = MeshRuntimeHandle::from(mesh.clone());
            for scope in runtime.active_scopes() {
                let dht_name = scoped_session(&scope, &session_id);
                let runtime = runtime.clone();
                let actor_ref = actor_ref.clone();
                tokio::spawn(async move {
                    runtime.register_actor(actor_ref, dht_name).await;
                });
            }
        }

        Ok(SessionMaterialization { actor_ref, runtime })
    }

    /// Remove a session actor from the registry.
    ///
    /// Also deregisters the session's re-registration closure from the mesh
    /// (if available) so dead actors don't accumulate (Phase 4 of Bug 1 fix).
    pub fn remove(&mut self, session_id: &str) -> Option<SessionActorRef> {
        #[cfg(feature = "remote")]
        if let Some(ref mesh) = self.mesh {
            let runtime = MeshRuntimeHandle::from(mesh.clone());
            for scope in runtime.active_scopes() {
                let session_dht_name = scoped_session(&scope, session_id);
                runtime.deregister_actor(&session_dht_name);
            }
        }
        self.local_actor_refs.remove(session_id);
        self.sessions.remove(session_id)
    }

    /// List all session IDs in the registry.
    pub fn session_ids(&self) -> Vec<String> {
        self.sessions.keys().cloned().collect()
    }

    /// Number of active sessions.
    pub fn len(&self) -> usize {
        self.sessions.len()
    }

    pub fn is_empty(&self) -> bool {
        self.sessions.is_empty()
    }

    /// Return `(session_id, peer_label, remote_node_id)` tuples for all remote sessions.
    ///
    /// Used by the UI session-list handler to include remote sessions in the
    /// session picker alongside local (persisted) sessions.
    #[cfg(feature = "remote")]
    pub fn remote_sessions(&self) -> Vec<(String, String, Option<String>)> {
        self.sessions
            .iter()
            .filter_map(|(id, r)| match r {
                SessionActorRef::Remote {
                    peer_label,
                    remote_node_id,
                    ..
                } => Some((id.clone(), peer_label.clone(), remote_node_id.clone())),
                SessionActorRef::Local(_) => None,
            })
            .collect()
    }

    /// Attach a remote session to this registry.
    ///
    /// Wraps the remote actor ref in a `SessionActorRef::Remote`, spawns a local
    /// `EventRelayActor`, registers it in the swarm (via `into_remote_ref`), and
    /// sends `SubscribeEvents` to the remote `SessionActor` so events stream back.
    ///
    /// # Event relay
    ///
    /// The `EventRelayActor` is spawned locally and its `ActorId` is sent to the
    /// remote session via `SubscribeEvents`. The remote handler constructs a
    /// `RemoteActorRef<EventRelayActor>` from that id (requires swarm to be
    /// bootstrapped — Phase 6). Until then the relay is spawned but the
    /// `SubscribeEvents` call returns Ok without installing the forwarder.
    ///
    /// # Returns
    ///
    /// The `SessionActorRef::Remote` for the attached session.
    #[cfg(feature = "remote")]
    pub async fn attach_remote_session(
        &mut self,
        session_id: String,
        remote_ref: kameo::actor::RemoteActorRef<SessionActor>,
        peer_label: String,
        mesh: Option<crate::agent::remote::MeshHandle>,
        preferred_scope: Option<MeshScopeId>,
        remote_node_id: Option<String>,
    ) -> SessionActorRef {
        log::debug!(
            "attach_remote_session: called for session_id={} peer='{}' \
             (registry currently has {} session(s))",
            session_id,
            peer_label,
            self.sessions.len(),
        );
        use crate::agent::remote::{EventRelayActor, SessionActorRef};
        use kameo::actor::Spawn;

        // 1. Wrap in SessionActorRef::Remote
        let session_ref = SessionActorRef::Remote {
            actor_ref: remote_ref.clone(),
            peer_label: peer_label.clone(),
            remote_node_id: remote_node_id.clone(),
        };

        // 2. Spawn a local EventRelayActor for this session.
        //    It persists durable remote events to the journal, publishes to fanout,
        //    and serves as the local endpoint for remote actor linking.
        let relay_actor = EventRelayActor::new(
            self.config.event_sink.clone(),
            session_id.clone(),
            peer_label.clone(),
            remote_node_id.clone(),
            self.remote_disconnect_tx.clone(),
        );
        let relay_ref = EventRelayActor::spawn(relay_actor);
        let relay_id = relay_ref.id().sequence_id();

        // 3. Register the relay in REMOTE_REGISTRY + DHT so the remote
        //    SessionActor can look it up by name and install an EventForwarder.
        //    Use peer-scoped name so multiple peers can attach to the same
        //    session without overwriting each other's relay (Bug 3 fix).
        let mesh_active = mesh.is_some();
        let relay_dht_name = if let Some(ref mesh) = mesh {
            let runtime = MeshRuntimeHandle::from(mesh.clone());
            let active_scopes = runtime.active_scopes();
            for scope in &active_scopes {
                let name = scoped_event_relay(scope, &session_id, mesh.peer_id());
                runtime.register_actor(relay_ref.clone(), name).await;
            }

            let selected_scope = select_relay_scope(&active_scopes, preferred_scope);
            scoped_event_relay(&selected_scope, &session_id, mesh.peer_id())
        } else {
            log::debug!(
                "attach_remote_session: no mesh, DHT registration skipped for relay (session {})",
                session_id
            );
            // Fallback name when there is no mesh (shouldn't happen in practice
            // but keeps the type system happy).
            format!("event_relay::{}::local", session_id)
        };

        // 4. Link the local relay to the remote session actor so we can observe
        //    peer disconnects and other remote actor lifecycle failures.
        let linked = match relay_ref.link_remote(&remote_ref).await {
            Ok(()) => {
                log::debug!(
                    "attach_remote_session: linked relay actor {} to remote session {}",
                    relay_id,
                    session_id
                );
                true
            }
            Err(e) => {
                log::warn!(
                    "attach_remote_session: failed to link relay actor {} to remote session {}: {}",
                    relay_id,
                    session_id,
                    e
                );
                false
            }
        };

        // 5. Send SubscribeEvents to the remote session.
        //    The remote SubscribeEvents handler uses mesh.lookup_actor to find
        //    the relay and install an EventForwarder on its EventBus.
        if let Err(e) = session_ref
            .subscribe_events(relay_id, relay_dht_name.clone())
            .await
        {
            log::warn!(
                "attach_remote_session: SubscribeEvents failed for {} (event relay may not be active): {}",
                session_id,
                e
            );
        }

        // 6. Insert into registry and track relay metadata for later cleanup.
        self.sessions
            .insert(session_id.clone(), session_ref.clone());
        self.relay_actor_ids.insert(
            session_id.clone(),
            RemoteRelayRegistration {
                relay_actor_id: relay_id,
                relay_dht_name: relay_dht_name.clone(),
                relay_ref: relay_ref.clone(),
                linked,
            },
        );

        // Persist a bookmark so this remote session survives server restart.
        if let Some(node_id) = remote_node_id {
            let title = self
                .config
                .provider
                .history_store()
                .get_session(&session_id)
                .await
                .ok()
                .flatten()
                .and_then(|session| session.name);
            let bookmark = crate::session::store::RemoteSessionBookmark {
                session_id: session_id.clone(),
                node_id,
                peer_label: peer_label.clone(),
                cwd: None,
                created_at: std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .map(|d| d.as_secs() as i64)
                    .unwrap_or(0),
                title,
            };
            let store = self.config.provider.history_store();
            tokio::spawn(async move {
                if let Err(e) = store.save_remote_session_bookmark(&bookmark).await {
                    log::warn!("Failed to persist remote session bookmark: {}", e);
                }
            });
        }

        log::info!(
            "Attached remote session {} from {} (relay_actor_id={}, event relay {})",
            session_id,
            peer_label,
            relay_id,
            if mesh_active {
                "active"
            } else {
                "pending mesh bootstrap"
            }
        );

        session_ref
    }

    #[cfg(feature = "remote")]
    async fn detach_remote_session_inner(
        &mut self,
        session_id: &str,
        options: RemoteDetachOptions,
    ) -> Option<SessionActorRef> {
        // Unlink and unsubscribe before removing so the remote forwarder is aborted
        // and normal local detach does not look like an unexpected disconnect.
        if options.notify_remote
            && let (Some(session_ref), Some(relay_reg)) = (
                self.sessions.get(session_id),
                self.relay_actor_ids.get(session_id).cloned(),
            )
            && session_ref.is_remote()
        {
            if let SessionActorRef::Remote { actor_ref, .. } = session_ref
                && relay_reg.linked
            {
                if let Err(e) = relay_reg.relay_ref.unlink_remote(actor_ref).await {
                    log::warn!(
                        "detach_remote_session: unlink_remote failed for {} (relay_actor_id={}): {}",
                        session_id,
                        relay_reg.relay_actor_id,
                        e
                    );
                } else {
                    log::debug!(
                        "detach_remote_session: unlinked relay actor {} from remote session {}",
                        relay_reg.relay_actor_id,
                        session_id
                    );
                }
            }

            if let Err(e) = session_ref
                .unsubscribe_events(relay_reg.relay_actor_id, relay_reg.relay_dht_name.clone())
                .await
            {
                log::warn!(
                    "detach_remote_session: UnsubscribeEvents failed for {} (relay_actor_id={}, relay_dht_name={}): {}",
                    session_id,
                    relay_reg.relay_actor_id,
                    relay_reg.relay_dht_name,
                    e
                );
            } else {
                log::info!(
                    "detach_remote_session: sent UnsubscribeEvents for {} (relay_actor_id={}, relay_dht_name={})",
                    session_id,
                    relay_reg.relay_actor_id,
                    relay_reg.relay_dht_name,
                );
            }
        }

        // Deregister the session and relay actors from the re-registration map
        // so dead closures don't accumulate (Phase 4 of Bug 1 fix).
        if let Some(ref mesh) = self.mesh {
            let runtime = MeshRuntimeHandle::from(mesh.clone());
            for scope in runtime.active_scopes() {
                let session_dht_name = scoped_session(&scope, session_id);
                runtime.deregister_actor(&session_dht_name);

                let relay_name = scoped_event_relay(&scope, session_id, mesh.peer_id());
                runtime.deregister_actor(&relay_name);
            }
        }

        if !options.preserve_bookmark {
            let store = self.config.provider.history_store();
            let sid = session_id.to_string();
            tokio::spawn(async move {
                if let Err(e) = store.remove_remote_session_bookmark(&sid).await {
                    log::warn!("Failed to remove remote session bookmark {}: {}", sid, e);
                }
            });
        }

        self.relay_actor_ids.remove(session_id);
        self.local_actor_refs.remove(session_id);
        self.sessions.remove(session_id)
    }

    /// Detach a remote session: send `UnsubscribeEvents` to stop the remote
    /// `EventForwarder`, then remove the session from the registry.
    ///
    /// This is the counterpart to [`attach_remote_session`](Self::attach_remote_session).
    /// Call this instead of bare `remove()` for remote sessions so the
    /// forwarder task on the remote node is properly cleaned up.
    ///
    /// For local sessions (or if the session is not in the registry) this
    /// falls back to a plain `remove()`.
    #[cfg(feature = "remote")]
    pub async fn detach_remote_session(&mut self, session_id: &str) -> Option<SessionActorRef> {
        self.detach_remote_session_inner(
            session_id,
            RemoteDetachOptions {
                preserve_bookmark: false,
                notify_remote: true,
            },
        )
        .await
    }

    /// Remove a remote session runtime from local tracking but keep its bookmark
    /// so the stopped session remains visible and resumable.
    #[cfg(feature = "remote")]
    pub async fn detach_remote_session_preserve_bookmark(
        &mut self,
        session_id: &str,
    ) -> Option<SessionActorRef> {
        self.detach_remote_session_inner(
            session_id,
            RemoteDetachOptions {
                preserve_bookmark: true,
                notify_remote: true,
            },
        )
        .await
    }

    #[cfg(feature = "remote")]
    pub(crate) async fn detach_remote_session_if_relay_matches(
        &mut self,
        session_id: &str,
        relay_actor_id: u64,
    ) -> Option<SessionActorRef> {
        let registration = self.relay_actor_ids.get(session_id)?;
        if registration.relay_actor_id != relay_actor_id {
            return None;
        }
        self.detach_remote_session_inner(
            session_id,
            RemoteDetachOptions {
                preserve_bookmark: true,
                notify_remote: false,
            },
        )
        .await
    }

    /// Fork an existing session at the latest message.
    pub async fn fork_session(
        &self,
        req: agent_client_protocol::schema::ForkSessionRequest,
    ) -> Result<agent_client_protocol::schema::ForkSessionResponse, Error> {
        let source_session_id = req.session_id.to_string();

        let _session = self
            .config
            .provider
            .history_store()
            .get_session(&source_session_id)
            .await
            .map_err(|e| Error::internal_error().data(e.to_string()))?
            .ok_or_else(|| {
                Error::invalid_params().data(serde_json::json!({
                    "message": "source session not found",
                    "session_id": source_session_id,
                }))
            })?;

        let history = self
            .config
            .provider
            .history_store()
            .get_history(&source_session_id)
            .await
            .map_err(|e| Error::internal_error().data(e.to_string()))?;

        let target_message_id = history
            .last()
            .map(|msg| msg.id.clone())
            .ok_or_else(|| Error::from(AgentError::EmptySessionFork))?;

        let new_session_id = self
            .config
            .provider
            .history_store()
            .fork_session(
                &source_session_id,
                &target_message_id,
                crate::session::domain::ForkOrigin::User,
            )
            .await
            .map_err(|e| Error::internal_error().data(e.to_string()))?;

        Ok(agent_client_protocol::schema::ForkSessionResponse::new(
            new_session_id,
        ))
    }

    /// Resume an existing session without history replay.
    pub async fn resume_session(
        &mut self,
        req: agent_client_protocol::schema::ResumeSessionRequest,
    ) -> Result<agent_client_protocol::schema::ResumeSessionResponse, Error> {
        let session_id = req.session_id.to_string();
        let _session = self
            .config
            .provider
            .history_store()
            .get_session(&session_id)
            .await
            .map_err(|e| Error::internal_error().data(e.to_string()))?
            .ok_or_else(|| {
                Error::invalid_params().data(serde_json::json!({
                    "message": "session not found",
                    "session_id": session_id,
                }))
            })?;

        let cwd = if req.cwd.as_os_str().is_empty() {
            None
        } else {
            if !req.cwd.is_absolute() {
                return Err(Error::invalid_params().data(serde_json::json!({
                    "message": "cwd must be an absolute path",
                    "cwd": req.cwd.display().to_string(),
                })));
            }
            Some(req.cwd.clone())
        };

        // If a local actor already exists for this session, reuse it instead
        // of spawning a duplicate.
        let already_materialized = self.local_actor_ref(&session_id).is_some();

        if !already_materialized {
            let _materialization = self
                .materialize_session_actor(
                    session_id.clone(),
                    cwd,
                    &req.mcp_servers,
                    false,
                    &mut SessionMaterializationOptions {
                        attach_mesh_handle: true,
                        register_in_dht: true,
                    },
                )
                .await?;

            self.config
                .emit_event(&session_id, crate::events::AgentEventKind::SessionCreated);
        }

        let current_mode = {
            let session_ref = self
                .sessions
                .get(&session_id)
                .ok_or_else(|| Error::internal_error().data("session actor missing after resume"))?
                .clone();
            session_ref.get_mode().await.map_err(Error::from)?
        };

        Ok(agent_client_protocol::schema::ResumeSessionResponse::new()
            .modes(mode_state(current_mode))
            .config_options(config_options(
                current_mode,
                **self.config.default_reasoning_effort.load(),
            )))
    }

    /// List all sessions (queries the store, not the actors).
    pub async fn list_sessions(
        &self,
        req: ListSessionsRequest,
    ) -> Result<ListSessionsResponse, Error> {
        let sessions = self
            .config
            .provider
            .history_store()
            .list_sessions()
            .await
            .map_err(|e| Error::internal_error().data(e.to_string()))?;

        let session_infos: Vec<SessionInfo> = sessions
            .into_iter()
            .map(|s| {
                let mut info = SessionInfo::new(
                    agent_client_protocol::schema::SessionId::from(s.public_id),
                    std::path::PathBuf::new(),
                );
                if let Some(name) = s.name {
                    info.title = Some(name);
                }
                if let Some(updated_at) = s.updated_at {
                    info.updated_at = Some(
                        updated_at
                            .format(&time::format_description::well_known::Rfc3339)
                            .unwrap_or_default(),
                    );
                }
                info
            })
            .collect();

        let filtered_infos = if let Some(_cwd) = req.cwd {
            // TODO: Filter by cwd once we store cwd per session
            session_infos
        } else {
            session_infos
        };

        let start_idx = req
            .cursor
            .as_ref()
            .and_then(|c| c.parse::<usize>().ok())
            .unwrap_or(0);
        let limit = 100;
        let end_idx = (start_idx + limit).min(filtered_infos.len());
        let paginated = filtered_infos[start_idx..end_idx].to_vec();
        let next_cursor = if end_idx < filtered_infos.len() {
            Some(end_idx.to_string())
        } else {
            None
        };

        Ok(ListSessionsResponse::new(paginated).next_cursor(next_cursor))
    }
}

// ══════════════════════════════════════════════════════════════════════════
//  Tests
// ══════════════════════════════════════════════════════════════════════════

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agent::agent_config_builder::AgentConfigBuilder;
    use crate::agent::core::ToolPolicy;
    use crate::agent::session_actor::SessionActor;
    use crate::session::backend::StorageBackend;
    use crate::session::store::SessionStore;
    use crate::test_utils::{
        MockLlmProvider, MockSessionStore, SharedLlmProvider, TestProviderFactory, mock_llm_config,
        mock_plugin_registry, mock_session,
    };
    use kameo::actor::Spawn;
    use querymt::LLMParams;
    use std::collections::HashMap;
    use std::sync::Arc;
    use tokio::sync::{Mutex, mpsc};

    // ── Fixture ──────────────────────────────────────────────────────────────

    struct RegistryFixture {
        registry: SessionRegistry,
        _temp_dir: tempfile::TempDir,
    }

    impl RegistryFixture {
        async fn new() -> Self {
            let provider = Arc::new(Mutex::new(MockLlmProvider::new()));
            let shared = SharedLlmProvider {
                inner: provider.clone(),
                tools: vec![].into_boxed_slice(),
            };
            let factory = Arc::new(TestProviderFactory { provider: shared });
            let (plugin_registry, temp_dir) =
                mock_plugin_registry(factory).expect("plugin registry");

            let mut store = MockSessionStore::new();
            let llm_config = mock_llm_config();
            let session = mock_session("test-session");
            store
                .expect_get_session()
                .returning(move |_| Ok(Some(session.clone())))
                .times(0..);
            store
                .expect_get_session_llm_config()
                .returning(move |_| Ok(Some(llm_config.clone())))
                .times(0..);
            store
                .expect_list_sessions()
                .returning(|| Ok(vec![]))
                .times(0..);

            let store: Arc<dyn SessionStore> = Arc::new(store);
            let storage = Arc::new(
                crate::session::sqlite_storage::SqliteStorage::connect(":memory:".into())
                    .await
                    .expect("create event store"),
            );

            let config = Arc::new(
                AgentConfigBuilder::new(
                    Arc::new(plugin_registry),
                    store.clone(),
                    storage.event_journal(),
                    LLMParams::new().provider("mock").model("mock-model"),
                )
                .with_tool_policy(ToolPolicy::ProviderOnly)
                .build(),
            );

            Self {
                registry: SessionRegistry::new(config),
                _temp_dir: temp_dir,
            }
        }

        fn spawn_actor(&self) -> kameo::actor::ActorRef<SessionActor> {
            let runtime = crate::agent::core::SessionRuntime::new(
                None,
                HashMap::new(),
                crate::agent::core::McpToolState::empty(),
            );
            let actor = SessionActor::new(
                self.registry.config.clone(),
                "test-session".to_string(),
                runtime,
            );
            SessionActor::spawn(actor)
        }
    }

    // ── Unit tests ───────────────────────────────────────────────────────────

    #[tokio::test]
    async fn test_new_registry_is_empty() {
        let f = RegistryFixture::new().await;
        assert!(f.registry.is_empty());
        assert_eq!(f.registry.len(), 0);
        assert!(f.registry.session_ids().is_empty());
    }

    #[test]
    fn test_mode_state_contains_all_supported_modes() {
        let state = mode_state(AgentMode::Build);
        assert_eq!(state.current_mode_id.0.as_ref(), "build");
        assert_eq!(state.available_modes.len(), 3);
    }

    #[test]
    fn test_config_options_include_mode_selector() {
        let options = config_options(AgentMode::Review, None);
        assert_eq!(options.len(), 2);
        assert_eq!(options[0].id.0.as_ref(), "mode");

        let select = match &options[0].kind {
            agent_client_protocol::schema::SessionConfigKind::Select(select) => select,
            _ => panic!("expected select mode option"),
        };
        assert_eq!(select.current_value.0.as_ref(), "review");

        // Verify reasoning effort option is present
        assert_eq!(options[1].id.0.as_ref(), "reasoning_effort");
        assert_eq!(
            options[1].category,
            Some(SessionConfigOptionCategory::ThoughtLevel)
        );
    }

    #[tokio::test]
    async fn test_get_nonexistent_session_returns_none() {
        let f = RegistryFixture::new().await;
        assert!(f.registry.get("no-such-session").is_none());
    }

    #[tokio::test]
    async fn test_insert_and_get_session() {
        let mut f = RegistryFixture::new().await;
        let actor_ref = f.spawn_actor();
        f.registry.insert("sess-1".to_string(), actor_ref);

        assert!(!f.registry.is_empty());
        assert_eq!(f.registry.len(), 1);
        assert!(f.registry.get("sess-1").is_some());
        assert!(f.registry.get("sess-2").is_none());
    }

    #[tokio::test]
    async fn test_insert_multiple_sessions() {
        let mut f = RegistryFixture::new().await;
        for i in 0..5 {
            let actor_ref = f.spawn_actor();
            f.registry.insert(format!("sess-{i}"), actor_ref);
        }
        assert_eq!(f.registry.len(), 5);
        let ids = f.registry.session_ids();
        assert_eq!(ids.len(), 5);
        for i in 0..5 {
            assert!(ids.contains(&format!("sess-{i}")));
        }
    }

    #[tokio::test]
    async fn test_remove_existing_session() {
        let mut f = RegistryFixture::new().await;
        let actor_ref = f.spawn_actor();
        f.registry.insert("sess-1".to_string(), actor_ref);
        assert_eq!(f.registry.len(), 1);

        let removed = f.registry.remove("sess-1");
        assert!(removed.is_some());
        assert!(f.registry.is_empty());
    }

    #[tokio::test]
    async fn test_remove_nonexistent_session_returns_none() {
        let mut f = RegistryFixture::new().await;
        let removed = f.registry.remove("no-such-session");
        assert!(removed.is_none());
        assert!(f.registry.is_empty());
    }

    #[tokio::test]
    async fn test_overwrite_existing_session_id() {
        let mut f = RegistryFixture::new().await;
        let a1 = f.spawn_actor();
        let a2 = f.spawn_actor();
        f.registry.insert("sess-1".to_string(), a1);
        f.registry.insert("sess-1".to_string(), a2);
        // Still one entry, not two
        assert_eq!(f.registry.len(), 1);
        assert!(f.registry.get("sess-1").is_some());
    }

    #[tokio::test]
    async fn test_session_ids_reflects_inserts_and_removes() {
        let mut f = RegistryFixture::new().await;
        let a1 = f.spawn_actor();
        let a2 = f.spawn_actor();
        f.registry.insert("alpha".to_string(), a1);
        f.registry.insert("beta".to_string(), a2);

        let ids = f.registry.session_ids();
        assert!(ids.contains(&"alpha".to_string()));
        assert!(ids.contains(&"beta".to_string()));

        f.registry.remove("alpha");
        let ids = f.registry.session_ids();
        assert!(!ids.contains(&"alpha".to_string()));
        assert!(ids.contains(&"beta".to_string()));
    }

    #[tokio::test]
    async fn test_list_sessions_empty_store() {
        let f = RegistryFixture::new().await;
        let req = ListSessionsRequest::new();
        let resp = f.registry.list_sessions(req).await.expect("list_sessions");
        assert!(resp.sessions.is_empty());
    }

    #[tokio::test]
    async fn test_fork_session_empty_history_fails() {
        let _f = RegistryFixture::new().await;

        // Override: get_history returns empty vec, get_session returns Some
        let mut store2 = MockSessionStore::new();
        let session = mock_session("source-session");
        store2
            .expect_get_session()
            .returning(move |_| Ok(Some(session.clone())))
            .times(0..);
        store2
            .expect_get_history()
            .returning(|_| Ok(vec![]))
            .times(0..);
        store2
            .expect_get_session_llm_config()
            .returning(|_| Ok(None))
            .times(0..);
        store2
            .expect_list_sessions()
            .returning(|| Ok(vec![]))
            .times(0..);

        let store2: Arc<dyn SessionStore> = Arc::new(store2);
        let provider = Arc::new(Mutex::new(MockLlmProvider::new()));
        let shared = SharedLlmProvider {
            inner: provider,
            tools: vec![].into_boxed_slice(),
        };
        let factory = Arc::new(TestProviderFactory { provider: shared });
        let (plugin_registry, _temp) = mock_plugin_registry(factory).expect("registry");
        let storage = Arc::new(
            crate::session::sqlite_storage::SqliteStorage::connect(":memory:".into())
                .await
                .expect("create event store"),
        );
        let config = Arc::new(
            AgentConfigBuilder::new(
                Arc::new(plugin_registry),
                store2.clone(),
                storage.event_journal(),
                LLMParams::new().provider("mock").model("mock-model"),
            )
            .build(),
        );
        let registry = SessionRegistry::new(config);

        let req = agent_client_protocol::schema::ForkSessionRequest::new(
            agent_client_protocol::schema::SessionId::from("source-session".to_string()),
            std::path::PathBuf::from("/tmp"),
        );
        let result = registry.fork_session(req).await;
        // Should fail with EmptySessionFork since history is empty
        assert!(result.is_err(), "expected error for empty session fork");
        let err = result.unwrap_err();
        assert!(
            err.to_string().contains("empty")
                || err.code == agent_client_protocol::ErrorCode::InternalError,
            "unexpected error: {}",
            err
        );
    }

    // ── merged_mcp_servers tests ─────────────────────────────────────────────

    #[tokio::test]
    async fn merged_mcp_servers_no_config_no_request() {
        let f = RegistryFixture::new().await;
        let merged = f.registry.merged_mcp_servers(&[]);
        assert!(merged.is_empty());
    }

    #[tokio::test]
    async fn merged_mcp_servers_config_only() {
        let mut f = RegistryFixture::new().await;
        // Inject an MCP server via config
        Arc::get_mut(&mut f.registry.config)
            .expect("single owner")
            .mcp_servers
            .push(crate::config::McpServerConfig::Http {
                name: "config-server".to_string(),
                url: "https://mcp.example.com/mcp".to_string(),
                headers: std::collections::HashMap::new(),
            });

        let merged = f.registry.merged_mcp_servers(&[]);
        assert_eq!(merged.len(), 1);
        assert!(matches!(&merged[0], McpServer::Http(s) if s.name == "config-server"));
    }

    #[tokio::test]
    async fn merged_mcp_servers_request_only() {
        let f = RegistryFixture::new().await;
        let req_server = agent_client_protocol::schema::McpServerHttp::new(
            "req-server".to_string(),
            "https://req.example.com/mcp".to_string(),
        );
        let req_servers = vec![McpServer::Http(req_server)];

        let merged = f.registry.merged_mcp_servers(&req_servers);
        assert_eq!(merged.len(), 1);
        assert!(matches!(&merged[0], McpServer::Http(s) if s.name == "req-server"));
    }

    #[tokio::test]
    async fn merged_mcp_servers_request_overrides_config_by_name() {
        let mut f = RegistryFixture::new().await;
        Arc::get_mut(&mut f.registry.config)
            .expect("single owner")
            .mcp_servers
            .push(crate::config::McpServerConfig::Http {
                name: "shared-name".to_string(),
                url: "https://config.example.com/mcp".to_string(),
                headers: std::collections::HashMap::new(),
            });

        // Request provides a different URL for the same name — should win.
        let req_server = agent_client_protocol::schema::McpServerHttp::new(
            "shared-name".to_string(),
            "https://override.example.com/mcp".to_string(),
        );
        let req_servers = vec![McpServer::Http(req_server)];

        let merged = f.registry.merged_mcp_servers(&req_servers);
        // Still one entry, not two
        assert_eq!(merged.len(), 1);
        assert!(
            matches!(&merged[0], McpServer::Http(s) if s.url == "https://override.example.com/mcp"),
            "expected the request server to override the config server"
        );
    }

    #[tokio::test]
    async fn merged_mcp_servers_both_different_names() {
        let mut f = RegistryFixture::new().await;
        Arc::get_mut(&mut f.registry.config)
            .expect("single owner")
            .mcp_servers
            .push(crate::config::McpServerConfig::Http {
                name: "config-server".to_string(),
                url: "https://config.example.com/mcp".to_string(),
                headers: std::collections::HashMap::new(),
            });

        let req_server = agent_client_protocol::schema::McpServerHttp::new(
            "req-server".to_string(),
            "https://req.example.com/mcp".to_string(),
        );
        let req_servers = vec![McpServer::Http(req_server)];

        let merged = f.registry.merged_mcp_servers(&req_servers);
        assert_eq!(merged.len(), 2);
        let names: Vec<&str> = merged
            .iter()
            .map(|s| match s {
                McpServer::Http(h) => h.name.as_str(),
                McpServer::Stdio(s) => s.name.as_str(),
                _ => "unknown",
            })
            .collect();
        assert!(names.contains(&"config-server"));
        assert!(names.contains(&"req-server"));
    }

    #[cfg(feature = "remote")]
    #[test]
    fn select_relay_scope_uses_preferred_scope_when_present() {
        let scopes = vec![
            MeshScopeId::lan_default(),
            MeshScopeId::Iroh {
                mesh_id: "team-a".to_string(),
            },
        ];
        let selected = select_relay_scope(
            &scopes,
            Some(MeshScopeId::Iroh {
                mesh_id: "team-a".to_string(),
            }),
        );
        assert_eq!(
            selected,
            MeshScopeId::Iroh {
                mesh_id: "team-a".to_string()
            }
        );
    }

    #[cfg(feature = "remote")]
    #[test]
    fn select_relay_scope_falls_back_to_first_scope_when_preferred_missing() {
        let scopes = vec![
            MeshScopeId::lan_default(),
            MeshScopeId::Iroh {
                mesh_id: "team-a".to_string(),
            },
        ];
        let selected = select_relay_scope(
            &scopes,
            Some(MeshScopeId::Iroh {
                mesh_id: "other".to_string(),
            }),
        );
        assert_eq!(selected, MeshScopeId::lan_default());
    }

    #[cfg(feature = "remote")]
    #[tokio::test]
    async fn detach_remote_session_if_relay_matches_ignores_stale_relay_id() {
        let mut f = RegistryFixture::new().await;
        let actor_ref = f.spawn_actor();
        f.registry.sessions.insert(
            "remote-session".to_string(),
            SessionActorRef::Local(actor_ref),
        );
        f.registry.relay_actor_ids.insert(
            "remote-session".to_string(),
            RemoteRelayRegistration {
                relay_actor_id: 42,
                relay_dht_name: "relay::remote-session".to_string(),
                relay_ref: crate::agent::remote::EventRelayActor::spawn(
                    crate::agent::remote::EventRelayActor::new(
                        f.registry.config.event_sink.clone(),
                        "remote-session".to_string(),
                        "peer".to_string(),
                        None,
                        None,
                    ),
                ),
                linked: false,
            },
        );

        let detached = f
            .registry
            .detach_remote_session_if_relay_matches("remote-session", 99)
            .await;

        assert!(detached.is_none());
        assert!(f.registry.sessions.contains_key("remote-session"));
        assert!(f.registry.relay_actor_ids.contains_key("remote-session"));
    }

    #[cfg(feature = "remote")]
    #[tokio::test]
    async fn detach_remote_session_if_relay_matches_detaches_current_registration() {
        let mut f = RegistryFixture::new().await;
        let actor_ref = f.spawn_actor();
        let session_id = "remote-session".to_string();
        f.registry
            .sessions
            .insert(session_id.clone(), SessionActorRef::Local(actor_ref));
        let relay_ref = crate::agent::remote::EventRelayActor::spawn(
            crate::agent::remote::EventRelayActor::new(
                f.registry.config.event_sink.clone(),
                session_id.clone(),
                "peer".to_string(),
                None,
                None,
            ),
        );
        let relay_actor_id = relay_ref.id().sequence_id();
        f.registry.relay_actor_ids.insert(
            session_id.clone(),
            RemoteRelayRegistration {
                relay_actor_id,
                relay_dht_name: "relay::remote-session".to_string(),
                relay_ref,
                linked: false,
            },
        );
        let (disconnect_tx, _disconnect_rx) = mpsc::unbounded_channel();
        f.registry.set_remote_disconnect_tx(disconnect_tx);

        let detached = f
            .registry
            .detach_remote_session_if_relay_matches(&session_id, relay_actor_id)
            .await;

        assert!(detached.is_some());
        assert!(!f.registry.sessions.contains_key(&session_id));
        assert!(!f.registry.relay_actor_ids.contains_key(&session_id));
    }
}
