plum

use std::collections::HashSet;

use crate::ast::{
    Arg, AssignTarget, AttrKind, CasePattern, Expr, FnBody, Item, ParamType, Source, Stmt,
    StringPart, Type,
};

/// The builtin type names this walk tracks: primitive value types with no
/// enum-variant/name-resolution step that naturally fails when unimported
/// (unlike `Bool`, whose `True`/`False` variants simply don't resolve without
/// `import std/Bool`) — so nothing else catches a missing `import std/Int` /
/// `import std/Float` / `import std/Str`.
const TRACKED_TYPES: [&str; 3] = ["Int", "Float", "Str"];

fn isTracked(name: &str) -> bool {
    TRACKED_TYPES.contains(&name)
}

/// Every builtin type name (`Int`, `Float`, `Str` — literal or type-annotation)
/// mentioned anywhere in `source` — used by `loader::loadAndMerge` to enforce
/// that a file using one of them explicitly `import`s it, exactly like every
/// other stdlib type (`Bool`, ...) already must be.
pub fn usedBuiltinTypeNames(source: &Source) -> HashSet<String> {
    let mut names = HashSet::new();
    for item in &source.items {
        match item {
            Item::Trait(t) => {
                for m in &t.methods {
                    for p in &m.params {
                        walkParamType(&p.ty, &mut names);
                    }
                    if let Some(r) = &m.returns {
                        walkType(r, &mut names);
                    }
                }
            }
            Item::Enum(e) => {
                for p in &e.params {
                    walkType(&p.ty, &mut names);
                }
                for v in &e.variants {
                    for f in &v.fields {
                        walkType(f, &mut names);
                    }
                    for val in &v.values {
                        walkExpr(val, &mut names);
                    }
                }
            }
            Item::Fn(f) => {
                for p in &f.params {
                    walkParamType(&p.ty, &mut names);
                    if let Some(d) = &p.default {
                        walkExpr(d, &mut names);
                    }
                }
                if let Some(r) = &f.returns {
                    walkType(r, &mut names);
                }
                match &f.body {
                    FnBody::Expr(e) => walkExpr(e, &mut names),
                    FnBody::Block(b) => walkStmts(&b.stmts, &mut names),
                    FnBody::Extern => {}
                }
            }
            Item::Const(c) => walkExpr(&c.value, &mut names),
            Item::Test(t) => walkStmts(&t.body.stmts, &mut names),
        }
    }
    names
}

fn walkType(ty: &Type, names: &mut HashSet<String>) {
    if isTracked(&ty.name) {
        names.insert(ty.name.clone());
    }
    for g in &ty.generics {
        walkType(g, names);
    }
}

fn walkParamType(pt: &ParamType, names: &mut HashSet<String>) {
    match pt {
        ParamType::Type(t) => walkType(t, names),
        ParamType::Variadic(t) => walkType(t, names),
        ParamType::Fn(params, ret) => {
            for p in params {
                walkType(p, names);
            }
            if let Some(r) = ret {
                walkType(r, names);
            }
        }
    }
}

fn walkStmts(stmts: &[Stmt], names: &mut HashSet<String>) {
    for s in stmts {
        walkStmt(s, names);
    }
}

fn walkStmt(stmt: &Stmt, names: &mut HashSet<String>) {
    match stmt {
        Stmt::Assign(a) => {
            for t in &a.targets {
                if let AssignTarget::Field(obj, _) = t {
                    walkExpr(obj, names);
                }
            }
            for v in &a.values {
                walkExpr(v, names);
            }
        }
        Stmt::Break | Stmt::Continue | Stmt::Todo => {}
        Stmt::Assert(check) => walkExpr(&check.cond, names),
        Stmt::For(f) => {
            walkExpr(&f.iter, names);
            walkStmts(&f.body.stmts, names);
        }
        Stmt::While(w) => {
            walkExpr(&w.condition, names);
            walkStmts(&w.body.stmts, names);
        }
        Stmt::If(i) => {
            walkExpr(&i.condition, names);
            walkStmts(&i.body.stmts, names);
            for ei in &i.else_ifs {
                walkExpr(&ei.condition, names);
                walkStmts(&ei.body.stmts, names);
            }
            if let Some(e) = &i.else_ {
                walkStmts(&e.stmts, names);
            }
        }
        Stmt::Match(m) => {
            for s in &m.subjects {
                walkExpr(s, names);
            }
            for c in &m.cases {
                for p in &c.patterns {
                    walkCasePattern(p, names);
                }
                if let Some(g) = &c.guard {
                    walkExpr(g, names);
                }
                walkStmts(&c.body.stmts, names);
            }
        }
        Stmt::Return(e) => {
            if let Some(e) = e {
                walkExpr(e, names);
            }
        }
        Stmt::Expr(e) => walkExpr(e, names),
    }
}

fn walkCasePattern(pat: &CasePattern, names: &mut HashSet<String>) {
    match pat {
        CasePattern::Class { fields, .. } => {
            for f in fields {
                walkCasePattern(f, names);
            }
        }
        CasePattern::Int(_) => {
            names.insert("Int".to_string());
        }
        CasePattern::Float(_) => {
            names.insert("Float".to_string());
        }
        CasePattern::String(_) => {
            names.insert("Str".to_string());
        }
        CasePattern::Name(_) | CasePattern::Wildcard => {}
    }
}

fn walkExpr(expr: &Expr, names: &mut HashSet<String>) {
    match expr {
        Expr::Binary(b) => {
            walkExpr(&b.left, names);
            walkExpr(&b.right, names);
        }
        Expr::Unary(u) => walkExpr(&u.operand, names),
        Expr::Bool(b) => {
            walkExpr(&b.left, names);
            walkExpr(&b.right, names);
        }
        Expr::Not(e) => walkExpr(e, names),
        Expr::Try(e) => walkExpr(e, names),
        Expr::Compare(c) => {
            walkExpr(&c.left, names);
            walkExpr(&c.right, names);
        }
        Expr::Ternary(t) => {
            walkExpr(&t.condition, names);
            walkExpr(&t.then, names);
            walkExpr(&t.else_, names);
        }
        Expr::Elvis(e) => {
            walkExpr(&e.left, names);
            walkExpr(&e.right, names);
        }
        Expr::FnCall(call) => {
            if isTracked(&call.name) {
                names.insert(call.name.clone());
            }
            for a in &call.args {
                walkArg(a, names);
            }
        }
        Expr::ClassCall(c) => {
            if isTracked(&c.type_name) {
                names.insert(c.type_name.clone());
            }
            for g in &c.generics {
                walkType(g, names);
            }
            for f in &c.fields {
                walkExpr(&f.value, names);
            }
        }
        Expr::Attribute(a) => {
            walkExpr(&a.object, names);
            if let AttrKind::Method(call) = &a.attr {
                for arg in &call.args {
                    walkArg(arg, names);
                }
            }
        }
        Expr::Paren(e) => walkExpr(e, names),
        Expr::String(s) => {
            names.insert("Str".to_string());
            for part in &s.parts {
                if let StringPart::Interp(e) = part {
                    walkExpr(e, names);
                }
            }
        }
        Expr::Int(_) => {
            names.insert("Int".to_string());
        }
        Expr::Float(_) => {
            names.insert("Float".to_string());
        }
        Expr::Self_ | Expr::Var(_) => {}
        Expr::TypeName(name) => {
            if isTracked(name) {
                names.insert(name.clone());
            }
        }
        Expr::Closure(c) => walkStmts(&c.body.stmts, names),
    }
}

fn walkArg(arg: &Arg, names: &mut HashSet<String>) {
    match arg {
        Arg::Positional(e) => walkExpr(e, names),
        Arg::Keyword { value, .. } => walkExpr(value, names),
        Arg::Pair { value, .. } => walkExpr(value, names),
    }
}