foxhunt/crates/ml-checkpoint/src/versioning.rs

//! Version management for model checkpoints
//!
//! Handles semantic versioning, compatibility checks, and migration paths.

use std::cmp::Ordering;
use std::collections::HashMap;

use serde::{Deserialize, Serialize};
use tracing::{debug, info};

use super::{CheckpointMetadata, ModelType};
use crate::MLError;

/// Semantic version representation
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SemanticVersion {
    /// Major version - breaking changes
    pub major: u32,

    /// Minor version - new features, backward compatible
    pub minor: u32,

    /// Patch version - bug fixes
    pub patch: u32,

    /// Pre-release identifier (e.g., "alpha", "beta", "rc1")
    pub pre_release: Option<String>,

    /// Build metadata
    pub build_metadata: Option<String>,
}

impl SemanticVersion {
    /// Create a new semantic version
    pub const fn new(major: u32, minor: u32, patch: u32) -> Self {
        Self {
            major,
            minor,
            patch,
            pre_release: None,
            build_metadata: None,
        }
    }

    /// Parse version string (e.g., "1.2.3-alpha+build1")
    pub fn parse(version_str: &str) -> Result<Self, MLError> {
        let mut parts = version_str.split('+');
        let version_part = parts.next().unwrap_or("");
        let build_metadata = parts.next().map(|s| s.to_owned());

        let mut version_pre_parts = version_part.split('-');
        let version_core = version_pre_parts.next().unwrap_or("");
        let pre_release = version_pre_parts.next().map(|s| s.to_owned());

        let version_nums: Vec<&str> = version_core.split('.').collect();
        if version_nums.len() != 3 {
            return Err(MLError::ModelError(format!(
                "Invalid version format: {} (expected major.minor.patch)",
                version_str
            )));
        }

        let major = version_nums[0].parse::<u32>().map_err(|e| {
            MLError::ModelError(format!("Invalid major version '{}': {e}", version_nums[0]))
        })?;

        let minor = version_nums[1].parse::<u32>().map_err(|e| {
            MLError::ModelError(format!("Invalid minor version '{}': {e}", version_nums[1]))
        })?;

        let patch = version_nums[2].parse::<u32>().map_err(|e| {
            MLError::ModelError(format!("Invalid patch version '{}': {e}", version_nums[2]))
        })?;

        Ok(Self {
            major,
            minor,
            patch,
            pre_release,
            build_metadata,
        })
    }

    /// Check if this version is compatible with another version
    pub const fn is_compatible_with(&self, other: &SemanticVersion) -> bool {
        // Major version must match for compatibility
        if self.major != other.major {
            return false;
        }

        // If major versions match, it's considered compatible
        // (minor/patch differences are handled separately)
        true
    }

    /// Check if this version is newer than another
    pub fn is_newer_than(&self, other: &SemanticVersion) -> bool {
        self > other
    }

    /// Get compatibility risk level
    pub const fn compatibility_risk(&self, other: &SemanticVersion) -> CompatibilityRisk {
        if self.major != other.major {
            CompatibilityRisk::High
        } else if self.minor != other.minor {
            CompatibilityRisk::Medium
        } else if self.patch != other.patch {
            CompatibilityRisk::Low
        } else {
            CompatibilityRisk::None
        }
    }
}

impl std::fmt::Display for SemanticVersion {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}.{}.{}", self.major, self.minor, self.patch)?;
        if let Some(ref pre_release) = self.pre_release {
            write!(f, "-{pre_release}")?;
        }
        if let Some(ref build_metadata) = self.build_metadata {
            write!(f, "+{build_metadata}")?;
        }
        Ok(())
    }
}

impl PartialOrd for SemanticVersion {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for SemanticVersion {
    fn cmp(&self, other: &Self) -> Ordering {
        match self.major.cmp(&other.major) {
            Ordering::Equal => {
                match self.minor.cmp(&other.minor) {
                    Ordering::Equal => {
                        match self.patch.cmp(&other.patch) {
                            Ordering::Equal => {
                                // Handle pre-release comparison
                                match (&self.pre_release, &other.pre_release) {
                                    (None, None) => Ordering::Equal,
                                    (Some(_), None) => Ordering::Less, // Pre-release is less than release
                                    (None, Some(_)) => Ordering::Greater,
                                    (Some(a), Some(b)) => a.cmp(b),
                                }
                            },
                            other @ Ordering::Less | other @ Ordering::Greater => other,
                        }
                    },
                    other @ Ordering::Less | other @ Ordering::Greater => other,
                }
            },
            other @ Ordering::Less | other @ Ordering::Greater => other,
        }
    }
}

/// Compatibility risk levels
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CompatibilityRisk {
    /// No compatibility issues expected
    None,

    /// Low risk - patch version differences
    Low,

    /// Medium risk - minor version differences
    Medium,

    /// High risk - major version differences
    High,
}

/// Version manager for handling model versioning
#[derive(Debug)]
pub struct VersionManager {
    /// Migration handlers for version upgrades
    migration_handlers: HashMap<(ModelType, String, String), MigrationHandler>,
}

impl VersionManager {
    /// Create a new version manager
    pub fn new() -> Self {
        Self {
            migration_handlers: HashMap::new(),
        }
    }

    /// Check version compatibility between two checkpoints
    pub fn check_compatibility(
        &self,
        current_version: &str,
        checkpoint_version: &str,
        model_type: ModelType,
    ) -> Result<CompatibilityInfo, MLError> {
        let current = SemanticVersion::parse(current_version)?;
        let checkpoint = SemanticVersion::parse(checkpoint_version)?;

        let risk = current.compatibility_risk(&checkpoint);
        let compatible = current.is_compatible_with(&checkpoint);

        let info = CompatibilityInfo {
            compatible,
            risk,
            current_version: current.clone(),
            checkpoint_version: checkpoint.clone(),
            migration_required: !compatible || risk == CompatibilityRisk::High,
            warnings: self.generate_compatibility_warnings(&current, &checkpoint, model_type),
        };

        debug!("Compatibility check: {:?}", info);
        Ok(info)
    }

    /// Generate compatibility warnings
    fn generate_compatibility_warnings(
        &self,
        current: &SemanticVersion,
        checkpoint: &SemanticVersion,
        model_type: ModelType,
    ) -> Vec<String> {
        let mut warnings = Vec::new();

        if current.major != checkpoint.major {
            warnings.push(format!(
                "Major version mismatch for {:?}: current {}, checkpoint {}. This may cause loading failures.",
                model_type, current.major, checkpoint.major
            ));
        }

        if current.minor < checkpoint.minor {
            warnings.push(format!(
                "Loading newer minor version for {:?}: current {}.{}, checkpoint {}.{}. Some features may be unavailable.",
                model_type, current.major, current.minor, checkpoint.major, checkpoint.minor
            ));
        }

        if current.minor > checkpoint.minor + 2 {
            warnings.push(format!(
                "Loading significantly older checkpoint for {:?}: current {}.{}, checkpoint {}.{}. Consider retraining.",
                model_type, current.major, current.minor, checkpoint.major, checkpoint.minor
            ));
        }

        if checkpoint.pre_release.is_some() {
            warnings.push("Loading pre-release checkpoint. Stability not guaranteed.".to_owned());
        }

        warnings
    }

    /// Register a migration handler for version transitions
    pub fn register_migration_handler(
        &mut self,
        model_type: ModelType,
        from_version: String,
        to_version: String,
        handler: MigrationHandler,
    ) {
        let key = (model_type, from_version.clone(), to_version.clone());
        self.migration_handlers.insert(key, handler);
        info!(
            "Registered migration handler for {:?} -> {:?}",
            from_version, to_version
        );
    }

    /// Get migration path between versions
    pub fn get_migration_path(
        &self,
        _model_type: ModelType,
        from_version: &str,
        to_version: &str,
    ) -> Result<Vec<MigrationStep>, MLError> {
        let from = SemanticVersion::parse(from_version)?;
        let to = SemanticVersion::parse(to_version)?;

        // For now, simple direct migration
        // In a more complex system, this could handle multi-step migrations
        if from == to {
            return Ok(vec![]);
        }

        let step = MigrationStep {
            from_version: from.clone(),
            to_version: to.clone(),
            migration_type: if from.major != to.major {
                MigrationType::Major
            } else if from.minor != to.minor {
                MigrationType::Minor
            } else {
                MigrationType::Patch
            },
            description: format!("Migrate from {} to {}", from_version, to_version),
            required: from.major != to.major,
        };

        Ok(vec![step])
    }

    /// Find the latest compatible version from a list of checkpoints
    pub fn find_latest_compatible_version(
        &self,
        current_version: &str,
        checkpoints: &[CheckpointMetadata],
        model_type: ModelType,
        model_name: &str,
    ) -> Result<Option<CheckpointMetadata>, MLError> {
        let current = SemanticVersion::parse(current_version)?;

        let mut compatible_checkpoints: Vec<_> = checkpoints
            .iter()
            .filter(|checkpoint| {
                checkpoint.model_type == model_type && checkpoint.model_name == model_name
            })
            .filter_map(|checkpoint| {
                SemanticVersion::parse(&checkpoint.version)
                    .ok()
                    .map(|version| (checkpoint, version))
            })
            .filter(|(_, version)| current.is_compatible_with(version))
            .collect();

        // Sort by version (newest first)
        compatible_checkpoints.sort_by(|a, b| b.1.cmp(&a.1));

        Ok(compatible_checkpoints
            .into_iter()
            .next()
            .map(|(checkpoint, _)| checkpoint.clone()))
    }

    /// Suggest next version for a checkpoint
    pub fn suggest_next_version(
        &self,
        current_version: &str,
        change_type: VersionChangeType,
    ) -> Result<String, MLError> {
        let mut version = SemanticVersion::parse(current_version)?;

        match change_type {
            VersionChangeType::Major => {
                version.major += 1;
                version.minor = 0;
                version.patch = 0;
            },
            VersionChangeType::Minor => {
                version.minor += 1;
                version.patch = 0;
            },
            VersionChangeType::Patch => {
                version.patch += 1;
            },
        }

        // Clear pre-release and build metadata for releases
        version.pre_release = None;
        version.build_metadata = None;

        Ok(version.to_string())
    }
}

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

/// Compatibility information between versions
#[derive(Debug, Clone)]
pub struct CompatibilityInfo {
    /// Whether versions are compatible
    pub compatible: bool,

    /// Risk level of using the checkpoint
    pub risk: CompatibilityRisk,

    /// Current model version
    pub current_version: SemanticVersion,

    /// Checkpoint version
    pub checkpoint_version: SemanticVersion,

    /// Whether migration is required
    pub migration_required: bool,

    /// Compatibility warnings
    pub warnings: Vec<String>,
}

/// Migration step between versions
#[derive(Debug, Clone)]
pub struct MigrationStep {
    /// Source version
    pub from_version: SemanticVersion,

    /// Target version
    pub to_version: SemanticVersion,

    /// Type of migration
    pub migration_type: MigrationType,

    /// Human-readable description
    pub description: String,

    /// Whether migration is required (or optional)
    pub required: bool,
}

/// Types of version changes
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VersionChangeType {
    /// Breaking changes
    Major,

    /// New features, backward compatible
    Minor,

    /// Bug fixes
    Patch,
}

/// Types of migrations
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MigrationType {
    /// Major version migration - potentially breaking
    Major,

    /// Minor version migration - new features
    Minor,

    /// Patch version migration - bug fixes
    Patch,
}

/// Migration handler function type
pub type MigrationHandler = fn(&[u8]) -> Result<Vec<u8>, MLError>;

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


    #[test]
    fn test_semantic_version_parsing() -> Result<(), MLError> {
        // Basic version
        let v1 = SemanticVersion::parse("1.2.3")?;
        assert_eq!(v1.major, 1);
        assert_eq!(v1.minor, 2);
        assert_eq!(v1.patch, 3);
        assert_eq!(v1.pre_release, None);
        assert_eq!(v1.build_metadata, None);

        // Version with pre-release
        let v2 = SemanticVersion::parse("1.0.0-alpha")?;
        assert_eq!(v2.pre_release, Some("alpha".to_owned()));

        // Version with build metadata
        let v3 = SemanticVersion::parse("1.0.0+build1")?;
        assert_eq!(v3.build_metadata, Some("build1".to_owned()));

        // Version with both
        let v4 = SemanticVersion::parse("2.1.0-beta+exp.1")?;
        assert_eq!(v4.pre_release, Some("beta".to_owned()));
        assert_eq!(v4.build_metadata, Some("exp.1".to_owned()));

        // Invalid versions
        assert!(SemanticVersion::parse("1.2").is_err());
        assert!(SemanticVersion::parse("a.b.c").is_err());
        assert!(SemanticVersion::parse("1.2.3.4").is_err());
        Ok(())
    }

    #[test]
    fn test_semantic_version_comparison() -> Result<(), MLError> {
        let v1_0_0 = SemanticVersion::parse("1.0.0")?;
        let v1_1_0 = SemanticVersion::parse("1.1.0")?;
        let v2_0_0 = SemanticVersion::parse("2.0.0")?;
        let v1_0_0_alpha = SemanticVersion::parse("1.0.0-alpha")?;

        // Basic comparisons
        assert!(v2_0_0 > v1_1_0);
        assert!(v1_1_0 > v1_0_0);
        assert!(v1_0_0 > v1_0_0_alpha);

        // Compatibility checks
        assert!(v1_0_0.is_compatible_with(&v1_1_0));
        assert!(!v1_0_0.is_compatible_with(&v2_0_0));

        // Newer checks
        assert!(v2_0_0.is_newer_than(&v1_0_0));
        assert!(!v1_0_0.is_newer_than(&v2_0_0));
        Ok(())
    }

    #[test]
    fn test_compatibility_risk() -> Result<(), MLError> {
        let v1_0_0 = SemanticVersion::parse("1.0.0")?;
        let v1_0_1 = SemanticVersion::parse("1.0.1")?;
        let v1_1_0 = SemanticVersion::parse("1.1.0")?;
        let v2_0_0 = SemanticVersion::parse("2.0.0")?;

        assert_eq!(v1_0_0.compatibility_risk(&v1_0_0), CompatibilityRisk::None);
        assert_eq!(v1_0_0.compatibility_risk(&v1_0_1), CompatibilityRisk::Low);
        assert_eq!(
            v1_0_0.compatibility_risk(&v1_1_0),
            CompatibilityRisk::Medium
        );
        assert_eq!(v1_0_0.compatibility_risk(&v2_0_0), CompatibilityRisk::High);
        Ok(())
    }

    #[test]
    fn test_version_manager() -> Result<(), MLError> {
        let manager = VersionManager::new();

        // Test compatibility check
        let compat_info = manager.check_compatibility("1.0.0", "1.1.0", ModelType::DQN)?;

        assert!(compat_info.compatible);
        assert_eq!(compat_info.risk, CompatibilityRisk::Medium);
        // Test incompatible versions
        let incompat_info = manager.check_compatibility("1.0.0", "2.0.0", ModelType::DQN)?;

        assert!(!incompat_info.compatible);
        assert_eq!(incompat_info.risk, CompatibilityRisk::High);
        Ok(())
    }

    #[test]
    fn test_version_suggestions() -> Result<(), MLError> {
        let manager = VersionManager::new();

        assert_eq!(
            manager.suggest_next_version("1.2.3", VersionChangeType::Patch)?,
            "1.2.4"
        );

        assert_eq!(
            manager.suggest_next_version("1.2.3", VersionChangeType::Minor)?,
            "1.3.0"
        );

        assert_eq!(
            manager.suggest_next_version("1.2.3", VersionChangeType::Major)?,
            "2.0.0"
        );
        Ok(())
    }

    #[test]
    fn test_migration_path() -> Result<(), MLError> {
        let manager = VersionManager::new();

        let path = manager.get_migration_path(ModelType::MAMBA, "1.0.0", "2.0.0")?;

        assert_eq!(path.len(), 1);
        assert_eq!(path[0].migration_type, MigrationType::Major);
        assert!(path[0].required);
        Ok(())
    }
}