unit Myc.Ast.MacroExpander;

interface

uses
    System.SysUtils,
    System.Classes,
    System.Generics.Collections,
    Myc.Data.Scalar,
    Myc.Data.Value,
    Myc.Ast.Nodes,
    Myc.Ast.Visitor,
    Myc.Ast.Scope,
    Myc.Ast.Types,
    Myc.Ast,
    Myc.Ast.Binding; // Required for TAstBinder.Bind

type
    IAstMacroExpander = interface(IAstVisitor)
        function Execute(const RootNode: IAstNode; out Descriptor: IScopeDescriptor): IAstNode;
    end;

    // Callback type for compile-time evaluation of Unquote expressions
    TEvaluateProc = reference to function(const Node: IAstNode): TDataValue;

    // --- TExpansionVisitor (Moved from Binding) ---
    // This visitor handles the expansion of a *single* macro body (` `...).
    TExpansionVisitor = class(TAstTransformer)
    private
        FEvaluate: TEvaluateProc;
        FMacroScope: IExecutionScope;
        function TransformAndSpliceNodes(const ANodes: TArray<IAstNode>): TArray<IAstNode>;
    protected
        function VisitUnquote(const Node: IUnquoteNode): TDataValue; override;
        function VisitUnquoteSplicing(const Node: IUnquoteSplicingNode): TDataValue; override;
        function VisitFunctionCall(const Node: IFunctionCallNode): TDataValue; override;
        function VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue; override;
        function VisitRecordLiteral(const Node: IRecordLiteralNode): TDataValue; override;
    public
        constructor Create(const AMacroScope: IExecutionScope; const AEvaluate: TEvaluateProc);
        class function Expand(const MacroScope: IExecutionScope; const RootNode: IAstNode; const AEvaluate: TEvaluateProc): IAstNode;
    end;

    // --- TMacroExpander (New main class for Phase 1) ---
    // This transformer traverses the entire AST *before* the binder,
    // to find and expand macros.
    TMacroExpander = class(TAstTransformer, IAstMacroExpander)
    private
        FInitialScope: IExecutionScope;
        FCurrentDescriptor: IScopeDescriptor;
        FEvaluatorFactory: TEvaluatorFactory;
    protected
        function Accept(const Node: IAstNode): TDataValue; override;
        // Finds macro definitions
        function VisitMacroDefinition(const Node: IMacroDefinitionNode): TDataValue; override;
        // Finds and expands macro calls
        function VisitFunctionCall(const Node: IFunctionCallNode): TDataValue; override;
    public
        constructor Create(const AInitialScope: IExecutionScope; const AEvaluatorFactory: TEvaluatorFactory);
        function Execute(const RootNode: IAstNode; out Descriptor: IScopeDescriptor): IAstNode;

        class function ExpandMacros(
            const InitialScope: IExecutionScope;
            const RootNode: IAstNode;
            out Descriptor: IScopeDescriptor;
            const EvaluatorFactory: TEvaluatorFactory
        ): IAstNode; static;
    end;

implementation

uses
    System.Generics.Defaults,
    Myc.Ast.Evaluator, // Required for TEvaluatorVisitor.Execute
    Myc.Data.Keyword;

{ TExpansionVisitor }

constructor TExpansionVisitor.Create(const AMacroScope: IExecutionScope; const AEvaluate: TEvaluateProc);
begin
    inherited Create;
    FMacroScope := AMacroScope;
    FEvaluate := AEvaluate;
end;

class function TExpansionVisitor.Expand(
    const MacroScope: IExecutionScope;
    const RootNode: IAstNode;
    const AEvaluate: TEvaluateProc
): IAstNode;
begin
    var expander := TExpansionVisitor.Create(MacroScope, AEvaluate) as IAstTransformer;
    Result := expander.Execute(RootNode);
end;

function TExpansionVisitor.TransformAndSpliceNodes(const ANodes: TArray<IAstNode>): TArray<IAstNode>;
var
    newList: TList<IAstNode>;
    nodeToSplice: IAstNode;
begin
    newList := TList<IAstNode>.Create;
    try
        for var node in ANodes do
        begin
            if (node is TUnquoteSplicingNode) then
            begin
                var spliceExpr := (node as TUnquoteSplicingNode).Expression;
                // Evaluate the unquoted expression
                var evaluatedSpliceValue := FEvaluate(spliceExpr);

                if (not evaluatedSpliceValue.IsVoid) and (evaluatedSpliceValue.Kind = vkInterface) then
                begin
                    nodeToSplice := evaluatedSpliceValue.AsIntf<IAstNode>;
                    if (nodeToSplice is TBlockExpressionNode) then
                        newList.AddRange((nodeToSplice as TBlockExpressionNode).Expressions)
                    else
                        // Allow splicing single nodes, not just blocks
                        newList.Add(nodeToSplice);
                end
                else if (not evaluatedSpliceValue.IsVoid) then
                begin
                    // If the value is not an AST node (e.g. a scalar), wrap it in a constant
                    newList.Add(TAst.Constant(evaluatedSpliceValue));
                end;
                // else (if void, splice nothing)
            end
            else
            begin
                var transformedValue := node.Accept(self);
                if not transformedValue.IsVoid then
                    newList.Add(transformedValue.AsIntf<IAstNode>);
            end;
        end;
        Result := newList.ToArray;
    finally
        newList.Free;
    end;
end;

function TExpansionVisitor.VisitUnquote(const Node: IUnquoteNode): TDataValue;
var
    value: TDataValue;
    expr: IAstNode;
    symbol: TResolvedSymbol;
begin
    expr := Node.Expression;

    // Check if we are unquoting a macro parameter (simple identifier)
    if (expr is TIdentifierNode) then
    begin
        symbol := FMacroScope.CreateDescriptor.FindSymbol((expr as TIdentifierNode).Name);
        if (symbol.Address.Kind = akLocalOrParent) and (symbol.Address.ScopeDepth = 0) then
        begin
            // It's a parameter. Return the TDataValue containing the AST node passed as argument.
            var argValue := FMacroScope.Values[symbol.Address];
            if argValue.Kind = vkInterface then
            begin
                Result := argValue;
                exit;
            end;
        end;
    end;

    // If not a simple parameter, evaluate the expression at compile-time.
    value := FEvaluate(expr);

    // If the result is an AST node (e.g. from a helper function), return it directly.
    if value.Kind = vkInterface then
    begin
        Result := value;
        exit;
    end;

    // If the result is a scalar, text, or void, wrap it in a TConstantNode.
    if value.Kind in [vkScalar, vkText, vkVoid] then
    begin
        Result := TDataValue.FromIntf<IAstNode>(TAst.Constant(value));
    end
    else
        raise Exception.CreateFmt('Cannot unquote complex runtime value of type %s at compile time.', [value.Kind.ToString]);
end;

function TExpansionVisitor.VisitUnquoteSplicing(const Node: IUnquoteSplicingNode): TDataValue;
begin
    // This should be handled by the caller (VisitFunctionCall / VisitBlockExpression)
    raise Exception.Create('Unquote-splicing (`~@`) can only appear inside a list form (e.g., a function call or a `do` block).');
end;

function TExpansionVisitor.VisitFunctionCall(const Node: IFunctionCallNode): TDataValue;
var
    newArgs: TArray<IAstNode>;
    transformedCallee: IAstNode;
begin
    transformedCallee := Self.Accept(Node.Callee).AsIntf<IAstNode>;
    newArgs := TransformAndSpliceNodes(Node.Arguments);
    Result := TDataValue.FromIntf<IFunctionCallNode>(TAst.FunctionCall(transformedCallee, newArgs));
end;

function TExpansionVisitor.VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue;
var
    newExprs: TArray<IAstNode>;
begin
    newExprs := TransformAndSpliceNodes(Node.Expressions);
    Result := TDataValue.FromIntf<IBlockExpressionNode>(TAst.Block(newExprs));
end;

function TExpansionVisitor.VisitRecordLiteral(const Node: IRecordLiteralNode): TDataValue;
begin
    // Record literals do not support splicing.
    // We just use the default TAstTransformer implementation which visits child values.
    Result := inherited VisitRecordLiteral(Node);
end;

{ TMacroExpander }

constructor TMacroExpander.Create(const AInitialScope: IExecutionScope; const AEvaluatorFactory: TEvaluatorFactory);
begin
    inherited Create;
    Assert(Assigned(AInitialScope));
    Assert(Assigned(AEvaluatorFactory));
    FInitialScope := AInitialScope;
    FEvaluatorFactory := AEvaluatorFactory;
    // Create the top-level descriptor for macros
    FCurrentDescriptor := AInitialScope.CreateDescriptor;
end;

class function TMacroExpander.ExpandMacros(
    const InitialScope: IExecutionScope;
    const RootNode: IAstNode;
    out Descriptor: IScopeDescriptor;
    const EvaluatorFactory: TEvaluatorFactory
): IAstNode;
begin
    var expander := TMacroExpander.Create(InitialScope, EvaluatorFactory) as IAstMacroExpander;
    Result := expander.Execute(RootNode, Descriptor);
end;

function TMacroExpander.Execute(const RootNode: IAstNode; out Descriptor: IScopeDescriptor): IAstNode;
begin
    // This executes the transformation (Accept)
    var transformedValue := Accept(RootNode);

    if transformedValue.IsVoid then
        Result := TAst.Block([])
    else
        Result := transformedValue.AsIntf<IAstNode>;

    Descriptor := FCurrentDescriptor;
end;

function TMacroExpander.Accept(const Node: IAstNode): TDataValue;
begin
    // Standard transformer Accept
    if (not Assigned(Node)) or Done then
        exit;
    Result := Node.Accept(Self);
end;

function TMacroExpander.VisitMacroDefinition(const Node: IMacroDefinitionNode): TDataValue;
begin
    // Register the macro in the current descriptor
    FCurrentDescriptor.DefineMacro(Node.Name.Name, Node);
    // Remove the (defmacro ...) node from the AST; it's compile-time only.
    Result := TDataValue.Void;
end;

function TMacroExpander.VisitFunctionCall(const Node: IFunctionCallNode): TDataValue;
var
    calleeIdentifier: TIdentifierNode;
    macroDef: IMacroDefinitionNode;
    i: Integer;
begin
    // Check if this is an identifier
    if not (Node.Callee is TIdentifierNode) then
        exit(inherited VisitFunctionCall(Node));

    calleeIdentifier := Node.Callee as TIdentifierNode;

    // Check if this identifier is a macro
    macroDef := FCurrentDescriptor.FindMacro(calleeIdentifier.Name);
    if macroDef = nil then
        exit(inherited VisitFunctionCall(Node)); // Not a macro, let transformer continue

    // --- It is a macro, expand it ---

    // 1. Create the temporary scope for the macro parameters
    var expansionScope := TAst.CreateScope(nil);
    var params := macroDef.Parameters;
    if Length(Node.Arguments) <> Length(params) then
        raise Exception
            .CreateFmt('Macro %s expects %d arguments, but got %d', [calleeIdentifier.Name, Length(params), Length(Node.Arguments)]);

    // 2. Populate scope: Map parameter names to the *un-evaluated* AST nodes
    for i := 0 to High(params) do
        expansionScope.Define(params[i].Name, TDataValue.FromIntf<IAstNode>(Node.Arguments[i]));

    // 3. Create the evaluation callback (TEvaluateProc)
    var evaluatorProc: TEvaluateProc;
    evaluatorProc :=
        function(const ANodeToEvaluate: IAstNode): TDataValue
        var
            subDescriptor: IScopeDescriptor;
        begin
            // This is the compile-time evaluation (Binder + Evaluator)
            // It runs in the *current* context (FInitialScope + FCurrentDescriptor)
            // not the expansionScope.
            var tempInitScope := TScope.CreateScope(FInitialScope.Parent, FCurrentDescriptor, nil);
            var boundSubAst := TAstBinder.Bind(tempInitScope, ANodeToEvaluate, subDescriptor);
            var evalScope := subDescriptor.CreateScope(tempInitScope);
            var evaluator := FEvaluatorFactory(evalScope);
            Result := evaluator.Execute(boundSubAst);
        end;

    // 4. Expand the macro body using the TExpansionVisitor
    var expandedBody := TExpansionVisitor.Expand(expansionScope, macroDef.Body.Expression, evaluatorProc);

    // 5. Wrap the result in a MacroExpansionNode (for debugging/tracing)
    var macroNode := TMacroExpansionNode.Create(Node, expandedBody);

    // 6. IMPORTANT: Re-run the expander on the *new* AST fragment
    // to handle macros that expand into other macros.
    Result := Self.Accept(macroNode);
end;

end.