PicoBot/src/gateway/reload.rs

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, RwLock};

use serde::Serialize;
use tokio::sync::{Notify, mpsc, oneshot};

use crate::config::Config;

const RELOAD_QUEUE_CAPACITY: usize = 8;

#[derive(Clone)]
pub(crate) struct RuntimeAdmission {
    inner: Arc<AdmissionInner>,
}

struct AdmissionInner {
    accepting: AtomicBool,
    active: AtomicUsize,
    idle: Notify,
}

pub(crate) struct ActivityGuard {
    admission: RuntimeAdmission,
}

impl RuntimeAdmission {
    pub fn open() -> Self {
        Self {
            inner: Arc::new(AdmissionInner {
                accepting: AtomicBool::new(true),
                active: AtomicUsize::new(0),
                idle: Notify::new(),
            }),
        }
    }

    pub fn try_enter(&self) -> Option<ActivityGuard> {
        if !self.inner.accepting.load(Ordering::Acquire) {
            return None;
        }
        self.inner.active.fetch_add(1, Ordering::AcqRel);
        if !self.inner.accepting.load(Ordering::Acquire) {
            self.leave();
            return None;
        }
        Some(ActivityGuard {
            admission: self.clone(),
        })
    }

    pub fn close(&self) {
        self.inner.accepting.store(false, Ordering::Release);
        if self.inner.active.load(Ordering::Acquire) == 0 {
            self.inner.idle.notify_waiters();
        }
    }

    pub fn is_accepting(&self) -> bool {
        self.inner.accepting.load(Ordering::Acquire)
    }

    pub async fn wait_for_idle(&self) {
        loop {
            let notified = self.inner.idle.notified();
            if self.inner.active.load(Ordering::Acquire) == 0 {
                return;
            }
            notified.await;
        }
    }

    fn leave(&self) {
        if self.inner.active.fetch_sub(1, Ordering::AcqRel) == 1 {
            self.inner.idle.notify_waiters();
        }
    }
}

impl Drop for ActivityGuard {
    fn drop(&mut self) {
        self.admission.leave();
    }
}

pub(crate) struct ReloadRequest {
    pub generation: u64,
    pub response: oneshot::Sender<Result<ReloadAccepted, ReloadError>>,
}

#[derive(Clone)]
pub struct ReloadHandle {
    sender: mpsc::Sender<ReloadRequest>,
    next_generation: Arc<AtomicU64>,
    pending: Arc<AtomicBool>,
    status: Arc<RwLock<ReloadStatus>>,
}

pub(crate) struct ReloadController {
    pub handle: ReloadHandle,
    pub receiver: mpsc::Receiver<ReloadRequest>,
    pub startup_process_env: HashMap<String, String>,
    pub startup_cwd: PathBuf,
    status: Arc<RwLock<ReloadStatus>>,
}

#[derive(Debug, Clone, Serialize)]
pub struct ReloadAccepted {
    pub generation: u64,
    pub message: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ReloadPhase {
    Active,
    Preparing,
    Draining,
    Activating,
    Failed,
}

#[derive(Debug, Clone, Serialize)]
pub struct ReloadStatus {
    pub generation: u64,
    pub phase: ReloadPhase,
    pub requested_at: Option<i64>,
    pub activated_at: Option<i64>,
    pub last_error: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ReloadError {
    AlreadyPending,
    ShuttingDown,
    InvalidConfig(String),
    ImmutableField(String),
    PreparationFailed(String),
}

impl std::fmt::Display for ReloadError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::AlreadyPending => {
                write!(formatter, "another configuration reload is already pending")
            }
            Self::ShuttingDown => write!(formatter, "gateway is shutting down"),
            Self::InvalidConfig(error)
            | Self::ImmutableField(error)
            | Self::PreparationFailed(error) => formatter.write_str(error),
        }
    }
}

impl std::error::Error for ReloadError {}

impl ReloadController {
    pub fn new(startup_process_env: HashMap<String, String>, startup_cwd: PathBuf) -> Self {
        let (sender, receiver) = mpsc::channel(RELOAD_QUEUE_CAPACITY);
        let status = Arc::new(RwLock::new(ReloadStatus {
            generation: 1,
            phase: ReloadPhase::Active,
            requested_at: None,
            activated_at: Some(chrono::Utc::now().timestamp_millis()),
            last_error: None,
        }));
        Self {
            handle: ReloadHandle {
                sender,
                next_generation: Arc::new(AtomicU64::new(2)),
                pending: Arc::new(AtomicBool::new(false)),
                status: status.clone(),
            },
            receiver,
            startup_process_env,
            startup_cwd,
            status,
        }
    }

    pub fn set_phase(&self, generation: u64, phase: ReloadPhase) {
        let mut status = self
            .status
            .write()
            .unwrap_or_else(|error| error.into_inner());
        status.generation = generation;
        status.phase = phase;
        if phase == ReloadPhase::Preparing {
            status.requested_at = Some(chrono::Utc::now().timestamp_millis());
            status.activated_at = None;
            status.last_error = None;
        }
        if phase == ReloadPhase::Active {
            status.activated_at = Some(chrono::Utc::now().timestamp_millis());
            self.handle.pending.store(false, Ordering::Release);
        }
    }

    pub fn set_failed(&self, generation: u64, error: impl Into<String>) {
        let mut status = self
            .status
            .write()
            .unwrap_or_else(|error| error.into_inner());
        status.generation = generation;
        status.phase = ReloadPhase::Failed;
        status.last_error = Some(error.into());
        self.handle.pending.store(false, Ordering::Release);
    }
}

impl ReloadHandle {
    pub(crate) fn unavailable() -> Self {
        let (sender, receiver) = mpsc::channel(1);
        drop(receiver);
        Self {
            sender,
            next_generation: Arc::new(AtomicU64::new(1)),
            pending: Arc::new(AtomicBool::new(false)),
            status: Arc::new(RwLock::new(ReloadStatus {
                generation: 0,
                phase: ReloadPhase::Failed,
                requested_at: None,
                activated_at: None,
                last_error: Some("reload controller is unavailable".to_string()),
            })),
        }
    }

    pub async fn request(&self) -> Result<ReloadAccepted, ReloadError> {
        self.pending
            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
            .map_err(|_| ReloadError::AlreadyPending)?;
        let generation = self.next_generation.fetch_add(1, Ordering::Relaxed);
        let (response, receiver) = oneshot::channel();
        if let Err(error) = self.sender.try_send(ReloadRequest {
            generation,
            response,
        }) {
            self.pending.store(false, Ordering::Release);
            return Err(match error {
                mpsc::error::TrySendError::Full(_) => ReloadError::AlreadyPending,
                mpsc::error::TrySendError::Closed(_) => ReloadError::ShuttingDown,
            });
        }
        match receiver.await {
            Ok(result) => result,
            Err(_) => {
                self.pending.store(false, Ordering::Release);
                Err(ReloadError::ShuttingDown)
            }
        }
    }

    pub fn status(&self) -> ReloadStatus {
        self.status
            .read()
            .unwrap_or_else(|error| error.into_inner())
            .clone()
    }
}

pub(crate) fn load_candidate(
    config_path: &Path,
    startup_process_env: &HashMap<String, String>,
    startup_cwd: &std::path::Path,
    current: &Config,
    current_workspace: &std::path::Path,
) -> Result<Config, ReloadError> {
    let mut candidate = Config::load_for_reload(config_path, startup_process_env, startup_cwd)
        .map_err(|error| {
            ReloadError::InvalidConfig(format!("configuration reload failed: {error}"))
        })?;
    candidate
        .get_provider_config("default")
        .map_err(|error| ReloadError::InvalidConfig(format!("invalid default agent: {error}")))?;
    if let Some(feishu) = candidate.channels.get("feishu")
        && feishu.enabled
        && (feishu.app_id.trim().is_empty() || feishu.app_secret.trim().is_empty())
    {
        return Err(ReloadError::InvalidConfig(
            "enabled channels.feishu requires non-empty app_id and app_secret".to_string(),
        ));
    }

    let mut candidate_workspace = crate::config::expand_path(&candidate.workspace_dir);
    if candidate_workspace.is_relative() {
        candidate_workspace = startup_cwd.join(candidate_workspace);
    }
    let candidate_workspace = candidate_workspace
        .canonicalize()
        .unwrap_or(candidate_workspace);
    if current_workspace != candidate_workspace {
        return Err(ReloadError::ImmutableField(
            "workspace_dir cannot be reloaded; restart the gateway".to_string(),
        ));
    }
    candidate.workspace_dir = current_workspace.to_string_lossy().to_string();
    if effective_db_path(current, current_workspace)
        != effective_db_path(&candidate, current_workspace)
    {
        return Err(ReloadError::ImmutableField(
            "gateway.session_db_path cannot be reloaded; restart the gateway".to_string(),
        ));
    }
    if current.gateway.host != candidate.gateway.host
        || current.gateway.port != candidate.gateway.port
    {
        return Err(ReloadError::ImmutableField(
            "gateway.host and gateway.port cannot be reloaded; restart the gateway".to_string(),
        ));
    }
    Ok(candidate)
}

fn effective_db_path(config: &Config, workspace: &Path) -> PathBuf {
    let path = config
        .gateway
        .session_db_path
        .as_deref()
        .map(crate::config::expand_path)
        .unwrap_or_else(|| workspace.join("picobot.db"));
    let path = if path.is_relative() {
        workspace.join(path)
    } else {
        path
    };
    path.canonicalize().unwrap_or(path)
}

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

    fn config_json(workspace: &std::path::Path, model_id: &str) -> String {
        serde_json::json!({
            "providers": {
                "provider": {
                    "type": "openai",
                    "base_url": "https://example.invalid/v1",
                    "api_key": "test"
                }
            },
            "models": { "model": { "model_id": model_id } },
            "agents": {
                "default": { "provider": "provider", "model": "model" }
            },
            "workspace_dir": workspace
        })
        .to_string()
    }

    #[test]
    fn candidate_accepts_runtime_changes_and_rejects_workspace_changes() {
        let temp = tempfile::tempdir().unwrap();
        let workspace = temp.path().join("workspace");
        std::fs::create_dir_all(&workspace).unwrap();
        let config_path = temp.path().join("config.json");
        let current: Config = serde_json::from_str(&config_json(&workspace, "old-model")).unwrap();
        std::fs::write(
            &config_path,
            config_json(std::path::Path::new("workspace"), "new-model"),
        )
        .unwrap();

        let candidate = load_candidate(
            &config_path,
            &HashMap::new(),
            temp.path(),
            &current,
            &workspace,
        )
        .unwrap();
        assert_eq!(
            candidate.get_provider_config("default").unwrap().model_id,
            "new-model"
        );

        let other_workspace = temp.path().join("other");
        std::fs::create_dir_all(&other_workspace).unwrap();
        std::fs::write(&config_path, config_json(&other_workspace, "new-model")).unwrap();
        let error = load_candidate(
            &config_path,
            &HashMap::new(),
            temp.path(),
            &current,
            &workspace,
        )
        .unwrap_err();
        assert!(
            error
                .to_string()
                .contains("workspace_dir cannot be reloaded")
        );
    }

    #[test]
    fn equivalent_default_database_paths_are_reloadable() {
        let temp = tempfile::tempdir().unwrap();
        let workspace = temp.path().join("workspace");
        std::fs::create_dir_all(&workspace).unwrap();
        std::fs::write(workspace.join("picobot.db"), []).unwrap();
        let config_path = temp.path().join("config.json");
        let current: Config = serde_json::from_str(&config_json(&workspace, "old-model")).unwrap();
        let mut candidate: serde_json::Value =
            serde_json::from_str(&config_json(&workspace, "new-model")).unwrap();
        candidate["gateway"] = serde_json::json!({ "session_db_path": "./picobot.db" });
        std::fs::write(&config_path, serde_json::to_vec(&candidate).unwrap()).unwrap();

        load_candidate(
            &config_path,
            &HashMap::new(),
            temp.path(),
            &current,
            &workspace,
        )
        .unwrap();
    }

    #[tokio::test]
    async fn admission_closes_and_waits_for_existing_activity() {
        let admission = RuntimeAdmission::open();
        let activity = admission.try_enter().unwrap();
        admission.close();
        assert!(admission.try_enter().is_none());

        let waiting = admission.wait_for_idle();
        tokio::pin!(waiting);
        assert!(
            tokio::time::timeout(std::time::Duration::from_millis(10), &mut waiting)
                .await
                .is_err()
        );
        drop(activity);
        tokio::time::timeout(std::time::Duration::from_secs(1), waiting)
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn one_reload_remains_pending_until_its_generation_is_terminal() {
        let mut controller = ReloadController::new(HashMap::new(), PathBuf::from("."));
        let handle = controller.handle.clone();
        let first = tokio::spawn({
            let handle = handle.clone();
            async move { handle.request().await }
        });
        let request = controller.receiver.recv().await.unwrap();
        assert_eq!(request.generation, 2);
        request
            .response
            .send(Ok(ReloadAccepted {
                generation: 2,
                message: "accepted".to_string(),
            }))
            .unwrap();
        first.await.unwrap().unwrap();

        assert_eq!(
            handle.request().await.unwrap_err(),
            ReloadError::AlreadyPending
        );
        controller.set_phase(2, ReloadPhase::Active);

        let next = tokio::spawn({
            let handle = handle.clone();
            async move { handle.request().await }
        });
        let request = controller.receiver.recv().await.unwrap();
        assert_eq!(request.generation, 3);
        controller.set_failed(3, "invalid candidate");
        request
            .response
            .send(Err(ReloadError::InvalidConfig(
                "invalid candidate".to_string(),
            )))
            .unwrap();
        assert!(matches!(
            next.await.unwrap(),
            Err(ReloadError::InvalidConfig(_))
        ));
        let status = handle.status();
        assert_eq!(status.generation, 3);
        assert_eq!(status.phase, ReloadPhase::Failed);
    }
}