tuitbot-server 0.1.49

//! Shared application state for the tuitbot server.

use std::collections::HashMap;
use std::net::IpAddr;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Instant, SystemTime};

use tokio::sync::{broadcast, Mutex, RwLock};
use tokio_util::sync::CancellationToken;
use tuitbot_core::automation::circuit_breaker::CircuitBreaker;
use tuitbot_core::automation::Runtime;
use tuitbot_core::automation::WatchtowerLoop;
use tuitbot_core::config::{
    effective_config, Config, ConnectorConfig, ContentSourcesConfig, DeploymentMode,
};
use tuitbot_core::content::ContentGenerator;
use tuitbot_core::context::semantic_index::SemanticIndex;
use tuitbot_core::llm::embedding::EmbeddingProvider;
use tuitbot_core::llm::factory::create_provider;
use tuitbot_core::storage::accounts::{self, DEFAULT_ACCOUNT_ID};
use tuitbot_core::storage::DbPool;
use tuitbot_core::x_api::auth::TokenManager;
use tuitbot_core::x_api::ScraperHealth;

use tuitbot_core::error::XApiError;
use tuitbot_core::x_api::auth;

use crate::ws::AccountWsEvent;

/// Pending OAuth PKCE state for connector link flows.
pub struct PendingOAuth {
    /// The PKCE code verifier needed to complete the token exchange.
    pub code_verifier: String,
    /// When this entry was created (for 10-minute expiry).
    pub created_at: Instant,
    /// The account ID that initiated this OAuth flow (empty for connectors).
    pub account_id: String,
    /// The X API client ID used for this flow (for callback token exchange).
    pub client_id: String,
}

/// Shared application state accessible by all route handlers.
pub struct AppState {
    /// SQLite connection pool.
    pub db: DbPool,
    /// Path to the configuration file.
    pub config_path: PathBuf,
    /// Data directory for media storage (parent of config file).
    pub data_dir: PathBuf,
    /// Broadcast channel sender for real-time WebSocket events.
    pub event_tx: broadcast::Sender<AccountWsEvent>,
    /// Local bearer token for API authentication.
    pub api_token: String,
    /// Bcrypt hash of the web login passphrase (None if not configured).
    pub passphrase_hash: RwLock<Option<String>>,
    /// Last-observed mtime of the `passphrase_hash` file (for detecting out-of-band resets).
    pub passphrase_hash_mtime: RwLock<Option<SystemTime>>,
    /// Host address the server is bound to.
    pub bind_host: String,
    /// Port the server is listening on.
    pub bind_port: u16,
    /// Per-IP login attempt tracking for rate limiting: (count, window_start).
    pub login_attempts: Mutex<HashMap<IpAddr, (u32, Instant)>>,
    /// Per-account automation runtimes (keyed by account_id).
    pub runtimes: Mutex<HashMap<String, Runtime>>,
    /// Per-account content generators for AI assist endpoints.
    pub content_generators: Mutex<HashMap<String, Arc<ContentGenerator>>>,
    /// Optional circuit breaker for X API rate-limit protection.
    pub circuit_breaker: Option<Arc<CircuitBreaker>>,
    /// Optional scraper health tracker (populated when provider_backend = "scraper").
    pub scraper_health: Option<ScraperHealth>,
    /// Cancellation token for the Watchtower filesystem watcher (None if not running).
    pub watchtower_cancel: RwLock<Option<CancellationToken>>,
    /// Content sources configuration for the Watchtower.
    pub content_sources: RwLock<ContentSourcesConfig>,
    /// Connector configuration for remote source OAuth flows.
    pub connector_config: ConnectorConfig,
    /// Deployment mode (desktop, self_host, or cloud).
    pub deployment_mode: DeploymentMode,
    /// Pending OAuth PKCE challenges keyed by state parameter.
    pub pending_oauth: Mutex<HashMap<String, PendingOAuth>>,
    /// Per-account X API token managers for automatic token refresh.
    pub token_managers: Mutex<HashMap<String, Arc<TokenManager>>>,
    /// X API client ID from config (needed to create token managers).
    pub x_client_id: String,
    /// In-memory semantic search index (None if embedding not configured).
    pub semantic_index: Option<Arc<RwLock<SemanticIndex>>>,
    /// Embedding provider for semantic indexing (None if not configured).
    pub embedding_provider: Option<Arc<dyn EmbeddingProvider>>,
}

impl AppState {
    /// Get a fresh X API access token for the given account.
    ///
    /// Lazily creates a `TokenManager` on first use (loading tokens from disk),
    /// then returns a token that is automatically refreshed before expiry.
    pub async fn get_x_access_token(
        &self,
        token_path: &std::path::Path,
        account_id: &str,
    ) -> Result<String, XApiError> {
        // Fast path: token manager already exists.
        {
            let managers = self.token_managers.lock().await;
            if let Some(tm) = managers.get(account_id) {
                return tm.get_access_token().await;
            }
        }

        // Load tokens from disk and create a new manager.
        let tokens = auth::load_tokens(token_path)?.ok_or(XApiError::AuthExpired)?;

        let tm = Arc::new(TokenManager::new(
            tokens,
            self.x_client_id.clone(),
            token_path.to_path_buf(),
        ));

        let access_token = tm.get_access_token().await?;

        self.token_managers
            .lock()
            .await
            .insert(account_id.to_string(), tm);

        Ok(access_token)
    }

    /// Load the effective config for a given account.
    ///
    /// Default account: reads config.toml directly (backward compat).
    /// Non-default: merges config.toml base with account's `config_overrides` from DB.
    pub async fn load_effective_config(&self, account_id: &str) -> Result<Config, String> {
        let contents = std::fs::read_to_string(&self.config_path).unwrap_or_default();
        let base: Config = toml::from_str(&contents).unwrap_or_default();

        if account_id == DEFAULT_ACCOUNT_ID {
            return Ok(base);
        }

        let account = accounts::get_account(&self.db, account_id)
            .await
            .map_err(|e| e.to_string())?
            .ok_or_else(|| format!("account not found: {account_id}"))?;

        effective_config(&base, &account.config_overrides)
            .map(|r| r.config)
            .map_err(|e| e.to_string())
    }

    /// Lazily create or return a cached `ContentGenerator` for the given account.
    ///
    /// Loads effective config, creates the LLM provider, and caches the generator.
    pub async fn get_or_create_content_generator(
        &self,
        account_id: &str,
    ) -> Result<Arc<ContentGenerator>, String> {
        // Fast path: already cached.
        {
            let generators = self.content_generators.lock().await;
            if let Some(gen) = generators.get(account_id) {
                return Ok(gen.clone());
            }
        }

        let config = self.load_effective_config(account_id).await?;

        let provider =
            create_provider(&config.llm).map_err(|e| format!("LLM not configured: {e}"))?;

        let gen = Arc::new(ContentGenerator::new(provider, config.business));

        self.content_generators
            .lock()
            .await
            .insert(account_id.to_string(), gen.clone());

        Ok(gen)
    }

    /// Returns `true` if the current deployment mode is local-first (Desktop).
    pub fn is_local_first(&self) -> bool {
        self.deployment_mode.is_local_first()
    }

    /// Cancel the running Watchtower (if any), reload config from disk,
    /// and spawn a new Watchtower loop with the updated sources.
    ///
    /// Called after `PATCH /api/settings` modifies `content_sources` or
    /// `deployment_mode`.
    pub async fn restart_watchtower(&self) {
        // 1. Cancel existing watchtower.
        if let Some(cancel) = self.watchtower_cancel.write().await.take() {
            cancel.cancel();
            tracing::info!("Watchtower cancelled for config reload");
        }

        // 2. Reload config from disk.
        let loaded_config = Config::load(Some(&self.config_path.to_string_lossy())).ok();
        let new_sources = loaded_config
            .as_ref()
            .map(|c| c.content_sources.clone())
            .unwrap_or_default();
        let connector_config = loaded_config
            .as_ref()
            .map(|c| c.connectors.clone())
            .unwrap_or_default();
        let deployment_mode = loaded_config
            .as_ref()
            .map(|c| c.deployment_mode.clone())
            .unwrap_or_default();

        // 3. Check if any sources are enabled and eligible.
        let has_enabled: Vec<_> = new_sources
            .sources
            .iter()
            .filter(|s| {
                s.is_enabled()
                    && deployment_mode.allows_source_type(&s.source_type)
                    && (s.path.is_some() || s.folder_id.is_some())
            })
            .collect();

        if has_enabled.is_empty() {
            tracing::info!("Watchtower restart: no enabled sources, not spawning");
            *self.content_sources.write().await = new_sources;
            return;
        }

        // Surface privacy envelope for operators: local_fs in non-Desktop mode
        // means data is user-controlled but not same-machine local-first.
        if !deployment_mode.is_local_first()
            && has_enabled.iter().any(|s| s.source_type == "local_fs")
        {
            tracing::info!(
                mode = %deployment_mode,
                "local_fs source in {} mode — data is user-controlled but not local-first",
                deployment_mode
            );
        }

        // 4. Spawn new WatchtowerLoop.
        let cancel = CancellationToken::new();
        let watchtower = WatchtowerLoop::new(
            self.db.clone(),
            new_sources.clone(),
            connector_config,
            self.data_dir.clone(),
        );
        let cancel_clone = cancel.clone();
        tokio::spawn(async move {
            watchtower.run(cancel_clone).await;
        });

        tracing::info!(
            sources = has_enabled.len(),
            "Watchtower restarted with updated config"
        );

        // 5. Update state.
        *self.watchtower_cancel.write().await = Some(cancel);
        *self.content_sources.write().await = new_sources;
    }
}