use migration_helpers::{migrate, Migration, MigrationData, Result};
use std::process;

const MPS_PREFIX: &str = "settings.kubelet-device-plugins.nvidia.mps";
const DEVICE_SHARING_STRATEGY_SETTING: &str =
    "settings.kubelet-device-plugins.nvidia.device-sharing-strategy";

pub struct MpsMigration;

impl Migration for MpsMigration {
    fn forward(&mut self, input: MigrationData) -> Result<MigrationData> {
        println!("MpsMigration has no work to do on upgrade.");
        Ok(input)
    }

    fn backward(&mut self, mut input: MigrationData) -> Result<MigrationData> {
        // Remove all settings with the MPS prefix
        let to_remove: Vec<_> = input.data.keys()
            .filter(|k| k.starts_with(MPS_PREFIX))
            .cloned()
            .collect();
        for key in to_remove {
            if let Some(data) = input.data.remove(&key) {
                println!("Removed {key}, which was set to '{data}'");
            }
        }

        // Change device-sharing-strategy from "mps" to "none"
        if let Some(data) = input.data.get_mut(DEVICE_SHARING_STRATEGY_SETTING) {
            if let serde_json::Value::String(s) = data {
                if s == "mps" {
                    *data = serde_json::Value::String("none".to_string());
                    println!("Changed device-sharing-strategy from 'mps' to 'none' on downgrade.");
                }
            }
        }
        Ok(input)
    }
}

fn run() -> Result<()> {
    migrate(MpsMigration)
}

fn main() {
    if let Err(e) = run() {
        eprintln!("{e}");
        process::exit(1);
    }
}