// Copyright © 2025 The µcad authors <info@ucad.xyz>
// SPDX-License-Identifier: AGPL-3.0-or-later

//! Collect names from within definitions.

mod name_list;

use crate::syntax::*;
use name_list::*;

pub(super) trait Names {
    fn names(&self) -> NameList;
}

impl Names for SourceFile {
    fn names(&self) -> NameList {
        self.statements.names()
    }
}

impl Names for ModuleDefinition {
    fn names(&self) -> NameList {
        if let Some(body) = &self.body {
            body.names()
        } else {
            Default::default()
        }
    }
}

impl Names for StatementList {
    fn names(&self) -> NameList {
        let mut names = NameList::default();
        self.iter()
            .for_each(|statement| names.merge_in_place(statement.names()));
        names
    }
}

impl Names for Statement {
    fn names(&self) -> NameList {
        match self {
            // Names of these will be collected directly from symbol
            Statement::Workbench(_) => NameList::default(),
            Statement::Module(_) => NameList::default(),
            Statement::Function(_) => NameList::default(),
            Statement::InnerAttribute(_) => NameList::default(),

            Statement::Init(i) => i.names().drop_locals(),
            Statement::If(i) => i.names().drop_locals(),

            Statement::Use(u) => u.names(),
            Statement::Return(r) => r.names(),
            Statement::Assignment(a) => a.names(),
            Statement::Expression(e) => e.names(),
        }
    }
}

impl Names for WorkbenchDefinition {
    fn names(&self) -> NameList {
        self.plan.names().add_as_name(&self.id)
    }
}

impl Names for ParameterList {
    fn names(&self) -> NameList {
        NameList::from_iter(self.iter().map(|param| &param.id)).merge_many(
            self.iter()
                .filter_map(|param| param.default_value.as_ref())
                .map(|expr| expr.names()),
        )
    }
}

impl Names for FunctionDefinition {
    fn names(&self) -> NameList {
        self.signature.names()
    }
}

impl Names for FunctionSignature {
    fn names(&self) -> NameList {
        self.parameters.names()
    }
}

impl Names for InitDefinition {
    fn names(&self) -> NameList {
        self.parameters.names()
    }
}

impl Names for ReturnStatement {
    fn names(&self) -> NameList {
        if let Some(result) = &self.result {
            result.names()
        } else {
            NameList::default()
        }
    }
}

impl Names for IfStatement {
    fn names(&self) -> NameList {
        let mut result = NameList::default();
        result.merge_in_place(self.cond.names());
        result.merge_in_place(self.body.names());
        result
    }
}

impl Names for AssignmentStatement {
    fn names(&self) -> NameList {
        self.assignment.names()
    }
}

impl Names for Assignment {
    fn names(&self) -> NameList {
        let names = self.expression.names();
        if matches!(self.qualifier(), Qualifier::Const) {
            names.add_as_name(&self.id)
        } else {
            names.add_local(&self.id)
        }
    }
}

impl Names for ExpressionStatement {
    fn names(&self) -> NameList {
        self.expression.names()
    }
}

impl Names for Expression {
    fn names(&self) -> NameList {
        match self {
            Expression::Invalid | Expression::Literal(_) | Expression::Marker(_) => {
                NameList::default()
            }

            Expression::FormatString(fs) => fs.names(),
            Expression::ArrayExpression(ae) => ae.names(),
            Expression::TupleExpression(te) => te.names(),
            Expression::Body(body) => body.names(),
            Expression::Call(call) => call.names(),
            Expression::QualifiedName(name) => name.into(),
            Expression::BinaryOp {
                lhs, op: _, rhs, ..
            } => lhs.names().merge(rhs.names()),
            Expression::UnaryOp { op: _, rhs, .. } => rhs.names(),
            Expression::ArrayElementAccess(e, e1, ..) => e.names().merge(e1.names()),
            Expression::PropertyAccess(e, ..) => e.names(),
            Expression::AttributeAccess(e, ..) => e.names(),
            Expression::MethodCall(e, mc, ..) => e.names().merge(mc.names()),
        }
    }
}

impl Names for Body {
    fn names(&self) -> NameList {
        self.statements.names()
    }
}

impl Names for UseStatement {
    fn names(&self) -> NameList {
        match &self.decl {
            UseDeclaration::Use(name) | UseDeclaration::UseAll(name) => name.into(),
            UseDeclaration::UseAs(name, _id) => NameList::default().add_name(name),
        }
    }
}

impl Names for FormatString {
    fn names(&self) -> NameList {
        NameList::default().merge_many(self.0.iter().filter_map(|inner| {
            if let FormatStringInner::FormatExpression(expr) = inner {
                Some(expr.expression.names())
            } else {
                None
            }
        }))
    }
}

impl Names for ArrayExpression {
    fn names(&self) -> NameList {
        self.inner.names()
    }
}

impl Names for ArrayExpressionInner {
    fn names(&self) -> NameList {
        match self {
            ArrayExpressionInner::List(expressions) => {
                NameList::default().merge_many(expressions.iter().map(|expr| expr.names()))
            }
            ArrayExpressionInner::Range(range_expression) => range_expression
                .first
                .names()
                .merge(range_expression.last.names()),
        }
    }
}

impl Names for RangeFirst {
    fn names(&self) -> NameList {
        self.0.names()
    }
}

impl Names for RangeLast {
    fn names(&self) -> NameList {
        self.0.names()
    }
}

impl Names for TupleExpression {
    fn names(&self) -> NameList {
        self.args.names()
    }
}

impl Names for Call {
    fn names(&self) -> NameList {
        self.argument_list.names().add_name(&self.name)
    }
}

impl Names for ArgumentList {
    fn names(&self) -> NameList {
        NameList::default().merge_many(
            self.iter()
                // get expressions out of arguments
                .map(|arg| arg.expression.as_ref())
                .map(|expr| expr.names()),
        )
    }
}

impl Names for MethodCall {
    fn names(&self) -> NameList {
        self.argument_list.names().add_name(&self.name)
    }
}
