flake-edit 0.3.4

use std::collections::HashMap;

use crate::change::Change;
use crate::error::FlakeEditError;
use crate::input::{Follows, Input};
use crate::validate;
use crate::walk::Walker;

pub struct FlakeEdit {
    changes: Vec<Change>,
    walker: Walker,
}

#[derive(Default, Debug)]
pub enum Outputs {
    #[default]
    None,
    Multiple(Vec<String>),
    Any(Vec<String>),
}

pub type InputMap = HashMap<String, Input>;

pub fn sorted_input_ids(inputs: &InputMap) -> Vec<&String> {
    let mut keys: Vec<_> = inputs.keys().collect();
    keys.sort();
    keys
}

/// Returns owned strings, for contexts where references won't work.
pub fn sorted_input_ids_owned(inputs: &InputMap) -> Vec<String> {
    let mut keys: Vec<String> = inputs.keys().cloned().collect();
    keys.sort();
    keys
}

#[derive(Default, Debug)]
pub enum OutputChange {
    #[default]
    None,
    Add(String),
    Remove(String),
}

impl FlakeEdit {
    pub fn new(changes: Vec<Change>, walker: Walker) -> Self {
        Self { changes, walker }
    }

    pub fn from_text(stream: &str) -> Result<Self, FlakeEditError> {
        let validation = validate::validate(stream);
        if validation.has_errors() {
            return Err(FlakeEditError::Validation(validation.errors));
        }

        let walker = Walker::new(stream);
        Ok(Self::new(Vec::new(), walker))
    }

    pub fn changes(&self) -> &[Change] {
        self.changes.as_ref()
    }

    pub fn add_change(&mut self, change: Change) {
        self.changes.push(change);
    }

    pub fn source_text(&self) -> String {
        self.walker.root.to_string()
    }

    pub fn curr_list(&self) -> &InputMap {
        &self.walker.inputs
    }

    /// Will walk and then list the inputs, for listing the current inputs,
    /// use `curr_list()`.
    pub fn list(&mut self) -> &InputMap {
        self.walker.inputs.clear();
        // Walk returns Ok(None) when no changes are made (expected for listing)
        assert!(self.walker.walk(&Change::None).ok().flatten().is_none());
        &self.walker.inputs
    }
    /// Apply a specific change to a walker, on some inputs it will need to walk
    /// multiple times, will error, if the edit could not be applied successfully.
    pub fn apply_change(&mut self, change: Change) -> Result<Option<String>, FlakeEditError> {
        match change {
            Change::None => Ok(None),
            Change::Add { .. } => {
                // Check for duplicate input before adding
                if let Some(input_id) = change.id() {
                    self.ensure_inputs_populated()?;

                    let input_id_string = input_id.to_string();
                    if self.walker.inputs.contains_key(&input_id_string) {
                        return Err(FlakeEditError::DuplicateInput(input_id_string));
                    }
                }

                if let Some(maybe_changed_node) = self.walker.walk(&change.clone())? {
                    let outputs = self.walker.list_outputs()?;
                    match outputs {
                        Outputs::Multiple(out) => {
                            let id = change.id().unwrap().to_string();
                            if !out.contains(&id) {
                                self.walker.root = maybe_changed_node.clone();
                                if let Some(maybe_changed_node) =
                                    self.walker.change_outputs(OutputChange::Add(id))?
                                {
                                    return Ok(Some(maybe_changed_node.to_string()));
                                }
                            }
                        }
                        Outputs::None | Outputs::Any(_) => {}
                    }
                    Ok(Some(maybe_changed_node.to_string()))
                } else {
                    self.walker.add_toplevel = true;
                    let maybe_changed_node = self.walker.walk(&change)?;
                    Ok(maybe_changed_node.map(|n| n.to_string()))
                }
            }
            Change::Remove { .. } => {
                self.ensure_inputs_populated()?;

                let removed_id = change.id().unwrap().to_string();

                // If we remove a node, it could be a flat structure,
                // we want to remove all of the references to its toplevel.
                let mut res = None;
                while let Some(changed_node) = self.walker.walk(&change)? {
                    if res == Some(changed_node.clone()) {
                        break;
                    }
                    res = Some(changed_node.clone());
                    self.walker.root = changed_node.clone();
                }
                // Removed nodes should be removed from the outputs
                let outputs = self.walker.list_outputs()?;
                match outputs {
                    Outputs::Multiple(out) | Outputs::Any(out) => {
                        if out.contains(&removed_id)
                            && let Some(changed_node) = self
                                .walker
                                .change_outputs(OutputChange::Remove(removed_id.clone()))?
                        {
                            res = Some(changed_node.clone());
                            self.walker.root = changed_node.clone();
                        }
                    }
                    Outputs::None => {}
                }

                // Remove orphaned follows references that point to the removed input
                let orphaned_follows = self.collect_orphaned_follows(&removed_id);
                for orphan_change in orphaned_follows {
                    while let Some(changed_node) = self.walker.walk(&orphan_change)? {
                        if res == Some(changed_node.clone()) {
                            break;
                        }
                        res = Some(changed_node.clone());
                        self.walker.root = changed_node.clone();
                    }
                }

                Ok(res.map(|n| n.to_string()))
            }
            Change::Follows { ref input, .. } => {
                self.ensure_inputs_populated()?;

                let parent_id = input.input();
                if !self.walker.inputs.contains_key(parent_id) {
                    return Err(FlakeEditError::InputNotFound(parent_id.to_string()));
                }

                if let Some(maybe_changed_node) = self.walker.walk(&change)? {
                    Ok(Some(maybe_changed_node.to_string()))
                } else {
                    Ok(None)
                }
            }
            Change::Change { .. } => {
                if let Some(input_id) = change.id() {
                    self.ensure_inputs_populated()?;

                    let input_id_string = input_id.to_string();
                    if !self.walker.inputs.contains_key(&input_id_string) {
                        return Err(FlakeEditError::InputNotFound(input_id_string));
                    }
                }

                if let Some(maybe_changed_node) = self.walker.walk(&change)? {
                    Ok(Some(maybe_changed_node.to_string()))
                } else {
                    Ok(None)
                }
            }
        }
    }

    pub fn walker(&self) -> &Walker {
        &self.walker
    }

    /// Ensure the inputs map is populated by walking if empty.
    fn ensure_inputs_populated(&mut self) -> Result<(), FlakeEditError> {
        if self.walker.inputs.is_empty() {
            let _ = self.walker.walk(&Change::None)?;
        }
        Ok(())
    }

    /// Collect follows references that point to a removed input.
    /// Returns a list of Change::Remove for orphaned follows.
    fn collect_orphaned_follows(&self, removed_id: &str) -> Vec<Change> {
        let mut orphaned = Vec::new();
        for (input_id, input) in &self.walker.inputs {
            for follows in input.follows() {
                if let Follows::Indirect(follows_name, target) = follows {
                    // target is the RHS of `follows = "target"`
                    if target.trim_matches('"') == removed_id {
                        let nested_id = format!("{}.{}", input_id, follows_name);
                        orphaned.push(Change::Remove {
                            ids: vec![nested_id.into()],
                        });
                    }
                }
            }
        }
        orphaned
    }
}

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

    #[test]
    fn already_follows_is_noop() {
        let flake = r#"{
  inputs = {
    nixpkgs.url = "github:nixos/nixpkgs";
    crane = {
      url = "github:ipetkov/crane";
      inputs.nixpkgs.follows = "nixpkgs";
    };
  };
  outputs = { ... }: { };
}"#;
        let mut fe = FlakeEdit::from_text(flake).unwrap();
        let original = fe.source_text();
        let change = Change::Follows {
            input: "crane.nixpkgs".to_string().into(),
            target: "nixpkgs".to_string(),
        };
        let result = fe.apply_change(change).unwrap();
        // Walker returns the unchanged text; caller detects the no-op
        match result {
            Some(text) => assert_eq!(text, original, "text should be unchanged"),
            None => {} // also acceptable
        }
    }

    #[test]
    fn new_follows_succeeds() {
        let flake = r#"{
  inputs = {
    nixpkgs.url = "github:nixos/nixpkgs";
    crane = {
      url = "github:ipetkov/crane";
    };
  };
  outputs = { ... }: { };
}"#;
        let mut fe = FlakeEdit::from_text(flake).unwrap();
        let change = Change::Follows {
            input: "crane.nixpkgs".to_string().into(),
            target: "nixpkgs".to_string(),
        };
        let result = fe.apply_change(change);
        assert!(result.is_ok(), "expected Ok, got: {:?}", result);
        let text = result.unwrap().unwrap();
        assert!(text.contains("inputs.nixpkgs.follows = \"nixpkgs\""));
    }
}