interactive_print.rs

use std::sync::atomic::{AtomicBool, Ordering};

use anyhow::{Context, Result};
use ast_grep_config::RuleConfig;

use super::{ColoredPrinter, Diff, Printer};
use crate::error::ErrorContext as EC;
use crate::interaction;
use ast_grep_core::NodeMatch;
use ast_grep_language::SupportLang;

pub use codespan_reporting::{files::SimpleFile, term::ColorArg};

use std::borrow::Cow;
use std::path::{Path, PathBuf};

// add this macro because neither trait_alias nor type_alias_impl is supported.
macro_rules! Matches {
  ($lt: lifetime) => { impl Iterator<Item = NodeMatch<$lt, SupportLang>> };
}
macro_rules! Diffs {
  ($lt: lifetime) => { impl Iterator<Item = Diff<$lt>> };
}

static ACCEPT_ALL: AtomicBool = AtomicBool::new(false);

pub struct InteractivePrinter {
  accept_all: bool,
  inner: ColoredPrinter,
}
impl InteractivePrinter {
  pub fn new(inner: ColoredPrinter) -> Self {
    Self {
      accept_all: false,
      inner,
    }
  }

  pub fn accept_all(self, accept_all: bool) -> Self {
    Self { accept_all, ..self }
  }
}

impl Printer for InteractivePrinter {
  fn before_print(&self) {
    ACCEPT_ALL.store(self.accept_all, Ordering::SeqCst);
  }
  fn print_rule<'a>(
    &self,
    matches: Matches!('a),
    file: SimpleFile<Cow<str>, &String>,
    rule: &RuleConfig<SupportLang>,
  ) {
    interaction::run_in_alternate_screen(|| {
      self.inner.print_rule(matches, file, rule);
      let resp = interaction::prompt(VIEW_PROMPT, "q", Some('\n')).expect("cannot fail");
      if resp == 'q' {
        Err(anyhow::anyhow!("Exit interactive editing"))
      } else {
        Ok(())
      }
    })
    .unwrap();
  }

  fn print_matches<'a>(&self, matches: Matches!('a), path: &Path) -> Result<()> {
    interaction::run_in_alternate_screen(|| {
      print_matches_and_confirm_next(&self.inner, matches, path)
    })
  }

  fn print_diffs<'a>(&self, diffs: Diffs!('a), path: &Path) -> Result<()> {
    interaction::run_in_alternate_screen(|| {
      print_diffs_and_prompt_action(&self.inner, &path.to_path_buf(), diffs)
    })
  }
  fn print_rule_diffs<'a>(
    &self,
    diffs: Diffs!('a),
    path: &Path,
    _rule: &RuleConfig<SupportLang>,
  ) -> Result<()> {
    self.print_diffs(diffs, path)
  }
}

const EDIT_PROMPT: &str = "Accept change? (Yes[y], No[n], Accept All[a], Quit[q], Edit[e])";
const VIEW_PROMPT: &str = "Next[enter], Quit[q]";

fn print_diffs_and_prompt_action<'a>(
  printer: &impl Printer,
  path: &PathBuf,
  diffs: Diffs!('a),
) -> Result<()> {
  let diffs: Vec<_> = diffs.collect();
  let rewrite_action = || {
    let new_content = apply_rewrite(diffs.clone().into_iter());
    std::fs::write(path, new_content).with_context(|| EC::WriteFile(path.clone()))?;
    Ok(())
  };
  if ACCEPT_ALL.load(Ordering::SeqCst) {
    return rewrite_action();
  }
  let first_match = match diffs.first() {
    Some(n) => n.node_match.start_pos().0,
    None => return Ok(()),
  };
  printer.print_diffs(diffs.clone().into_iter(), path)?;
  let response =
    interaction::prompt(EDIT_PROMPT, "ynaqe", Some('n')).expect("Error happened during prompt");
  match response {
    'y' => rewrite_action(),
    'a' => {
      ACCEPT_ALL.store(true, Ordering::SeqCst);
      rewrite_action()
    }
    'n' => Ok(()),
    'e' => interaction::open_in_editor(path, first_match),
    'q' => Err(anyhow::anyhow!("Exit interactive editing")),
    _ => Ok(()),
  }
}

fn print_matches_and_confirm_next<'a>(
  printer: &impl Printer,
  matches: Matches!('a),
  path: &Path,
) -> Result<()> {
  printer.print_matches(matches, path)?;
  let resp = interaction::prompt(VIEW_PROMPT, "q", Some('\n')).expect("cannot fail");
  if resp == 'q' {
    Err(anyhow::anyhow!("Exit interactive editing"))
  } else {
    Ok(())
  }
}

fn apply_rewrite<'a>(diffs: Diffs!('a)) -> String {
  let mut diffs = diffs.peekable();
  let mut new_content = String::new();
  let Some(first) = diffs.peek() else {
    return new_content
  };
  let old_content = first.node_match.ancestors().last().unwrap().text();
  let mut start = 0;
  for diff in diffs {
    let range = diff.node_match.range();
    new_content.push_str(&old_content[start..range.start]);
    new_content.push_str(&diff.replacement);
    start = range.end;
  }
  // add trailing statements
  new_content.push_str(&old_content[start..]);
  new_content
}

#[cfg(test)]
mod test {
  use super::*;
  use ast_grep_config::from_yaml_string;
  use ast_grep_core::{AstGrep, Matcher, Pattern};

  fn make_rule(rule: &str) -> RuleConfig<SupportLang> {
    from_yaml_string(&format!(
      r"
id: test
message: test rule
severity: info
language: TypeScript
{rule}"
    ))
    .unwrap()
    .pop()
    .unwrap()
  }

  fn make_diffs<'a>(
    grep: &'a AstGrep<SupportLang>,
    matcher: impl Matcher<SupportLang>,
    fixer: &Pattern<SupportLang>,
  ) -> Vec<Diff<'a>> {
    let root = grep.root();
    root
      .find_all_without_nesting(&matcher)
      .into_iter()
      .map(|nm| Diff::generate(nm, &matcher, fixer))
      .collect()
  }

  #[test]
  fn test_apply_rewrite() {
    let root = AstGrep::new("let a = () => c++", SupportLang::TypeScript);
    let config = make_rule(
      r"
rule:
  all:
    - pattern: $B
    - any:
        - pattern: $A++
fix: ($B, lifecycle.update(['$A']))",
    );
    let matcher = config.get_matcher();
    let fixer = config.get_fixer().unwrap();
    let diffs = make_diffs(&root, matcher, &fixer);
    let ret = apply_rewrite(diffs.into_iter());
    assert_eq!(ret, "let a = () => (c++, lifecycle.update(['c']))");
  }

  #[test]
  fn test_rewrite_nested() {
    let root = AstGrep::new("Some(Some(1))", SupportLang::TypeScript);
    let diffs = make_diffs(
      &root,
      "Some($A)",
      &Pattern::new("$A", SupportLang::TypeScript),
    );
    let ret = apply_rewrite(diffs.into_iter());
    assert_eq!("Some(1)", ret);
  }
}