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),

    }

}