vibe-action 0.1.4

//! AST modifier — parses source code into structured JSON via vibe-ast.

use anyhow::Result;
use vibe_ast::ArkTSNode;
use vibe_ast::BashNode;
use vibe_ast::BatchNode;
use vibe_ast::CSharpNode;
use vibe_ast::DartNode;
use vibe_ast::GoNode;
use vibe_ast::JavaNode;
use vibe_ast::JavaScriptNode;
use vibe_ast::KotlinNode;
use vibe_ast::Language;
use vibe_ast::MarkdownNode;
use vibe_ast::PythonNode;
use vibe_ast::RustNode;
use vibe_ast::SwiftNode;
use vibe_ast::TypeScriptNode;
use vibe_ast::nodes::nodes::Nodes;
use vibe_ast::parse_file;
use vibe_ast::parse_text;

use super::modifier::Modifier;
use super::modifier::ModifierKey;
use crate::models::context::ContextModel;

pub struct AstModifier;

impl AstModifier {
    fn lang_from_arg(arg: &str) -> Result<Language> {
        match arg {
            "rs" => Ok(Language::Rust),
            "py" => Ok(Language::Python),
            "ts" => Ok(Language::TypeScript),
            "js" => Ok(Language::JavaScript),
            "java" => Ok(Language::Java),
            "go" => Ok(Language::Go),
            "cs" => Ok(Language::CSharp),
            "kt" => Ok(Language::Kotlin),
            "swift" => Ok(Language::Swift),
            "dart" => Ok(Language::Dart),
            "sh" => Ok(Language::Bash),
            "bat" => Ok(Language::Batch),
            "ets" => Ok(Language::ArkTS),
            "md" => Ok(Language::Markdown),
            _ => anyhow::bail!("Unknown language: {}", arg),
        }
    }

    /// Remove non-essential data from AST based on language.
    fn clean(nodes: &mut Nodes) {
        match nodes {
            Nodes::ArkTS(n) => Self::clean_arkts(n),
            Nodes::Bash(n) => Self::clean_bash(n),
            Nodes::Batch(n) => Self::clean_batch(n),
            Nodes::CSharp(n) => Self::clean_csharp(n),
            Nodes::Dart(n) => Self::clean_dart(n),
            Nodes::Go(n) => Self::clean_go(n),
            Nodes::Java(n) => Self::clean_java(n),
            Nodes::JavaScript(n) => Self::clean_javascript(n),
            Nodes::Kotlin(n) => Self::clean_kotlin(n),
            Nodes::Markdown(n) => Self::clean_markdown(n),
            Nodes::Python(n) => Self::clean_python(n),
            Nodes::Rust(n) => Self::clean_rust(n),
            Nodes::Swift(n) => Self::clean_swift(n),
            Nodes::TypeScript(n) => Self::clean_typescript(n),
        }
    }

    /// Remove non-essential data from ArkTS AST.
    fn clean_arkts(n: &mut ArkTSNode) {
        n.imports.clear();
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for f in &mut n.functions {
            f.body = None;
        }
        for c in &mut n.classes {
            c.body = None;
        }
        for s in &mut n.structs {
            s.body = None;
        }
        for e in &mut n.enums {
            e.body = None;
        }
        for i in &mut n.interfaces {
            i.body = None;
        }
    }

    /// Remove non-essential data from Bash AST.
    fn clean_bash(n: &mut BashNode) {
        n.imports.clear();
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for f in &mut n.functions {
            f.body = None;
        }
    }

    /// Remove non-essential data from Batch AST.
    fn clean_batch(n: &mut BatchNode) {
        n.imports.clear();
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for f in &mut n.functions {
            f.body = None;
        }
    }

    /// Remove non-essential data from C# AST.
    fn clean_csharp(n: &mut CSharpNode) {
        n.imports.clear();
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for f in &mut n.functions {
            f.body = None;
        }
        for c in &mut n.classes {
            c.body = None;
        }
        for s in &mut n.structs {
            s.body = None;
        }
        for e in &mut n.enums {
            e.body = None;
        }
        for i in &mut n.interfaces {
            i.body = None;
        }
    }

    /// Remove non-essential data from Dart AST.
    fn clean_dart(n: &mut DartNode) {
        n.imports.clear();
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for f in &mut n.functions {
            f.body = None;
        }
        for c in &mut n.classes {
            c.body = None;
        }
        for e in &mut n.enums {
            e.body = None;
        }
        for i in &mut n.interfaces {
            i.body = None;
        }
    }

    /// Remove non-essential data from Go AST.
    fn clean_go(n: &mut GoNode) {
        n.imports.clear();
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for f in &mut n.functions {
            f.body = None;
        }
        for s in &mut n.structs {
            s.body = None;
        }
        for e in &mut n.enums {
            e.body = None;
        }
    }

    /// Remove non-essential data from Java AST.
    fn clean_java(n: &mut JavaNode) {
        n.imports.clear();
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for f in &mut n.functions {
            f.body = None;
        }
        for c in &mut n.classes {
            c.body = None;
        }
        for s in &mut n.structs {
            s.body = None;
        }
        for e in &mut n.enums {
            e.body = None;
        }
        for i in &mut n.interfaces {
            i.body = None;
        }
    }

    /// Remove non-essential data from JavaScript AST.
    fn clean_javascript(n: &mut JavaScriptNode) {
        n.imports.clear();
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for f in &mut n.functions {
            f.body = None;
        }
        for c in &mut n.classes {
            c.body = None;
        }
    }

    /// Remove non-essential data from Kotlin AST.
    fn clean_kotlin(n: &mut KotlinNode) {
        n.imports.clear();
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for f in &mut n.functions {
            f.body = None;
        }
        for c in &mut n.classes {
            c.body = None;
        }
        for e in &mut n.enums {
            e.body = None;
        }
        for i in &mut n.interfaces {
            i.body = None;
        }
    }

    /// Remove non-essential data from Markdown AST.
    fn clean_markdown(n: &mut MarkdownNode) {
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for h in &mut n.headings {
            h.body = None;
        }
    }

    /// Remove non-essential data from Python AST.
    fn clean_python(n: &mut PythonNode) {
        n.imports.clear();
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for f in &mut n.functions {
            f.body = None;
        }
        for c in &mut n.classes {
            c.body = None;
        }
    }

    /// Remove non-essential data from Rust AST.
    fn clean_rust(n: &mut RustNode) {
        n.imports.clear();
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for f in &mut n.functions {
            f.body = None;
        }
        for s in &mut n.structs {
            s.body = None;
        }
        for e in &mut n.enums {
            e.body = None;
        }
        for i in &mut n.interfaces {
            i.body = None;
        }
    }

    /// Remove non-essential data from Swift AST.
    fn clean_swift(n: &mut SwiftNode) {
        n.imports.clear();
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for f in &mut n.functions {
            f.body = None;
        }
        for c in &mut n.classes {
            c.body = None;
        }
        for s in &mut n.structs {
            s.body = None;
        }
        for e in &mut n.enums {
            e.body = None;
        }
        for i in &mut n.interfaces {
            i.body = None;
        }
    }

    /// Remove non-essential data from TypeScript AST.
    fn clean_typescript(n: &mut TypeScriptNode) {
        n.imports.clear();
        n.tags.clear();
        n.markers.clear();
        n.warnings.clear();
        for f in &mut n.functions {
            f.body = None;
        }
        for c in &mut n.classes {
            c.body = None;
        }
        for s in &mut n.structs {
            s.body = None;
        }
        for e in &mut n.enums {
            e.body = None;
        }
        for i in &mut n.interfaces {
            i.body = None;
        }
    }
}

impl Modifier for AstModifier {
    fn key(&self) -> ModifierKey {
        ModifierKey::Ast
    }

    fn apply(&self, value: &ContextModel, arg: &str) -> Result<ContextModel> {
        match value {
            ContextModel::String(s) => {
                if arg.is_empty() || arg == "brief" {
                    match parse_file(s) {
                        Ok(mut nodes) => {
                            Self::clean(&mut nodes);
                            Ok(ContextModel::String(serde_json::to_string(&nodes)?))
                        }
                        Err(_) => Ok(ContextModel::String(String::new())),
                    }
                } else if arg == "full" {
                    match parse_file(s) {
                        Ok(nodes) => Ok(ContextModel::String(serde_json::to_string(&nodes)?)),
                        Err(_) => Ok(ContextModel::String(String::new())),
                    }
                } else {
                    let lang = Self::lang_from_arg(arg)?;
                    let nodes = parse_text(s, lang)?;
                    Ok(ContextModel::String(serde_json::to_string(&nodes)?))
                }
            }
            ContextModel::List(files) => {
                let results: Vec<String> = files
                    .iter()
                    .map(
                        |path| match self.apply(&ContextModel::String(path.clone()), arg)? {
                            ContextModel::String(json) => Ok(json),
                            _ => unreachable!(),
                        },
                    )
                    .collect::<Result<Vec<_>>>()?;
                Ok(ContextModel::List(results))
            }
        }
    }
}