vsec 0.0.1

// src/parser/ast_visitor.rs

use std::collections::HashMap;
use std::path::PathBuf;

use syn::visit::Visit;
use syn::{Expr, File, Item, ItemConst, ItemFn, ItemMod, ItemStatic};

use crate::models::Constant;
use crate::parser::context::{ParseContext, ScopeEntry, SymbolInfo, SymbolKind};

/// A local literal binding found in a let statement
#[derive(Debug, Clone)]
pub struct LocalLiteral {
    /// Variable name
    pub name: String,
    /// The string value
    pub value: String,
    /// Line number where defined
    pub line: u32,
    /// Function scope (if any)
    pub in_function: Option<String>,
}

/// A basic visitor for extracting information from Rust source files
pub struct SecretVisitor {
    /// Path to the file being visited
    file_path: PathBuf,

    /// Parse context for tracking scope
    context: ParseContext,

    /// Collected constants
    constants: Vec<Constant>,

    /// String values found (inline literals in comparisons)
    string_literals: Vec<(String, u32)>,

    /// Local constant map for resolution
    local_constants: HashMap<String, String>,

    /// Local let bindings with literal values (for tracking within functions)
    local_literals: Vec<LocalLiteral>,

    /// Current function name (for scoping)
    current_function: Option<String>,
}

impl SecretVisitor {
    pub fn new(file_path: PathBuf) -> Self {
        Self {
            file_path,
            context: ParseContext::new(),
            constants: Vec::new(),
            string_literals: Vec::new(),
            local_constants: HashMap::new(),
            local_literals: Vec::new(),
            current_function: None,
        }
    }

    /// Visit a parsed file
    pub fn visit(&mut self, file: &File) {
        // First pass: collect constants
        self.collect_constants(file);

        // Second pass: visit all items
        for item in &file.items {
            syn::visit::visit_item(self, item);
        }
    }

    fn collect_constants(&mut self, file: &File) {
        for item in &file.items {
            match item {
                Item::Const(c) => {
                    if let Some(value) = Self::extract_string_value(&c.expr) {
                        self.local_constants.insert(c.ident.to_string(), value);
                    }
                }
                Item::Static(s) => {
                    if let Some(value) = Self::extract_string_value(&s.expr) {
                        self.local_constants.insert(s.ident.to_string(), value);
                    }
                }
                _ => {}
            }
        }
    }

    /// Extract string value from an expression
    /// Handles method chains like "literal".to_string().parse().unwrap()
    pub fn extract_string_value(expr: &Expr) -> Option<String> {
        match expr {
            Expr::Lit(lit) => match &lit.lit {
                syn::Lit::Str(s) => Some(s.value()),
                syn::Lit::ByteStr(s) => String::from_utf8(s.value()).ok(),
                _ => None,
            },
            // Handle method call chains: "literal".to_string(), "literal".parse().unwrap(), etc.
            Expr::MethodCall(call) => {
                let method = call.method.to_string();
                // Methods that pass through the string value
                let passthrough_methods = [
                    "to_string",
                    "to_owned",
                    "into",
                    "parse",
                    "unwrap",
                    "unwrap_or",
                    "unwrap_or_default",
                    "expect",
                    "ok",
                    "as_str",
                    "as_ref",
                    "clone",
                    "trim",
                    "trim_start",
                    "trim_end",
                    "to_lowercase",
                    "to_uppercase",
                    "to_ascii_lowercase",
                    "to_ascii_uppercase",
                ];
                if passthrough_methods.contains(&method.as_str()) {
                    // Recursively extract from the receiver
                    Self::extract_string_value(&call.receiver)
                } else {
                    None
                }
            }
            // Handle try expressions: "literal".parse()?
            Expr::Try(try_expr) => Self::extract_string_value(&try_expr.expr),
            Expr::Reference(r) => Self::try_decode_referenced_value(&r.expr),
            Expr::Macro(mac) => {
                if mac
                    .mac
                    .path
                    .segments
                    .last()
                    .map(|s| s.ident == "vec")
                    .unwrap_or(false)
                {
                    Self::try_decode_vec_macro(&mac.mac.tokens)
                } else {
                    None
                }
            }
            Expr::Call(call) => {
                let func_name = quote::quote!(#call.func).to_string();
                if func_name.contains("from_utf8") {
                    call.args.first().and_then(Self::extract_string_value)
                } else if func_name.contains("String :: from") || func_name.contains("String::from")
                {
                    // Handle String::from("literal")
                    call.args.first().and_then(Self::extract_string_value)
                } else {
                    None
                }
            }
            Expr::Group(g) => Self::extract_string_value(&g.expr),
            Expr::Paren(p) => Self::extract_string_value(&p.expr),
            _ => None,
        }
    }

    /// Try to decode byte/char arrays
    fn try_decode_referenced_value(expr: &Expr) -> Option<String> {
        if let Expr::Array(arr) = expr {
            // Try to decode as byte values
            let bytes: Option<Vec<u8>> = arr
                .elems
                .iter()
                .map(|e| {
                    if let Expr::Lit(lit) = e {
                        match &lit.lit {
                            syn::Lit::Int(i) => i.base10_parse::<u8>().ok(),
                            syn::Lit::Byte(b) => Some(b.value()),
                            _ => None,
                        }
                    } else {
                        None
                    }
                })
                .collect();

            if let Some(bytes) = bytes {
                if let Ok(s) = String::from_utf8(bytes) {
                    return Some(s);
                }
            }

            // Try to decode as char values
            let chars: Option<String> = arr
                .elems
                .iter()
                .map(|e| {
                    if let Expr::Lit(lit) = e {
                        if let syn::Lit::Char(c) = &lit.lit {
                            return Some(c.value());
                        }
                    }
                    None
                })
                .collect();

            return chars;
        }
        None
    }

    /// Try to decode vec![...] with byte literals
    fn try_decode_vec_macro(tokens: &proc_macro2::TokenStream) -> Option<String> {
        use syn::parse::Parser;
        use syn::{Expr, ExprLit, Lit};

        let parser = syn::punctuated::Punctuated::<Expr, syn::Token![,]>::parse_terminated;
        let exprs: syn::punctuated::Punctuated<Expr, syn::Token![,]> =
            match parser.parse2(tokens.clone()) {
                Ok(exprs) => exprs,
                Err(_) => return None,
            };

        let bytes: Option<Vec<u8>> = exprs
            .iter()
            .map(|expr| {
                if let Expr::Lit(ExprLit {
                    lit: Lit::Int(int), ..
                }) = expr
                {
                    int.base10_parse::<u8>().ok()
                } else if let Expr::Lit(ExprLit {
                    lit: Lit::Byte(b), ..
                }) = expr
                {
                    Some(b.value())
                } else {
                    None
                }
            })
            .collect();

        bytes.and_then(|b| String::from_utf8(b).ok())
    }

    /// Get collected constants
    pub fn constants(&self) -> &[Constant] {
        &self.constants
    }

    /// Get collected string literals
    pub fn string_literals(&self) -> &[(String, u32)] {
        &self.string_literals
    }

    /// Get local constants map
    pub fn local_constants(&self) -> &HashMap<String, String> {
        &self.local_constants
    }

    /// Get local literals (let bindings with string values)
    pub fn local_literals(&self) -> &[LocalLiteral] {
        &self.local_literals
    }

    /// Take ownership of results
    pub fn into_results(
        self,
    ) -> (
        Vec<Constant>,
        Vec<(String, u32)>,
        HashMap<String, String>,
        Vec<LocalLiteral>,
    ) {
        (
            self.constants,
            self.string_literals,
            self.local_constants,
            self.local_literals,
        )
    }
}

impl<'ast> Visit<'ast> for SecretVisitor {
    fn visit_item_const(&mut self, node: &'ast ItemConst) {
        if let Some(value) = Self::extract_string_value(&node.expr) {
            let visibility = crate::models::Visibility::from_syn(&node.vis);

            let constant = Constant::new(
                node.ident.to_string(),
                value.clone(),
                self.file_path.clone(),
                node.ident.span().start().line as u32,
            )
            .with_visibility(visibility);

            self.constants.push(constant);

            // Register in context
            self.context.register_symbol(
                node.ident.to_string(),
                SymbolInfo {
                    name: node.ident.to_string(),
                    value: Some(value),
                    kind: SymbolKind::Constant,
                    line: node.ident.span().start().line as u32,
                },
            );
        }

        syn::visit::visit_item_const(self, node);
    }

    fn visit_item_static(&mut self, node: &'ast ItemStatic) {
        if let Some(value) = Self::extract_string_value(&node.expr) {
            let visibility = crate::models::Visibility::from_syn(&node.vis);

            let constant = Constant::new(
                node.ident.to_string(),
                value.clone(),
                self.file_path.clone(),
                node.ident.span().start().line as u32,
            )
            .with_visibility(visibility)
            .as_static();

            self.constants.push(constant);

            self.context.register_symbol(
                node.ident.to_string(),
                SymbolInfo {
                    name: node.ident.to_string(),
                    value: Some(value),
                    kind: SymbolKind::Static,
                    line: node.ident.span().start().line as u32,
                },
            );
        }

        syn::visit::visit_item_static(self, node);
    }

    fn visit_item_fn(&mut self, node: &'ast ItemFn) {
        let name = node.sig.ident.to_string();

        // Check if this is a test function
        let is_test = node.attrs.iter().any(|attr| {
            attr.path().is_ident("test")
                || attr
                    .path()
                    .segments
                    .last()
                    .map(|s| s.ident == "test")
                    .unwrap_or(false)
        });

        if is_test {
            self.context.push_scope(ScopeEntry::Test);
        }
        self.context.push_scope(ScopeEntry::Function(name.clone()));

        // Track current function for local literal scoping
        let old_function = self.current_function.take();
        self.current_function = Some(name);

        syn::visit::visit_item_fn(self, node);

        self.current_function = old_function;
        self.context.pop_scope();
        if is_test {
            self.context.pop_scope();
        }
    }

    fn visit_local(&mut self, node: &'ast syn::Local) {
        // Track let bindings with string literal values
        // e.g., let token = "rustfs rpc".parse().unwrap();
        if let Some(init) = &node.init {
            if let Some(value) = Self::extract_string_value(&init.expr) {
                // Extract variable name from the pattern
                if let syn::Pat::Ident(pat_ident) = &node.pat {
                    let name = pat_ident.ident.to_string();
                    let line = pat_ident.ident.span().start().line as u32;

                    self.local_literals.push(LocalLiteral {
                        name: name.clone(),
                        value: value.clone(),
                        line,
                        in_function: self.current_function.clone(),
                    });

                    // Also register in context for lookup
                    self.context.register_symbol(
                        name.clone(),
                        SymbolInfo {
                            name,
                            value: Some(value),
                            kind: SymbolKind::Local,
                            line,
                        },
                    );
                }
            }
        }

        syn::visit::visit_local(self, node);
    }

    fn visit_item_mod(&mut self, node: &'ast ItemMod) {
        let name = node.ident.to_string();

        // Check if this is a test module
        let is_test = name == "tests"
            || name == "test"
            || node.attrs.iter().any(|attr| {
                if attr.path().is_ident("cfg") {
                    if let Ok(meta) = attr.meta.require_list() {
                        return meta.tokens.to_string().contains("test");
                    }
                }
                false
            });

        if is_test {
            self.context.push_scope(ScopeEntry::Test);
        }
        self.context.push_scope(ScopeEntry::Module(name));

        syn::visit::visit_item_mod(self, node);

        self.context.pop_scope();
        if is_test {
            self.context.pop_scope();
        }
    }

    fn visit_expr_lit(&mut self, node: &'ast syn::ExprLit) {
        if let syn::Lit::Str(s) = &node.lit {
            let value = s.value();
            let line = s.span().start().line as u32;
            self.string_literals.push((value, line));
        }

        syn::visit::visit_expr_lit(self, node);
    }
}

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

    #[test]
    fn test_extract_string_literal() {
        let expr: syn::Expr = syn::parse_quote!("hello");
        assert_eq!(SecretVisitor::extract_string_value(&expr), Some("hello".into()));
    }

    #[test]
    fn test_extract_byte_string() {
        let expr: syn::Expr = syn::parse_quote!(b"hello");
        assert_eq!(SecretVisitor::extract_string_value(&expr), Some("hello".into()));
    }

    #[test]
    fn test_visitor_collects_constants() {
        let code = r#"
            const API_KEY: &str = "secret123";
            static DB_URL: &str = "localhost";
        "#;

        let file: syn::File = syn::parse_str(code).unwrap();
        let mut visitor = SecretVisitor::new(PathBuf::from("test.rs"));
        visitor.visit(&file);

        assert_eq!(visitor.constants().len(), 2);
        assert!(visitor.local_constants().contains_key("API_KEY"));
        assert!(visitor.local_constants().contains_key("DB_URL"));
    }

    #[test]
    fn test_extract_method_chain_to_string() {
        let expr: syn::Expr = syn::parse_quote!("secret".to_string());
        assert_eq!(
            SecretVisitor::extract_string_value(&expr),
            Some("secret".into())
        );
    }

    #[test]
    fn test_extract_method_chain_parse_unwrap() {
        // This is the RustFS pattern: "rustfs rpc".parse().unwrap()
        let expr: syn::Expr = syn::parse_quote!("rustfs rpc".parse().unwrap());
        assert_eq!(
            SecretVisitor::extract_string_value(&expr),
            Some("rustfs rpc".into())
        );
    }

    #[test]
    fn test_extract_string_from() {
        let expr: syn::Expr = syn::parse_quote!(String::from("secret"));
        assert_eq!(
            SecretVisitor::extract_string_value(&expr),
            Some("secret".into())
        );
    }

    #[test]
    fn test_visitor_collects_local_literals() {
        let code = r#"
            fn authenticate(t: &str) -> bool {
                let token = "rustfs rpc".parse().unwrap();
                t == token
            }
        "#;

        let file: syn::File = syn::parse_str(code).unwrap();
        let mut visitor = SecretVisitor::new(PathBuf::from("test.rs"));
        visitor.visit(&file);

        assert_eq!(visitor.local_literals().len(), 1);
        let local = &visitor.local_literals()[0];
        assert_eq!(local.name, "token");
        assert_eq!(local.value, "rustfs rpc");
        assert_eq!(local.in_function, Some("authenticate".into()));
    }

    #[test]
    fn test_visitor_collects_local_literal_simple() {
        let code = r#"
            fn check_auth(input: &str) -> bool {
                let secret = "hardcoded_password";
                input == secret
            }
        "#;

        let file: syn::File = syn::parse_str(code).unwrap();
        let mut visitor = SecretVisitor::new(PathBuf::from("test.rs"));
        visitor.visit(&file);

        assert_eq!(visitor.local_literals().len(), 1);
        let local = &visitor.local_literals()[0];
        assert_eq!(local.name, "secret");
        assert_eq!(local.value, "hardcoded_password");
    }
}