use crate::ast::*;
use std::collections::HashSet;

/// Optimizes the AST by removing unused states (properties) and methods.
/// 
/// # Arguments
/// * `ast` - The LiquidAST to be optimized.
pub fn optimize_ast(mut ast: LiquidAST) -> LiquidAST {
    for class in &mut ast.classes {
        if class.extends_class != "Store" {
            optimize_class(class);
        }
    }
    ast
}

/// Optimizes a single component class by tracking state and method dependencies.
fn optimize_class(class: &mut ClassDecl) {
    let mut used_identifiers = HashSet::new();
    let mut active_methods = HashSet::new();
    
    // 1. Collect active methods from event handlers (like onPressed) in the UI render tree
    collect_active_methods_from_element(&class.render.element, &mut active_methods);
    
    // 2. Collect states used directly inside the UI render tree
    collect_identifiers_from_element(&class.render.element, &mut used_identifiers);
    
    // 3. Collect states used/mutated inside active methods
    for method in &class.methods {
        if active_methods.contains(&method.name) {
            for stmt in &method.statements {
                match stmt {
                    Statement::Expr(expr) => {
                        used_identifiers.insert(expr.name.clone());
                    }
                    Statement::Return(val) => {
                        collect_identifiers_from_value(val, &mut used_identifiers);
                    }
                }
            }
        }
    }
    
    // 4. Eliminate unused states (properties)
    class.properties.retain(|prop| used_identifiers.contains(&prop.name));
    
    // 5. Eliminate unused logic methods
    class.methods.retain(|method| active_methods.contains(&method.name));
}

/// Tracks active methods based on event handler calls in UI elements.
fn collect_active_methods_from_element(el: &Element, active: &mut HashSet<String>) {
    for prop in &el.props {
        if prop.name == "onPressed" {
            if let Value::Identifier(id) = &prop.value {
                active.insert(id.clone());
            }
        } else if prop.name == "children" {
            if let Value::Array(arr) = &prop.value {
                for val in arr {
                    if let Value::Element(child_el) = val {
                        collect_active_methods_from_element(child_el, active);
                    }
                }
            }
        }
    }
}

/// Tracks variables used inside a single UI element.
fn collect_identifiers_from_element(el: &Element, used: &mut HashSet<String>) {
    for prop in &el.props {
        if prop.name != "onPressed" {
            collect_identifiers_from_value(&prop.value, used);
        }
    }
}

/// Tracks variables used inside a value.
fn collect_identifiers_from_value(value: &Value, used: &mut HashSet<String>) {
    match value {
        Value::Identifier(id) => {
            used.insert(id.clone());
        }
        Value::StringLit(s) => {
            collect_identifiers_from_string(s, used);
        }
        Value::Array(arr) => {
            for v in arr {
                collect_identifiers_from_value(v, used);
            }
        }
        Value::Element(el) => {
            collect_identifiers_from_element(el, used);
        }
        _ => {}
    }
}

/// Extracts variable interpolations (prefixed with $) from string literals.
fn collect_identifiers_from_string(s: &str, used: &mut HashSet<String>) {
    let chars: Vec<char> = s.chars().collect();
    let mut i = 0;
    while i < chars.len() {
        if chars[i] == '$' && i + 1 < chars.len() && chars[i+1].is_alphabetic() {
            let mut var_name = String::new();
            i += 1;
            while i < chars.len() && (chars[i].is_alphanumeric() || chars[i] == '_') {
                var_name.push(chars[i]);
                i += 1;
            }
            used.insert(var_name);
        } else {
            i += 1;
        }
    }
}
