geist_supervisor 0.1.28

use crate::app_config::AppConfig;
use crate::config::Config;
use crate::services::ota::OtaCoordinator;
use crate::services::{FileService, GcsService};
use anyhow::{Context, Result};
use serde_json;
use std::fs;
use std::path::{Path, PathBuf};
use tempfile::{self, TempDir};
use tracing::{info, warn};

pub struct UpdateService {
    app_binary_path: PathBuf,
    service_name: String,
    registry_url: String,
    release_bundle_name: String,
    checksum_file_name: String,
    binary_names: Vec<String>,
    health_check_url: Option<String>,
    health_check_timeout_secs: u32,
    health_check_retries: u32,
}

/// Context for an update operation, containing all necessary state
struct UpdateContext {
    current_version: String,
    target_version: Option<String>,
    bundle_path: PathBuf,
    signature_verified: bool,
    _temp_dir_handle: Option<TempDir>,
}

impl UpdateService {
    pub fn with_config(app: &AppConfig) -> Self {
        Self {
            app_binary_path: app.app_binary_path(),
            service_name: app.service_name.clone(),
            registry_url: app.registry_url(),
            release_bundle_name: app.release_bundle_name.clone(),
            checksum_file_name: app.checksum_file_name.clone(),
            binary_names: app.binary_names.clone(),
            health_check_url: app.health_check_url.clone(),
            health_check_timeout_secs: app.health_check_timeout_secs,
            health_check_retries: app.health_check_retries,
        }
    }

    pub fn new() -> Self {
        Self::with_config(&AppConfig::default())
    }

    /// Perform a dry run update - show what would be done without making changes
    pub fn dry_run_update(&self, source: Option<String>) -> Result<()> {
        let source = source.unwrap_or_else(|| "latest".to_string());
        let current_version = Config::format_version(&Config::get_current_version());

        println!("=== DRY RUN MODE ===");
        println!("Current version: {}", current_version);

        let source_path = Path::new(&source);

        if source_path.is_file() {
            println!("Would use local update bundle: {}", source_path.display());

            let (staged_bundle_path, _signature_verified, _staged_bundle_guard) =
                self.stage_local_bundle(source_path)?;
            let (release_bundle_dir, _temp_guard) = self.extract_bundle(&staged_bundle_path)?;

            if let Ok(Some(target_version)) = self.detect_version_from_bundle(&release_bundle_dir) {
                println!("Target version: {}", target_version);

                if self.is_same_version(&current_version, &target_version) {
                    println!(
                        "Already on target version {}, no update needed",
                        target_version
                    );
                    return Ok(());
                }

                println!(
                    "Would update from {} to {}",
                    current_version, target_version
                );
            } else {
                println!("Target version: Unable to determine from bundle");
                println!("Would proceed with update using local bundle");
            }

            // Check what's in the bundle
            self.analyze_bundle_contents(&release_bundle_dir)?;
        } else {
            // Remote update
            let gcs = GcsService::with_bundle_names(
                String::new(),
                self.registry_url.clone(),
                self.release_bundle_name.clone(),
                self.checksum_file_name.clone(),
            );

            let version_to_fetch = if source == "latest" {
                println!("Fetching latest version information...");
                match gcs.get_latest_version() {
                    Ok(version) => version,
                    Err(e) => {
                        println!("Error: Failed to fetch latest version: {}", e);
                        return Ok(());
                    }
                }
            } else {
                Config::format_version(&source)
            };

            println!("Target version: {}", version_to_fetch);

            if self.is_same_version(&current_version, &version_to_fetch) {
                println!(
                    "Already on target version {}, no update needed",
                    version_to_fetch
                );
                return Ok(());
            }

            println!(
                "Would update from {} to {}",
                current_version, version_to_fetch
            );

            // Verify version exists
            match gcs.verify_version(&version_to_fetch) {
                Ok(true) => println!("✓ Version {} is available for download", version_to_fetch),
                Ok(false) => {
                    println!("✗ Version {} not found", version_to_fetch);
                    return Ok(());
                }
                Err(e) => {
                    println!("✗ Error verifying version: {}", e);
                    return Ok(());
                }
            }

            println!(
                "Would download release bundle for version {}",
                version_to_fetch
            );
        }

        // Show what operations would be performed
        println!("\nOperations that would be performed:");
        println!("1. Create update lock file");
        println!("2. Extract release bundle");
        println!("3. Backup current installation");
        println!("4. Stop service: {}", self.service_name);

        if source_path.is_file() {
            let (staged_bundle_path, _signature_verified, _staged_bundle_guard) =
                self.stage_local_bundle(source_path)?;
            let (release_bundle_dir, _temp_guard) = self.extract_bundle(&staged_bundle_path)?;
            let install_script_path = release_bundle_dir.join("install.sh");

            if install_script_path.exists() {
                println!("5. Execute install.sh script (comprehensive installation)");
            } else {
                println!("5. Replace binary only");
            }
        } else {
            println!("5. Install new version");
        }

        println!("6. Start service: {}", self.service_name);
        println!("7. Verify installation");
        println!("8. Update version tracking");
        println!("9. Clean up temporary files");

        println!("\n=== END DRY RUN ===");
        println!("No changes were made to the system.");

        Ok(())
    }

    /// Perform update to specified version or latest if version is None
    pub fn update(&self, source: Option<String>) -> Result<()> {
        let _lock_guard = self.create_update_lock()?;
        let source = source.unwrap_or_else(|| "latest".to_string());
        let current_version = Config::format_version(&Config::get_current_version());
        info!("Current version: {}", current_version);

        let mut context = self.prepare_update_context(&source, current_version)?;

        // Early return if already on target version (for remote updates)
        if let Some(ref target_version) = context.target_version {
            if self.is_same_version(&context.current_version, target_version) {
                info!(
                    "Already on target version {}, no update needed",
                    target_version
                );
                println!("Already on version {}, no update needed.", target_version);
                return Ok(());
            }
        }

        let (release_bundle_dir, _temp_dir_guard) = self.extract_bundle(&context.bundle_path)?;

        // Detect version from bundle if not already known (for local files)
        if context.target_version.is_none() {
            context.target_version = self.detect_version_from_bundle(&release_bundle_dir)?;

            // Check version again for local files
            if let Some(ref target_version) = context.target_version {
                if self.is_same_version(&context.current_version, target_version) {
                    info!(
                        "Already on target version {}, no update needed",
                        target_version
                    );
                    println!("Already on version {}, no update needed.", target_version);
                    return Ok(());
                }
                info!(
                    "Updating from {} to {}",
                    context.current_version, target_version
                );
            }
        }

        self.perform_update_operation(&release_bundle_dir, context.signature_verified)?;
        self.finalize_version_update(context.target_version)?;

        info!("Update completed successfully!");
        Ok(())
    }

    /// Prepare the update context based on the source (local file or remote version)
    fn prepare_update_context(
        &self,
        source: &str,
        current_version: String,
    ) -> Result<UpdateContext> {
        let source_path = Path::new(source);

        if source_path.is_file() {
            info!("Using local update bundle: {}", source_path.display());
            let (staged_bundle_path, signature_verified, temp_dir_handle) =
                self.stage_local_bundle(source_path)?;

            Ok(UpdateContext {
                current_version,
                target_version: None,
                bundle_path: staged_bundle_path,
                signature_verified,
                _temp_dir_handle: Some(temp_dir_handle),
            })
        } else {
            self.prepare_remote_update_context(source, current_version)
        }
    }

    /// Prepare context for remote update (download from GCS)
    fn prepare_remote_update_context(
        &self,
        source: &str,
        current_version: String,
    ) -> Result<UpdateContext> {
        let gcs = GcsService::with_bundle_names(
            String::new(),
            self.registry_url.clone(),
            self.release_bundle_name.clone(),
            self.checksum_file_name.clone(),
        );

        let version_to_fetch = if source == "latest" {
            info!("No version specified, fetching latest version...");
            gcs.get_latest_version()?
        } else {
            Config::format_version(source)
        };

        info!("Target version: {}", version_to_fetch);

        // Verify version exists
        if !gcs.verify_version(&version_to_fetch)? {
            anyhow::bail!("Version {} not found", version_to_fetch);
        }

        // Download release bundle
        let temp_dir = tempfile::tempdir()?;
        let bundle_path_in_temp = temp_dir.path().join(&self.release_bundle_name);

        info!(
            "Downloading release bundle to: {}",
            bundle_path_in_temp.display()
        );
        gcs.download_release_bundle(&version_to_fetch, &bundle_path_in_temp)?;

        // Download and verify signature
        let sig_path = temp_dir.path().join("bundle.sig");
        let signature_verified = match gcs.download_signature(&version_to_fetch, &sig_path) {
            Ok(()) => {
                crate::services::signature::verify_bundle(&bundle_path_in_temp, &sig_path)?;
                true
            }
            Err(e) => {
                if crate::services::signature::REQUIRE_SIGNATURE {
                    anyhow::bail!("Signature required but not found: {}", e);
                }
                warn!("No signature available (transition period): {}", e);
                false
            }
        };

        Ok(UpdateContext {
            current_version,
            target_version: Some(version_to_fetch),
            bundle_path: bundle_path_in_temp,
            signature_verified,
            _temp_dir_handle: Some(temp_dir),
        })
    }

    fn signature_sidecar_path(bundle_path: &Path) -> PathBuf {
        let mut sig_path = bundle_path.as_os_str().to_owned();
        sig_path.push(".sig");
        PathBuf::from(sig_path)
    }

    /// Copy a local bundle and optional signature into a private temp directory,
    /// then verify the staged signature before returning.
    fn stage_local_bundle(&self, source_path: &Path) -> Result<(PathBuf, bool, TempDir)> {
        let bundle_name = source_path.file_name().ok_or_else(|| {
            anyhow::anyhow!(
                "Local bundle path must include a file name: {}",
                source_path.display()
            )
        })?;

        let temp_dir = tempfile::tempdir()?;
        let staged_bundle_path = temp_dir.path().join(bundle_name);
        fs::copy(source_path, &staged_bundle_path).with_context(|| {
            format!(
                "Failed to stage local bundle from {} to {}",
                source_path.display(),
                staged_bundle_path.display()
            )
        })?;

        let source_sig_path = Self::signature_sidecar_path(source_path);
        let staged_sig_path = Self::signature_sidecar_path(&staged_bundle_path);
        let signature_verified = match fs::copy(&source_sig_path, &staged_sig_path) {
            Ok(_) => {
                match crate::services::signature::verify_bundle(
                    &staged_bundle_path,
                    &staged_sig_path,
                ) {
                    Ok(()) => {
                        info!("Local bundle signature verified OK");
                        true
                    }
                    Err(e) => {
                        anyhow::bail!("Local bundle signature verification FAILED: {}", e);
                    }
                }
            }
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
                if crate::services::signature::REQUIRE_SIGNATURE {
                    anyhow::bail!(
                        "Signature required but no .sig file found at: {}",
                        source_sig_path.display()
                    );
                }
                warn!(
                    "No signature file found at {} — skipping verification",
                    source_sig_path.display()
                );
                false
            }
            Err(e) => {
                return Err(e).with_context(|| {
                    format!(
                        "Failed to stage local bundle signature from {} to {}",
                        source_sig_path.display(),
                        staged_sig_path.display()
                    )
                });
            }
        };

        Ok((staged_bundle_path, signature_verified, temp_dir))
    }

    /// Extract the bundle and return the extraction directory
    fn extract_bundle(&self, bundle_path: &Path) -> Result<(PathBuf, TempDir)> {
        let data_dir = Config::data_dir();
        let fs_service = FileService::new(data_dir);

        info!("Extracting release bundle from: {}", bundle_path.display());
        let temp_dir_for_extraction = tempfile::tempdir()?;
        let release_bundle_dir =
            fs_service.extract_bundle_with_details(bundle_path, temp_dir_for_extraction.path())?;
        Ok((release_bundle_dir, temp_dir_for_extraction))
    }

    /// Try to detect version from bundle metadata
    fn detect_version_from_bundle(&self, release_bundle_dir: &Path) -> Result<Option<String>> {
        let metadata_path = release_bundle_dir.join("ota_metadata.json");

        if let Ok(metadata_content) = fs::read_to_string(metadata_path) {
            if let Ok(json) = serde_json::from_str::<serde_json::Value>(&metadata_content) {
                if let Some(version) = json["version"].as_str() {
                    let formatted_version = Config::format_version(version);
                    info!(
                        "Detected version from bundle metadata: {}",
                        formatted_version
                    );
                    return Ok(Some(formatted_version));
                }
            }
        }

        warn!("Could not determine version from bundle metadata - update will proceed without version tracking");
        Ok(None)
    }

    /// Check if two versions are the same (normalized comparison)
    fn is_same_version(&self, version1: &str, version2: &str) -> bool {
        Config::format_version(version1) == Config::format_version(version2)
    }

    /// Perform the actual update operation using OtaCoordinator
    fn perform_update_operation(
        &self,
        release_bundle_dir: &Path,
        signature_verified: bool,
    ) -> Result<()> {
        info!("Using app binary path: {}", self.app_binary_path.display());
        let ota_coordinator = OtaCoordinator::with_binary_names(
            self.app_binary_path.clone(),
            self.service_name.clone(),
            self.binary_names.clone(),
        )
        .with_health_check(
            self.health_check_url.clone(),
            self.health_check_timeout_secs,
            self.health_check_retries,
        );

        let install_script_path = release_bundle_dir.join("install.sh");
        info!(
            "Checking for install script at: {}",
            install_script_path.display()
        );
        info!("Install script exists: {}", install_script_path.exists());

        if install_script_path.exists() {
            info!("Found install.sh script, performing comprehensive installation...");
            ota_coordinator.perform_install_script_update(
                &install_script_path,
                release_bundle_dir,
                signature_verified,
            )
        } else {
            info!("No install.sh found, performing binary-only update with asset copy...");
            let binary_path = ota_coordinator.find_binary_in_bundle(release_bundle_dir)?;
            ota_coordinator.perform_binary_update_with_assets(&binary_path, release_bundle_dir)
        }
    }

    /// Analyze bundle contents for dry-run mode
    fn analyze_bundle_contents(&self, release_bundle_dir: &Path) -> Result<()> {
        println!("\nBundle contents analysis:");

        let install_script_path = release_bundle_dir.join("install.sh");
        if install_script_path.exists() {
            println!("✓ install.sh script found - comprehensive installation mode");
        } else {
            println!("✗ No install.sh script - binary-only update mode");
        }

        let metadata_path = release_bundle_dir.join("ota_metadata.json");
        if metadata_path.exists() {
            println!("✓ ota_metadata.json found");
        } else {
            println!("✗ No ota_metadata.json found");
        }

        // Look for common files
        let common_files = ["config.toml", "assets/", "dependencies/"];
        for file in &common_files {
            let file_path = release_bundle_dir.join(file);
            if file_path.exists() {
                println!("✓ {} found", file);
            }
        }

        Ok(())
    }

    /// Log the finalized version (ota_metadata.json is deployed by install.sh)
    fn finalize_version_update(&self, target_version: Option<String>) -> Result<()> {
        if let Some(version) = target_version {
            let formatted_version = Config::format_version(&version);
            info!(
                "Update finalized to version: {} (version tracked via ota_metadata.json)",
                formatted_version
            );
        } else {
            warn!("No version information available from bundle metadata");
        }

        Ok(())
    }

    /// Create update lock file to prevent service restart during update
    fn create_update_lock(&self) -> Result<impl Drop> {
        FlockGuard::acquire(Path::new("/tmp/geist_supervisor_update_in_progress.lock"))
    }
}

/// Guard that holds an exclusive flock(2) on a lock file.
/// The kernel releases the lock when the file descriptor is closed (on drop).
#[derive(Debug)]
struct FlockGuard {
    _file: fs::File,
}

impl Drop for FlockGuard {
    fn drop(&mut self) {
        // Kernel releases the flock when fd is closed, which happens
        // automatically when _file is dropped after this.
    }
}

impl FlockGuard {
    fn acquire(lock_path: &Path) -> Result<Self> {
        let file = fs::OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(false)
            .open(lock_path)?;

        #[cfg(unix)]
        {
            use std::os::unix::io::AsRawFd;
            let rc = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) };
            if rc != 0 {
                let err = std::io::Error::last_os_error();
                match err.raw_os_error() {
                    // POSIX allows EAGAIN and EWOULDBLOCK to differ; on Linux they're the same
                    #[allow(unreachable_patterns)]
                    Some(libc::EWOULDBLOCK) | Some(libc::EAGAIN) => {
                        anyhow::bail!("Update already in progress (could not acquire lock)");
                    }
                    _ => {
                        anyhow::bail!("Failed to acquire update lock: {}", err);
                    }
                }
            }
        }

        info!("Acquired update lock: {}", lock_path.display());
        Ok(Self { _file: file })
    }
}

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

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

    #[test]
    fn test_flock_guard_acquires_lock() {
        let temp_dir = tempdir().unwrap();
        let lock_path = temp_dir.path().join("test.lock");

        let guard = FlockGuard::acquire(&lock_path);
        assert!(guard.is_ok(), "Should acquire lock: {:?}", guard.err());
        assert!(lock_path.exists(), "Lock file should exist");
    }

    #[test]
    fn test_flock_guard_contention() {
        let temp_dir = tempdir().unwrap();
        let lock_path = temp_dir.path().join("test.lock");

        let _guard1 = FlockGuard::acquire(&lock_path).unwrap();
        let guard2 = FlockGuard::acquire(&lock_path);

        assert!(guard2.is_err(), "Second acquire should fail");
        let err_msg = guard2.unwrap_err().to_string();
        assert!(
            err_msg.contains("already in progress"),
            "Error should mention contention, got: {}",
            err_msg
        );
    }

    #[test]
    fn test_flock_guard_release_on_drop() {
        let temp_dir = tempdir().unwrap();
        let lock_path = temp_dir.path().join("test.lock");

        {
            let _guard = FlockGuard::acquire(&lock_path).unwrap();
        } // guard dropped here

        // Should be able to re-acquire after drop
        let guard2 = FlockGuard::acquire(&lock_path);
        assert!(
            guard2.is_ok(),
            "Should re-acquire after drop: {:?}",
            guard2.err()
        );
    }
}