forge_guard/parser/

mod.rs

//! Solidity source file parser — parses source code into structured representations
//! for deeper vulnerability analysis (CEI, access control, etc.).

/// Visibility of a function or state variable.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Visibility {
    Public,
    Internal,
    External,
    Private,
}

/// State mutability of a function.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Mutability {
    Pure,
    View,
    Nonpayable,
    Payable,
}

/// Classification of a single logical statement within a function body.
#[derive(Debug, Clone, PartialEq)]
pub enum StatementKind {
    /// An external call: .call{..., .transfer(), .send(), or calling an interface method
    ExternalCall,
    /// A state write: assignment (=, +=, -=, etc.), delete, push, pop
    StateWrite,
    /// A state read: reading a state variable
    StateRead,
    /// A require() or revert() — flow control
    Guard,
    /// An if/while/for condition — flow control
    FlowControl,
    /// An internal function call
    InternalCall,
    /// An emit statement
    Emit,
    /// An event definition (inside contract, not function)
    EventDef,
    /// Inner assembly block
    InlineAssembly,
    /// Other / unclassified
    Other,
}

/// A statement with its content, line number, and kind.
#[derive(Debug, Clone)]
pub struct Statement {
    pub kind: StatementKind,
    pub line: usize,
    pub text: String,
}

/// A single state variable definition.
#[derive(Debug, Clone)]
pub struct StateVariable {
    pub name: String,
    pub type_name: String,
    pub visibility: Visibility,
    pub line: usize,
}

/// A parameter in a function or event.
#[derive(Debug, Clone)]
pub struct Param {
    pub name: String,
    pub type_name: String,
    pub indexed: bool,
}

/// A full function definition with parsed body.
#[derive(Debug, Clone)]
pub struct FunctionDef {
    pub name: String,
    pub visibility: Visibility,
    pub mutability: Mutability,
    pub modifiers: Vec<String>,
    pub params: Vec<Param>,
    pub return_params: Vec<Param>,
    pub body: Vec<Statement>,
    pub line: usize,
    pub is_constructor: bool,
    pub is_fallback: bool,
    pub is_receive: bool,
}

impl FunctionDef {
    /// Whether this function can modify state.
    pub fn modifies_state(&self) -> bool {
        matches!(
            self.mutability,
            Mutability::Nonpayable | Mutability::Payable
        )
    }

    /// Whether this function has a reentrancy guard modifier.
    pub fn has_reentrancy_guard(&self) -> bool {
        self.modifiers.iter().any(|m| {
            let m = m.to_lowercase();
            m == "nonreentrant"
                || m == "reentrancyguard"
                || m.contains("nonreentrant")
                || m == "reentrant"
        })
    }

    /// Whether this function has any access-control modifiers.
    pub fn has_access_control(&self, contract: &Contract) -> bool {
        if self.modifiers.is_empty() {
            return false;
        }
        for modifier in &self.modifiers {
            let m = modifier.to_lowercase();
            // Standard access control modifiers
            if m.contains("only")
                || m == "auth"
                || m.contains("role")
                || m == "whennotpaused"
                || m == "whenpaused"
                || m == "nonreentrant"
            {
                return true;
            }
            // Check if modifier is defined in the contract (could be custom access control)
            if contract.has_modifier_named(modifier) {
                return true;
            }
        }
        false
    }
}

/// A modifier definition.
#[derive(Debug, Clone)]
pub struct ModifierDef {
    pub name: String,
    pub params: Vec<Param>,
    pub body: Vec<Statement>,
    pub line: usize,
}

/// An event definition at the contract level.
#[derive(Debug, Clone)]
pub struct EventDef {
    pub name: String,
    pub params: Vec<Param>,
    pub line: usize,
}

/// A contract, interface, or library definition.
#[derive(Debug, Clone)]
pub struct Contract {
    pub name: String,
    pub inheritance: Vec<String>,
    pub state_variables: Vec<StateVariable>,
    pub functions: Vec<FunctionDef>,
    pub modifiers_defs: Vec<ModifierDef>,
    pub events: Vec<EventDef>,
    pub kind: ContractKind,
    pub line: usize,
    pub structs: Vec<StructDef>,
    pub errors: Vec<ErrorDef>,
    pub using_for: Vec<UsingForDef>,
}

impl Contract {
    /// Check if contract (or its inheritance chain) likely has a specific capability.
    pub fn inherits_access_control(&self) -> bool {
        self.inheritance.iter().any(|i| {
            let l = i.to_lowercase();
            l.contains("ownable")
                || l.contains("accesscontrol")
                || l.contains("ownableupgradeable")
                || l.contains("accesscontrolupgradeable")
                || l.contains("auth")
        })
    }

    /// Check if a modifier with the given name exists in this contract.
    pub fn has_modifier_named(&self, name: &str) -> bool {
        self.modifiers_defs.iter().any(|m| m.name == name)
    }

    /// Get a modifier by name.
    pub fn get_modifier(&self, name: &str) -> Option<&ModifierDef> {
        self.modifiers_defs.iter().find(|m| m.name == name)
    }

    /// Check if the contract has OpenZeppelin-style ReentrancyGuard.
    pub fn has_reentrancy_guard_modifier(&self) -> bool {
        self.inheritance.iter().any(|i| {
            let l = i.to_lowercase();
            l.contains("reentrancyguard") || l.contains("reentrancyguardupgradeable")
        }) || self.modifiers_defs.iter().any(|m| {
            let l = m.name.to_lowercase();
            l.contains("nonreentrant")
        })
    }

    /// Check if this contract defines ERC-165 `supportsInterface`.
    pub fn has_supports_interface(&self) -> bool {
        self.functions.iter().any(|f| f.name == "supportsInterface")
    }

    /// Check if this contract likely implements ERC-721 (has balanceOf, ownerOf).
    pub fn is_erc721(&self) -> bool {
        self.functions.iter().any(|f| f.name == "ownerOf")
            && self.functions.iter().any(|f| f.name == "balanceOf")
    }

    /// Check if this contract likely implements ERC-1155 (has balanceOfBatch).
    pub fn is_erc1155(&self) -> bool {
        self.functions.iter().any(|f| f.name == "balanceOfBatch")
    }

    /// Find a function by name.
    pub fn get_function(&self, name: &str) -> Option<&FunctionDef> {
        self.functions.iter().find(|f| f.name == name)
    }

    /// Find all state variables initialized to a specific value pattern.
    pub fn state_variable_by_name(&self, name: &str) -> Option<&StateVariable> {
        self.state_variables.iter().find(|sv| sv.name == name)
    }
}

/// A struct definition inside a contract.
#[derive(Debug, Clone)]
pub struct StructDef {
    pub name: String,
    pub fields: Vec<StateVariable>,
    pub line: usize,
}

/// A custom error definition.
#[derive(Debug, Clone)]
pub struct ErrorDef {
    pub name: String,
    pub params: Vec<Param>,
    pub line: usize,
}

/// A using-for directive.
#[derive(Debug, Clone)]
pub struct UsingForDef {
    pub type_name: String,
    pub library_name: String,
    pub line: usize,
}

/// The kind of top-level type.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ContractKind {
    Contract,
    Interface,
    Library,
    Abstract,
}

/// Metadata about a function call target (which contract/interface is being called).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CallTarget {
    /// Direct call to an address: address(...).call{...}
    Address,
    /// Call via an interface type: IERC20(token).transfer(...)
    Interface(String),
    /// Call via a known contract name
    Named(String),
    /// Call on msg.sender, tx.origin, etc.
    BuiltIn,
    /// Call on a local variable (could be anything)
    LocalVariable(String),
    /// Self-call: this.function()
    SelfCall,
    /// Unresolved target
    Unknown,
}

/// A parsed Solidity source file.
#[derive(Debug, Clone)]
pub struct SourceFile {
    pub contracts: Vec<Contract>,
    pub pragma: Option<String>,
    pub imports: Vec<String>,
}

impl SourceFile {
    /// Find a contract by name (if defined in this file).
    pub fn get_contract(&self, name: &str) -> Option<&Contract> {
        self.contracts.iter().find(|c| c.name == name)
    }

    /// Resolve whether an inherited contract is known to have access control.
    pub fn inheritance_includes_access_control(&self, contract: &Contract) -> bool {
        for parent_name in &contract.inheritance {
            let parent_name = parent_name.trim();
            // Direct name check
            let parent = self.get_contract(parent_name);
            if parent.map(|p| p.inherits_access_control()).unwrap_or(false) {
                return true;
            }
            // Check if the name itself is a known access-control base
            let l = parent_name.to_lowercase();
            if l.contains("ownable") || l.contains("accesscontrol") || l.contains("auth") {
                return true;
            }
        }
        false
    }

    /// Check if an inherited contract provides reentrancy guard.
    pub fn inheritance_includes_reentrancy_guard(&self, contract: &Contract) -> bool {
        for parent_name in &contract.inheritance {
            let parent_name = parent_name.trim();
            let parent = self.get_contract(parent_name);
            if parent
                .map(|p| p.has_reentrancy_guard_modifier())
                .unwrap_or(false)
            {
                return true;
            }
            let l = parent_name.to_lowercase();
            if l.contains("reentrancyguard") || l.contains("reentrancyguardupgradeable") {
                return true;
            }
        }
        false
    }

    /// Find all state variables across all contracts (useful for cross-contract analysis).
    pub fn all_state_variables(&self) -> Vec<&StateVariable> {
        self.contracts
            .iter()
            .flat_map(|c| c.state_variables.iter())
            .collect()
    }
}

/// Parse a full Solidity source string into a structured representation.
pub fn parse_source(content: &str) -> SourceFile {
    let lines: Vec<&str> = content.lines().collect();
    let cleaned = strip_comments(content);
    let clean_lines: Vec<&str> = cleaned.lines().collect();

    let mut contracts = Vec::new();
    let mut pragma = None;
    let mut imports = Vec::new();

    // First pass: collect pragma and imports
    for line in &clean_lines {
        let t = line.trim();
        if t.starts_with("pragma ") {
            pragma = Some(t.to_string());
        } else if t.starts_with("import ") {
            imports.push(t.to_string());
        }
    }

    // Second pass: find top-level contract/interface/library definitions
    let mut i = 0;
    while i < clean_lines.len() {
        let trimmed = clean_lines[i].trim();

        if let Some(kind) = detect_top_level_kind(trimmed) {
            let name = extract_name(trimmed).unwrap_or("Unknown");
            let inheritance = extract_inheritance(trimmed);

            // Find the opening brace of the contract body
            let mut brace_start = i;
            if trimmed.contains('{') {
                // brace is on the same line
            } else {
                // find next non-commented line with {
                let mut j = i + 1;
                while j < clean_lines.len() && !clean_lines[j].trim().contains('{') {
                    j += 1;
                }
                if j < clean_lines.len() {
                    brace_start = j;
                }
            }

            // Find the closing brace with depth counting
            let (body_start, body_end) = find_body_range(&clean_lines, brace_start);
            if body_end <= body_start {
                i += 1;
                continue;
            }

            let body_lines = &clean_lines[body_start..body_end];

            // Parse contract body
            let (state_vars, funcs, mods, evts, structs, errors, using_for) =
                parse_contract_body(body_lines, &original_lines(&lines, body_start, body_end));

            contracts.push(Contract {
                name: name.to_string(),
                inheritance,
                state_variables: state_vars,
                functions: funcs,
                modifiers_defs: mods,
                events: evts,
                structs,
                errors,
                using_for,
                kind,
                line: i + 1,
            });

            i = body_end; // skip past the contract body
        }
        i += 1;
    }

    SourceFile {
        contracts,
        pragma,
        imports,
    }
}

/// Extract the call target from a statement text.
pub fn extract_call_target(text: &str) -> CallTarget {
    let t = text.trim();

    // Check for self-call: this.function(...)
    if t.starts_with("this.") {
        return CallTarget::SelfCall;
    }

    // Check for address(...).call{...}
    if t.starts_with("address(") {
        return CallTarget::Address;
    }

    // Check for interface pattern: SomeInterface(address).func(...)
    if let Some(dot_pos) = t.find(".") {
        let before_dot = t[..dot_pos].trim();
        // If it contains a parenthesis, it's likely Interface(address)...
        if before_dot.ends_with(')') {
            // Extract the interface/contract name before '(' if present
            if let Some(paren_pos) = before_dot.find('(') {
                let name = before_dot[..paren_pos].trim();
                // Check if it looks like an interface (I prefix or known patterns)
                if name.starts_with('I') && name[1..].starts_with(|c: char| c.is_uppercase()) {
                    return CallTarget::Interface(name.to_string());
                }
                if name == "address" {
                    return CallTarget::Address;
                }
                return CallTarget::Named(name.to_string());
            }
            return CallTarget::Address;
        } else {
            // Direct property access: msg.sender, tx.origin, etc.
            let lower = before_dot.to_lowercase();
            if lower == "msg"
                || lower == "tx"
                || lower == "block"
                || lower == "abi"
                || lower == "gasleft"
            {
                return CallTarget::BuiltIn;
            }
            // If the part before dot is uppercase, it could be a contract name
            if before_dot.starts_with(|c: char| c.is_uppercase()) {
                return CallTarget::Named(before_dot.to_string());
            }
            return CallTarget::LocalVariable(before_dot.to_string());
        }
    }

    // Direct call: IInterface(address).function(...)
    if t.contains("(") && !t.starts_with("if") && !t.starts_with("for") {
        // Could be a bare function call
        let name = t.split('(').next().unwrap_or("").trim();
        if name.starts_with(|c: char| c.is_lowercase()) {
            return CallTarget::LocalVariable(name.to_string());
        }
    }

    CallTarget::Unknown
}

// ─────────────────────────────────────────────────────────────────
// Internal helpers
// ─────────────────────────────────────────────────────────────────

/// Detect contract/interface/library/abstract
fn detect_top_level_kind(trimmed: &str) -> Option<ContractKind> {
    // Skip lines that are not top-level definitions
    if trimmed.starts_with("contract ") && !trimmed.ends_with(';') {
        Some(ContractKind::Contract)
    } else if trimmed.starts_with("interface ") && !trimmed.ends_with(';') {
        Some(ContractKind::Interface)
    } else if trimmed.starts_with("library ") && !trimmed.ends_with(';') {
        Some(ContractKind::Library)
    } else if trimmed.starts_with("abstract contract ") {
        Some(ContractKind::Abstract)
    } else {
        None
    }
}

/// Extract the contract/interface/library name from a definition line.
fn extract_name(trimmed: &str) -> Option<&str> {
    let stripped = trimmed
        .strip_prefix("abstract contract ")
        .or_else(|| trimmed.strip_prefix("contract "))
        .or_else(|| trimmed.strip_prefix("interface "))
        .or_else(|| trimmed.strip_prefix("library "))?;

    // Take the first word (before whitespace, '(', or '{')
    let name = stripped.trim().split([' ', '(', '{']).next().unwrap_or("");
    if name.is_empty() {
        None
    } else {
        Some(name)
    }
}

/// Extract inherited contract names.
fn extract_inheritance(trimmed: &str) -> Vec<String> {
    let stripped = trimmed
        .strip_prefix("abstract contract ")
        .or_else(|| trimmed.strip_prefix("contract "))
        .or_else(|| trimmed.strip_prefix("interface "))
        .or_else(|| trimmed.strip_prefix("library "));

    let Some(after_kw) = stripped else {
        return vec![];
    };

    // Find 'is' keyword after the name
    let after_is = match after_kw.find(" is ") {
        Some(pos) => &after_kw[pos + 4..],
        None => return vec![],
    };

    // Take everything before '{' or end of line
    let inherits_str = if let Some(brace_pos) = after_is.find('{') {
        &after_is[..brace_pos]
    } else {
        after_is
    };

    inherits_str
        .split(',')
        .map(|s| s.trim().trim_end_matches([' ', ')']).trim().to_string())
        .filter(|s| !s.is_empty())
        .collect()
}

/// Find the range of lines that form the body of a contract (or function) given the opening brace line.
/// Returns (body_start, body_end) where body_end is exclusive.
fn find_body_range(lines: &[&str], open_brace_line: usize) -> (usize, usize) {
    let mut depth = 0u32;
    let body_start = open_brace_line + 1;
    let mut started = false;

    for (i, line) in lines.iter().enumerate() {
        if i < open_brace_line {
            continue;
        }
        for ch in line.chars() {
            match ch {
                '{' => {
                    depth += 1;
                    started = true;
                }
                '}' => {
                    depth = depth.saturating_sub(1);
                    if started && depth == 0 {
                        // Found closing brace — body is from open_brace+1 to here
                        return (body_start, i);
                    }
                }
                _ => {}
            }
        }
    }

    // Fallback: unbalanced braces, return empty
    (open_brace_line + 1, open_brace_line + 1)
}

/// Parse the body lines of a contract and extract state variables, functions, modifiers, events,
/// structs, errors, and using-for directives.
#[allow(clippy::type_complexity)]
fn parse_contract_body(
    body_lines: &[&str],
    original_lines: &[&str],
) -> (
    Vec<StateVariable>,
    Vec<FunctionDef>,
    Vec<ModifierDef>,
    Vec<EventDef>,
    Vec<StructDef>,
    Vec<ErrorDef>,
    Vec<UsingForDef>,
) {
    let mut state_vars = Vec::new();
    let mut functions = Vec::new();
    let mut mods = Vec::new();
    let mut events = Vec::new();
    let mut structs = Vec::new();
    let mut errors = Vec::new();
    let mut using_for = Vec::new();

    let mut i = 0;
    while i < body_lines.len() {
        let trimmed = body_lines[i].trim();

        // Skip empty/comment-only lines
        if trimmed.is_empty()
            || trimmed.starts_with("//")
            || trimmed.starts_with("/*")
            || trimmed.starts_with("*")
        {
            i += 1;
            continue;
        }

        // Parse structs
        if trimmed.starts_with("struct ") {
            if let Some(s) = parse_struct_def(body_lines, i) {
                structs.push(s);
            }
            i += 1;
            continue;
        }

        // Parse errors
        if trimmed.starts_with("error ") {
            if let Some(err) = parse_error_def(trimmed, i) {
                errors.push(err);
            }
            i += 1;
            continue;
        }

        // Parse using-for directives
        if trimmed.starts_with("using ") && trimmed.contains(" for ") {
            if let Some(uf) = parse_using_for(trimmed, i) {
                using_for.push(uf);
            }
            i += 1;
            continue;
        }

        // Parse events
        if trimmed.starts_with("event ") {
            if let Some(evt) = parse_event_line(trimmed, i) {
                events.push(evt);
            }
            i += 1;
            continue;
        }

        // Parse modifiers
        if trimmed.starts_with("modifier ") {
            let (mod_def, consumed) = parse_modifier(body_lines, i, original_lines);
            if let Some(m) = mod_def {
                mods.push(m);
            }
            i += consumed;
            continue;
        }

        // Parse functions
        if trimmed.starts_with("function ") || trimmed.contains("function(") {
            let (func, consumed) = parse_function(body_lines, i, original_lines);
            if let Some(f) = func {
                functions.push(f);
            }
            i += consumed;
            continue;
        }

        // Parse constructor
        if trimmed.starts_with("constructor(") {
            let (func, consumed) = parse_constructor(body_lines, i, original_lines);
            if let Some(f) = func {
                functions.push(f);
            }
            i += consumed;
            continue;
        }

        // Parse fallback
        if trimmed.starts_with("fallback(") || trimmed == "fallback()" {
            let (func, consumed) = parse_fallback(body_lines, i, original_lines);
            if let Some(f) = func {
                functions.push(f);
            }
            i += consumed;
            continue;
        }

        // Parse receive
        if trimmed.starts_with("receive()") || trimmed.starts_with("receive (") {
            let (func, consumed) = parse_receive(body_lines, i, original_lines);
            if let Some(f) = func {
                functions.push(f);
            }
            i += consumed;
            continue;
        }

        // Parse state variables (lines ending with ;, not function-like)
        if !trimmed.starts_with("function ")
            && !trimmed.starts_with("modifier ")
            && !trimmed.starts_with("event ")
            && !trimmed.starts_with("error ")
            && !trimmed.starts_with("using ")
            && !trimmed.starts_with("type ")
            && trimmed.ends_with(';')
            && !trimmed.starts_with("//")
        {
            if let Some(sv) = parse_state_variable(trimmed, i) {
                state_vars.push(sv);
            }
        }

        i += 1;
    }

    (
        state_vars, functions, mods, events, structs, errors, using_for,
    )
}

/// Parse a struct definition.
fn parse_struct_def(lines: &[&str], start: usize) -> Option<StructDef> {
    let first_line = lines[start].trim();
    let name = first_line
        .strip_prefix("struct ")?
        .trim()
        .split([' ', '{', '('])
        .next()?
        .to_string();
    if name.is_empty() {
        return None;
    }

    // Find opening brace
    let brace_line = if first_line.contains('{') {
        start
    } else {
        let mut j = start + 1;
        while j < lines.len() && !lines[j].trim().contains('{') {
            j += 1;
        }
        if j >= lines.len() {
            return None;
        }
        j
    };

    let (body_start, body_end) = find_body_range(lines, brace_line);
    if body_end <= body_start {
        return Some(StructDef {
            name,
            fields: Vec::new(),
            line: start + 1,
        });
    }

    let mut fields = Vec::new();
    for line in &lines[body_start..body_end] {
        let trimmed = line.trim();
        if trimmed.is_empty() || trimmed.starts_with("//") {
            continue;
        }
        // Fields end with semicolon
        if trimmed.ends_with(';') {
            if let Some(sv) = parse_state_variable(trimmed, body_start) {
                fields.push(sv);
            }
        }
    }

    Some(StructDef {
        name,
        fields,
        line: start + 1,
    })
}

/// Parse an error definition: error Unauthorized(address caller);
fn parse_error_def(trimmed: &str, line_idx: usize) -> Option<ErrorDef> {
    let content = trimmed.strip_prefix("error ")?;
    let paren_pos = content.find('(')?;
    let name = content[..paren_pos].trim().to_string();
    if name.is_empty() {
        return None;
    }
    let params_str = if let Some(close_paren) = content.rfind(')') {
        let start = paren_pos + 1;
        if start < close_paren {
            &content[start..close_paren]
        } else {
            ""
        }
    } else {
        ""
    };
    let params = parse_params(params_str);
    Some(ErrorDef {
        name,
        params,
        line: line_idx + 1,
    })
}

/// Parse a using-for directive: using Lib for Type;
fn parse_using_for(trimmed: &str, line_idx: usize) -> Option<UsingForDef> {
    let content = trimmed.strip_prefix("using ")?;
    let parts: Vec<&str> = content.splitn(2, " for ").collect();
    if parts.len() != 2 {
        return None;
    }
    let library_name = parts[0].trim().to_string();
    let type_name = parts[1].trim().trim_end_matches(';').to_string();
    if library_name.is_empty() || type_name.is_empty() {
        return None;
    }
    Some(UsingForDef {
        type_name,
        library_name,
        line: line_idx + 1,
    })
}

/// Get the original (uncommented) lines for the given range.
fn original_lines<'a>(lines: &[&'a str], start: usize, end: usize) -> Vec<&'a str> {
    if start < lines.len() && end <= lines.len() && start < end {
        lines[start..end].to_vec()
    } else {
        Vec::new()
    }
}

/// Parse an event definition line like: event Transfer(address indexed from, address indexed to, uint256 value);
fn parse_event_line(trimmed: &str, _line_idx: usize) -> Option<EventDef> {
    let content = trimmed.strip_prefix("event ")?;
    // Extract name before '('
    let paren_pos = content.find('(')?;
    let name = content[..paren_pos].trim().to_string();

    if name.is_empty() {
        return None;
    }

    // Extract params between '(' and ')'
    let params_str = if let Some(close_paren) = content.rfind(')') {
        let start = paren_pos + 1;
        if start < close_paren {
            &content[start..close_paren]
        } else {
            ""
        }
    } else {
        ""
    };

    let params = parse_params(params_str);

    Some(EventDef {
        name,
        params,
        line: 0,
    })
}

/// Parse a state variable declaration.
fn parse_state_variable(trimmed: &str, _line_idx: usize) -> Option<StateVariable> {
    // Strip modifiers like public, internal, private, constant, immutable
    let mut visibility = Visibility::Internal; // default
    let mut cleaned = trimmed.to_string();

    if cleaned.contains(" public ") || cleaned.starts_with("public ") {
        visibility = Visibility::Public;
        cleaned = cleaned.replace(" public ", " ");
    } else if cleaned.contains(" internal ") || cleaned.starts_with("internal ") {
        visibility = Visibility::Internal;
        cleaned = cleaned.replace(" internal ", " ");
    } else if cleaned.contains(" private ") || cleaned.starts_with("private ") {
        visibility = Visibility::Private;
        cleaned = cleaned.replace(" private ", " ");
    } else if cleaned.contains(" external ") || cleaned.starts_with("external ") {
        visibility = Visibility::External;
        cleaned = cleaned.replace(" external ", " ");
    }

    // Remove constant/immutable/override/virtual keywords
    for kw in &[" constant ", " immutable ", " override ", " virtual "] {
        cleaned = cleaned.replace(kw, " ");
    }

    // Remove initializer ( = ... ) — but skip => arrow
    let eq_check = cleaned.replace("=>", "  ");
    if let Some(eq_pos) = eq_check.find('=') {
        cleaned = cleaned[..eq_pos].trim().to_string();
    }

    // Remove trailing semicolon
    cleaned = cleaned.trim_end_matches(';').trim().to_string();

    // The last token before the semicolon is the variable name
    let parts: Vec<&str> = cleaned.split_whitespace().collect();
    if parts.len() < 2 {
        return None;
    }

    // The last part is likely the name; everything before is the type
    let name = parts.last()?.to_string();
    let type_name = parts[..parts.len() - 1].join(" ");

    if name.is_empty() || type_name.is_empty() {
        return None;
    }

    Some(StateVariable {
        name,
        type_name,
        visibility,
        line: 0,
    })
}

/// Parse a function definition, extracting signature and body lines.
fn parse_function(
    lines: &[&str],
    start: usize,
    _original: &[&str],
) -> (Option<FunctionDef>, usize) {
    let first_line = lines[start].trim();
    if !first_line.starts_with("function ") {
        // Might be a function on multiple lines — check next few
        if !first_line.starts_with("function") {
            return (None, 1);
        }
    }

    // Collect all lines until we hit the opening brace
    let mut sig_lines = Vec::new();
    let mut brace_line = start;
    let mut found_brace = false;

    for (offset, line) in lines[start..].iter().enumerate() {
        sig_lines.push(line.to_string());
        if line.contains('{') {
            brace_line = start + offset;
            found_brace = true;
            break;
        }
        // Semicolon-only lines (abstract/interface functions)
        if line.trim().ends_with(';') && !line.trim().contains('{') {
            // This is an abstract function declaration
            let sig = sig_lines
                .iter()
                .map(|l| l.trim())
                .collect::<Vec<_>>()
                .join(" ");
            return match build_function_from_sig(&sig, start + 1, true) {
                Some(func) => (Some(func), offset + 1),
                None => (None, offset + 1),
            };
        }
    }

    if !found_brace {
        // Abstract function — semicolon at end
        let sig = sig_lines
            .iter()
            .map(|l| l.trim())
            .collect::<Vec<_>>()
            .join(" ");
        return match build_function_from_sig(&sig, start + 1, true) {
            Some(func) => (Some(func), sig_lines.len()),
            None => (None, sig_lines.len()),
        };
    }

    // Reconstruct the full signature
    let sig = sig_lines
        .iter()
        .map(|l| l.trim())
        .collect::<Vec<_>>()
        .join(" ");
    let sig = sig.trim_end_matches('{').trim().to_string();

    // Find the body range (after the opening brace)
    let (body_start, body_end) = find_body_range(lines, brace_line);

    // Parse body statements
    let body = if body_end > body_start {
        let body_slice = &lines[body_start..body_end];
        parse_body_statements(body_slice, start + body_start + 1)
    } else {
        Vec::new()
    };

    let mut func = match build_function_from_sig(&sig, start + 1, false) {
        Some(f) => f,
        None => return (None, (body_end + 1).saturating_sub(start).max(1)),
    };
    func.body = body;

    let consumed = (body_end + 1).saturating_sub(start).max(1);

    (Some(func), consumed)
}

/// Parse a constructor.
fn parse_constructor(
    lines: &[&str],
    start: usize,
    _original: &[&str],
) -> (Option<FunctionDef>, usize) {
    let first_line = lines[start].trim();
    if !first_line.starts_with("constructor(") {
        return (None, 1);
    }

    let sig = first_line.trim_end_matches('{').trim().to_string();
    let params_str = extract_params_str(&sig);

    let (body_start, body_end) = find_body_range(lines, start);

    let body = if body_end > body_start {
        parse_body_statements(&lines[body_start..body_end], body_start + 1)
    } else {
        Vec::new()
    };

    let func = FunctionDef {
        name: "constructor".to_string(),
        visibility: Visibility::Internal,
        mutability: Mutability::Nonpayable,
        modifiers: Vec::new(),
        params: parse_params(&params_str),
        return_params: Vec::new(),
        body,
        line: start + 1,
        is_constructor: true,
        is_fallback: false,
        is_receive: false,
    };

    let consumed = (body_end + 1).saturating_sub(start).max(1);
    (Some(func), consumed)
}

/// Parse a fallback function.
fn parse_fallback(
    lines: &[&str],
    start: usize,
    _original: &[&str],
) -> (Option<FunctionDef>, usize) {
    let first_line = lines[start].trim();
    if !first_line.starts_with("fallback(") && first_line != "fallback()" {
        return (None, 1);
    }

    let sig = first_line.trim_end_matches('{').trim().to_string();
    let (body_start, body_end) = find_body_range(lines, start);

    let body = if body_end > body_start {
        parse_body_statements(&lines[body_start..body_end], body_start + 1)
    } else {
        Vec::new()
    };

    let visibility = if sig.contains("external") {
        Visibility::External
    } else {
        Visibility::Public
    };
    let mutability = if sig.contains("payable") {
        Mutability::Payable
    } else if sig.contains("view") {
        Mutability::View
    } else {
        Mutability::Nonpayable
    };
    let modifiers = extract_modifiers(&sig);

    let func = FunctionDef {
        name: "fallback".to_string(),
        visibility,
        mutability,
        modifiers,
        params: Vec::new(),
        return_params: Vec::new(),
        body,
        line: start + 1,
        is_constructor: false,
        is_fallback: true,
        is_receive: false,
    };

    let consumed = (body_end + 1).saturating_sub(start).max(1);
    (Some(func), consumed)
}

/// Parse a receive() function.
fn parse_receive(lines: &[&str], start: usize, _original: &[&str]) -> (Option<FunctionDef>, usize) {
    let first_line = lines[start].trim();
    if !first_line.starts_with("receive(") && !first_line.starts_with("receive (") {
        return (None, 1);
    }

    let (body_start, body_end) = find_body_range(lines, start);

    let body = if body_end > body_start {
        parse_body_statements(&lines[body_start..body_end], body_start + 1)
    } else {
        Vec::new()
    };

    let func = FunctionDef {
        name: "receive".to_string(),
        visibility: Visibility::External,
        mutability: Mutability::Payable,
        modifiers: Vec::new(),
        params: Vec::new(),
        return_params: Vec::new(),
        body,
        line: start + 1,
        is_constructor: false,
        is_fallback: false,
        is_receive: true,
    };

    let consumed = (body_end + 1).saturating_sub(start).max(1);
    (Some(func), consumed)
}

/// Parse a modifier definition.
fn parse_modifier(
    lines: &[&str],
    start: usize,
    _original: &[&str],
) -> (Option<ModifierDef>, usize) {
    let first_line = lines[start].trim();
    if !first_line.starts_with("modifier ") {
        return (None, 1);
    }

    // Collect signature lines until opening brace
    let mut brace_line = start;
    let mut found_brace = false;

    for (offset, line) in lines[start..].iter().enumerate() {
        if line.contains('{') {
            brace_line = start + offset;
            found_brace = true;
            break;
        }
        // Semicolon-only = abstract modifier
        if line.trim().ends_with(';') && offset > 0 {
            let name = first_line
                .strip_prefix("modifier ")
                .and_then(|s| s.trim().split('(').next())
                .unwrap_or("")
                .to_string();
            return (
                Some(ModifierDef {
                    name,
                    params: Vec::new(),
                    body: Vec::new(),
                    line: start + 1,
                }),
                offset + 1,
            );
        }
    }

    if !found_brace {
        return (None, 1);
    }

    let sig = first_line.to_string();
    let name = sig
        .strip_prefix("modifier ")
        .and_then(|s| s.trim().split('(').next())
        .unwrap_or("")
        .to_string();

    let (body_start, body_end) = find_body_range(lines, brace_line);
    let body = if body_end > body_start {
        parse_body_statements(&lines[body_start..body_end], body_start + 1)
    } else {
        Vec::new()
    };

    let consumed = (body_end + 1).saturating_sub(start).max(1);
    (
        Some(ModifierDef {
            name,
            params: Vec::new(),
            body,
            line: start + 1,
        }),
        consumed,
    )
}

/// Build a FunctionDef from a reconstructed signature string.
fn build_function_from_sig(sig: &str, line: usize, _is_abstract: bool) -> Option<FunctionDef> {
    // Extract function name
    let after_fn = sig.strip_prefix("function ")?;
    let name = after_fn.split('(').next().unwrap_or("").trim().to_string();

    if name.is_empty() {
        return None;
    }

    let params_str = extract_params_str(sig);
    let params = parse_params(&params_str);

    // Extract return parameters
    let return_params = if let Some(returns_pos) = sig.find(" returns (") {
        let ret_str = &sig[returns_pos + 9..];
        let ret_str = ret_str.trim_end_matches(';').trim_end_matches(')');
        parse_params(ret_str)
    } else {
        Vec::new()
    };

    // Determine visibility
    let visibility = if sig.contains(" external ")
        || sig.ends_with(" external")
        || sig.starts_with("external ")
    {
        Visibility::External
    } else if sig.contains(" public ") || sig.ends_with(" public") || sig.starts_with("public ") {
        Visibility::Public
    } else if sig.contains(" internal ") || sig.ends_with(" internal") {
        Visibility::Internal
    } else if sig.contains(" private ") || sig.ends_with(" private") {
        Visibility::Private
    } else {
        // Default visibility
        Visibility::Public
    };

    // Determine mutability
    let mutability = if sig.contains(" pure ") {
        Mutability::Pure
    } else if sig.contains(" view ") {
        Mutability::View
    } else if sig.contains(" payable ") {
        Mutability::Payable
    } else {
        Mutability::Nonpayable
    };

    let modifiers = extract_modifiers(sig);

    Some(FunctionDef {
        name,
        visibility,
        mutability,
        modifiers,
        params,
        return_params,
        body: Vec::new(), // filled in by caller
        line,
        is_constructor: false,
        is_fallback: false,
        is_receive: false,
    })
}

/// Extract parameters string from between parentheses (handles nesting).
fn extract_params_str(sig: &str) -> String {
    let open_paren = match sig.find('(') {
        Some(p) => p,
        None => return String::new(),
    };
    let mut depth = 0u32;
    let mut close_paren = sig.len();
    for (i, ch) in sig[open_paren..].char_indices() {
        match ch {
            '(' => depth += 1,
            ')' => {
                depth -= 1;
                if depth == 0 {
                    close_paren = open_paren + i;
                    break;
                }
            }
            _ => {}
        }
    }
    sig[open_paren + 1..close_paren].to_string()
}

/// Parse parameter definitions from a comma-separated string.
fn parse_params(params_str: &str) -> Vec<Param> {
    let mut params = Vec::new();
    let mut depth = 0u32;
    let mut current = String::new();

    for ch in params_str.chars() {
        match ch {
            '(' | '<' => {
                depth += 1;
                current.push(ch);
            }
            ')' | '>' => {
                depth = depth.saturating_sub(1);
                current.push(ch);
            }
            ',' if depth == 0 => {
                if let Some(param) = parse_single_param(current.trim()) {
                    params.push(param);
                }
                current.clear();
            }
            _ => current.push(ch),
        }
    }

    if !current.trim().is_empty() {
        if let Some(param) = parse_single_param(current.trim()) {
            params.push(param);
        }
    }

    params
}

/// Parse a single parameter like "address indexed from" or "uint256 amount".
fn parse_single_param(s: &str) -> Option<Param> {
    let s = s.trim();
    if s.is_empty() {
        return None;
    }

    let indexed = s.contains(" indexed ");
    let cleaned = s.replace(" indexed ", " ");

    let parts: Vec<&str> = cleaned.split_whitespace().collect();
    if parts.is_empty() {
        return None;
    }

    if parts.len() == 1 {
        // Only type, no name (common in events)
        return Some(Param {
            name: String::new(),
            type_name: parts[0].to_string(),
            indexed,
        });
    }

    let name = parts.last()?.to_string();
    let type_name = parts[..parts.len() - 1].join(" ");

    Some(Param {
        name,
        type_name,
        indexed,
    })
}

/// Extract modifier names from a function signature.
fn extract_modifiers(sig: &str) -> Vec<String> {
    // Remove visibility and mutability keywords
    let cleaned = sig
        .replace(" public ", " ")
        .replace(" external ", " ")
        .replace(" internal ", " ")
        .replace(" private ", " ")
        .replace(" pure ", " ")
        .replace(" view ", " ")
        .replace(" payable ", " ")
        .replace(" virtual ", " ")
        .replace(" override ", " ");

    // Modifiers appear after the first ')', before 'returns'
    let after_paren = match cleaned.find(')') {
        Some(p) => &cleaned[p + 1..],
        None => "",
    };

    let before_returns = match after_paren.find(" returns ") {
        Some(p) => &after_paren[..p],
        None => after_paren,
    };

    before_returns
        .split_whitespace()
        .map(|s| {
            s.trim()
                .trim_start_matches(['(', ')', ',', ' ', '\t'])
                .trim_end_matches(['(', ')', ',', ' '])
                .to_string()
        })
        .filter(|s| {
            !s.is_empty()
                && s != "public"
                && s != "external"
                && s != "internal"
                && s != "private"
                && s != "pure"
                && s != "view"
                && s != "payable"
                && s != "virtual"
                && s != "override"
                && s != "returns"
        })
        .collect()
}

/// Parse body lines into classified statements.
fn parse_body_statements(lines: &[&str], base_line: usize) -> Vec<Statement> {
    let mut statements = Vec::new();
    let mut stmt_buffer = String::new();
    let mut stmt_start = base_line;

    for (offset, line) in lines.iter().enumerate() {
        let trimmed = line.trim();
        let abs_line = base_line + offset;

        // Skip comment-only lines
        if trimmed.starts_with("//") {
            continue;
        }

        // Handle multi-line comments
        if trimmed.starts_with("/*") || trimmed.starts_with("*") || trimmed.ends_with("*/") {
            continue;
        }

        if trimmed.is_empty() {
            continue;
        }

        // Check for inline assembly
        if trimmed.starts_with("assembly ") || trimmed.starts_with("assembly{") {
            statements.push(Statement {
                kind: StatementKind::InlineAssembly,
                line: abs_line,
                text: trimmed.to_string(),
            });
            continue;
        }

        // Accumulate lines until we hit a semicolon or brace block
        stmt_buffer.push_str(trimmed);
        stmt_buffer.push(' ');

        // Is this a complete statement?
        if trimmed.ends_with(';') || trimmed.ends_with("{\"") || trimmed.ends_with('}') {
            let text = stmt_buffer.trim().to_string();
            let kind = classify_statement(&text);
            statements.push(Statement {
                kind,
                line: stmt_start,
                text,
            });
            stmt_buffer.clear();
            stmt_start = abs_line + 1;
        } else if trimmed.contains('{') || trimmed.ends_with(')') || trimmed.ends_with("){") {
            // Opening brace could be start of a block
            if trimmed.contains('{') && !trimmed.contains('}') {
                // Statement with opening brace - might be if/for/while
                let text = stmt_buffer.trim().to_string();
                let kind = classify_statement(&text);
                statements.push(Statement {
                    kind,
                    line: stmt_start,
                    text,
                });
                stmt_buffer.clear();
                stmt_start = abs_line + 1;
            }
        }
    }

    // Flush any remaining buffer
    if !stmt_buffer.is_empty() {
        let text = stmt_buffer.trim().to_string();
        let kind = classify_statement(&text);
        statements.push(Statement {
            kind,
            line: stmt_start,
            text,
        });
    }

    statements
}

/// Classify a single statement line into a StatementKind.
fn classify_statement(text: &str) -> StatementKind {
    let t = text.trim();

    // Inline assembly
    if t.starts_with("assembly ") || t.starts_with("assembly{") {
        return StatementKind::InlineAssembly;
    }

    // Guard statements
    if t.starts_with("require(")
        || t.starts_with("require (")
        || t.starts_with("revert ")
        || t.starts_with("assert(")
    {
        return StatementKind::Guard;
    }

    // Flow control
    if t.starts_with("if ")
        || t.starts_with("if(")
        || t.starts_with("while ")
        || t.starts_with("while(")
        || t.starts_with("for (")
        || t.starts_with("for(")
    {
        return StatementKind::FlowControl;
    }

    // Emit
    if t.starts_with("emit ") {
        return StatementKind::Emit;
    }

    // External calls: .call{, .delegatecall{, .staticcall{, .transfer(, .send(
    // Also detect interface calls: IERC20(token).transfer(...), etc.
    if t.contains(".call{") || t.contains(".call(") || t.contains(".delegatecall{")
        || t.contains(".delegatecall(") || t.contains(".staticcall{")
        || t.contains(".staticcall(") || t.contains(".transfer(")
        || t.contains(".send(")
        // Pattern: ContractName(...).functionName( or address(...).call(
        || (t.contains(").") && (t.contains(".call") || t.contains(".transfer") || t.contains(".send")))
    {
        return StatementKind::ExternalCall;
    }

    // State writes
    if t.contains('=')
        || t.starts_with("delete ")
        || t.contains("++")
        || t.contains("--")
        || t.starts_with("mapping")
        || t.contains(" += ")
        || t.contains(" -= ")
        || t.contains(" *= ")
        || t.contains(" /= ")
    {
        return StatementKind::StateWrite;
    }

    // Emit via event call (older Solidity style without emit keyword)
    if t.contains('(') && !t.starts_with("//") && !t.contains("function") && !t.contains("modifier")
    {
        // Could be an event emit or internal call
        return StatementKind::InternalCall;
    }

    StatementKind::Other
}

/// Strip comments from Solidity source code, preserving line count.
fn strip_comments(content: &str) -> String {
    let mut result = String::with_capacity(content.len());
    let chars: Vec<char> = content.chars().collect();
    let len = chars.len();
    let mut i = 0;
    let mut in_multiline = false;

    while i < len {
        if !in_multiline && i + 1 < len && chars[i] == '/' && chars[i + 1] == '/' {
            // Single-line comment: skip to end of line
            while i < len && chars[i] != '\n' {
                result.push(chars[i]);
                i += 1;
            }
            // keep the newline
            if i < len && chars[i] == '\n' {
                result.push('\n');
                i += 1;
            }
        } else if !in_multiline && i + 1 < len && chars[i] == '/' && chars[i + 1] == '*' {
            // Multi-line comment start
            in_multiline = true;
            // Preserve the /* to maintain line count
            result.push_str("/*");
            i += 2;
        } else if in_multiline && i + 1 < len && chars[i] == '*' && chars[i + 1] == '/' {
            in_multiline = false;
            result.push_str("*/");
            i += 2;
        } else if in_multiline {
            // Preserve whitespace/newlines inside multi-line comments to keep line count
            if chars[i] == '\n' {
                result.push('\n');
            } else if chars[i] == '\r' {
                // skip
            }
            i += 1;
        } else {
            result.push(chars[i]);
            i += 1;
        }
    }

    result
}

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

    #[test]
    fn test_simple_contract() {
        let source = r#"
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

contract MyToken is Ownable {
    uint256 public totalSupply;
    mapping(address => uint256) public balances;

    event Transfer(address indexed from, address indexed to, uint256 value);

    function mint(address to, uint256 amount) external onlyOwner {
        balances[to] += amount;
        totalSupply += amount;
        emit Transfer(address(0), to, amount);
    }

    function withdraw(uint256 amount) public {
        require(amount > 0, "Amount too low");
        balances[msg.sender] -= amount;
        totalSupply -= amount;
        payable(msg.sender).transfer(amount);
    }
}
"#;

        let parsed = parse_source(source);
        assert_eq!(parsed.contracts.len(), 1);
        let c = &parsed.contracts[0];
        assert_eq!(c.name, "MyToken");
        assert_eq!(c.inheritance, vec!["Ownable"]);
        assert!(c.inherits_access_control());
        assert_eq!(c.state_variables.len(), 2);
        assert_eq!(c.functions.len(), 2);
    }

    #[test]
    fn test_function_parsing() {
        let source = r#"
contract Test {
    function safeWithdraw(uint256 amount) external nonReentrant {
        uint256 balance = balances[msg.sender];
        require(balance >= amount, "Insufficient balance");
        balances[msg.sender] = balance - amount;
        payable(msg.sender).transfer(amount);
    }
}
"#;
        let parsed = parse_source(source);
        let func = &parsed.contracts[0].functions[0];
        assert_eq!(func.name, "safeWithdraw");
        assert_eq!(func.visibility, Visibility::External);
        assert_eq!(func.modifiers, vec!["nonReentrant"]);
        assert!(func.has_reentrancy_guard());
        assert_eq!(func.params.len(), 1);
        assert_eq!(func.params[0].name, "amount");
    }

    #[test]
    fn test_cei_violation_detection() {
        let source = r#"
contract Vulnerable {
    mapping(address => uint256) public balances;

    function withdraw(uint256 amount) public {
        require(amount > 0, "Invalid");
        payable(msg.sender).transfer(amount);
        balances[msg.sender] -= amount;
    }
}
"#;
        let parsed = parse_source(source);
        let func = &parsed.contracts[0].functions[0];
        assert_eq!(func.name, "withdraw");

        // Find statement kinds
        let kinds: Vec<&StatementKind> = func.body.iter().map(|s| &s.kind).collect();
        // require -> guard
        // .transfer -> external call
        // -= -> state write (after external call = CEI violation)
        assert!(kinds.contains(&&StatementKind::ExternalCall));
        assert!(kinds.contains(&&StatementKind::Guard));
        assert!(kinds.contains(&&StatementKind::StateWrite));
    }

    #[test]
    fn test_access_control_detection() {
        let source = r#"
contract NoAccess {
    function withdrawAll() public {
        payable(msg.sender).transfer(address(this).balance);
    }

    function setAdmin(address newAdmin) external {
        admin = newAdmin;
    }
}
"#;
        let parsed = parse_source(source);
        let contract = &parsed.contracts[0];
        assert!(!contract.inherits_access_control());

        let withdraw_func = &contract.functions[0];
        assert!(!withdraw_func.has_access_control(contract));

        let admin_func = &contract.functions[1];
        assert!(!admin_func.has_access_control(contract));
    }

    #[test]
    fn test_proper_access_control() {
        let source = r#"
contract Secured is Ownable {
    function withdrawAll() public onlyOwner {
        payable(msg.sender).transfer(address(this).balance);
    }
}
"#;
        let parsed = parse_source(source);
        let contract = &parsed.contracts[0];
        assert!(contract.inherits_access_control());

        let func = &contract.functions[0];
        assert!(func.has_access_control(contract));
        assert_eq!(func.modifiers[0], "onlyOwner");
    }

    #[test]
    fn test_state_variable_parsing() {
        let source = r#"
contract Storage {
    uint256 public count;
    address private owner;
    mapping(address => uint256) internal balances;
    bool public initialized;
}
"#;
        let parsed = parse_source(source);
        let state_vars = &parsed.contracts[0].state_variables;
        assert_eq!(state_vars.len(), 4);

        assert_eq!(state_vars[0].name, "count");
        assert_eq!(state_vars[0].type_name, "uint256");
        assert_eq!(state_vars[0].visibility, Visibility::Public);

        assert_eq!(state_vars[1].name, "owner");
        assert_eq!(state_vars[1].visibility, Visibility::Private);

        assert_eq!(state_vars[2].name, "balances");
        assert_eq!(state_vars[2].visibility, Visibility::Internal);

        assert_eq!(state_vars[3].name, "initialized");
        assert_eq!(state_vars[3].visibility, Visibility::Public);
    }

    #[test]
    fn test_event_parsing() {
        let source = r#"
contract Events {
    event Transfer(address indexed from, address indexed to, uint256 value);
    event Approval(address indexed owner, address indexed spender, uint256 value);
}
"#;
        let parsed = parse_source(source);
        let events = &parsed.contracts[0].events;
        assert_eq!(events.len(), 2);
        assert_eq!(events[0].name, "Transfer");
        assert_eq!(events[0].params[0].name, "from");
        assert!(events[0].params[0].indexed);
    }

    #[test]
    fn test_comment_stripping() {
        let source = "\
// This is a comment
contract Test {
    // another comment
    uint256 x;
    /* inline */
    uint256 y;
}
";
        let cleaned = strip_comments(source);
        assert!(cleaned.contains("// This is a comment"));
        assert!(cleaned.contains("// another comment"));
        assert!(cleaned.contains("/*"));
        assert!(cleaned.contains("*/"));

        let parsed = parse_source(&cleaned);
        assert_eq!(parsed.contracts.len(), 1);
        assert_eq!(parsed.contracts[0].state_variables.len(), 2);
    }

    #[test]
    fn test_multiline_function_sig() {
        let source = r#"
contract Multi {
    function complex(
        address param1,
        uint256 param2,
        bytes calldata data
    ) external payable onlyOwner returns (bool success) {
        return true;
    }
}
"#;
        let parsed = parse_source(source);
        let func = &parsed.contracts[0].functions[0];
        assert_eq!(func.name, "complex");
        assert_eq!(func.params.len(), 3);
        assert!(func.modifiers.contains(&"onlyOwner".to_string()));
        assert_eq!(func.return_params.len(), 1);
        assert_eq!(func.return_params[0].name, "success");
    }

    #[test]
    fn test_contract_with_no_functions() {
        let source = r#"
contract Empty {
    uint256 public constant VERSION = 1;
}
"#;
        let parsed = parse_source(source);
        assert_eq!(parsed.contracts[0].functions.len(), 0);
        assert_eq!(parsed.contracts[0].state_variables.len(), 1);
    }

    #[test]
    fn test_interface_parsing() {
        let source = r#"
interface IERC20 {
    function transfer(address to, uint256 amount) external returns (bool);
    function balanceOf(address account) external view returns (uint256);
}
"#;
        let parsed = parse_source(source);
        eprintln!(
            "IFACE test: contracts={}, pragma={:?}, imports={:?}",
            parsed.contracts.len(),
            parsed.pragma,
            parsed.imports
        );
        for (ci, c) in parsed.contracts.iter().enumerate() {
            eprintln!(
                "  contract[{}]: name={}, kind={:?}, functions={}",
                ci,
                c.name,
                c.kind,
                c.functions.len()
            );
            for fi in 0..c.functions.len() {
                eprintln!(
                    "    func[{}]: name={}, body_stmts={}",
                    fi,
                    c.functions[fi].name,
                    c.functions[fi].body.len()
                );
            }
        }
        assert_eq!(parsed.contracts.len(), 1, "Should have 1 contract");
        assert_eq!(parsed.contracts[0].kind, ContractKind::Interface);
        assert_eq!(
            parsed.contracts[0].functions.len(),
            2,
            "Should have 2 functions, got {}",
            parsed.contracts[0].functions.len()
        );
        // Interface functions should have empty bodies
        for func in &parsed.contracts[0].functions {
            assert!(func.body.is_empty());
        }
    }

    #[test]
    fn test_inheritance_multiple() {
        let source = r#"
contract MyContract is Ownable, ReentrancyGuard, AccessControl {
    uint256 dummy;
}
"#;
        let parsed = parse_source(source);
        let c = &parsed.contracts[0];
        assert!(c.inheritance.contains(&"Ownable".to_string()));
        assert!(c.inheritance.contains(&"ReentrancyGuard".to_string()));
        assert!(c.inheritance.contains(&"AccessControl".to_string()));
        assert!(c.inherits_access_control());
        assert!(c.has_reentrancy_guard_modifier());
    }
}