unit Myc.Ast.MacroExpander;

interface

uses
    System.SysUtils,
    System.Classes,
    System.Generics.Collections,
    Myc.Data.Value,
    Myc.Ast.Nodes,
    Myc.Ast.Scope,
    Myc.Ast.Visitor;

type
    // Defines the strategy for compiling and evaluating an AST node at compile time.
    TCompileTimeExecutor = reference to function(const ANode: IAstNode): TDataValue;

    IMacroExpander = interface(IAstVisitor)
        function Execute(const RootNode: IAstNode): IAstNode;
    end;

    // The macro expander is now a pure AST transformation engine.
    TMacroExpander = class(TAstTransformer, IMacroExpander)
    private
        FMacros: TDictionary<string, IMacroDefinitionNode>;
        FExecutor: TCompileTimeExecutor;
        function ExpandNode(const ANode: IAstNode): IAstNode;
    protected
        function VisitMacroDefinition(const Node: IMacroDefinitionNode): TDataValue; override;
        function VisitFunctionCall(const Node: IFunctionCallNode): TDataValue; override;
    public
        constructor Create(const AExecutor: TCompileTimeExecutor);
        destructor Destroy; override;
        function Execute(const RootNode: IAstNode): IAstNode;
    end;

implementation

uses
    Myc.Ast;

type
    TExpansionHelper = class(TAstTransformer)
    private
        FExpansionScope: IExecutionScope;
        FExpander: TMacroExpander;
        FExecutor: TCompileTimeExecutor;
    protected
        function VisitUnquote(const Node: IUnquoteNode): TDataValue; override;
        function VisitUnquoteSplicing(const Node: IUnquoteSplicingNode): TDataValue; override;
        function VisitFunctionCall(const Node: IFunctionCallNode): TDataValue; override;
    public
        constructor Create(AExpansionScope: IExecutionScope; AExpander: TMacroExpander; AExecutor: TCompileTimeExecutor);
        class function Expand(
            Node: IAstNode;
            ExpansionScope: IExecutionScope;
            Expander: TMacroExpander;
            Executor: TCompileTimeExecutor
        ): IAstNode; static;
    end;

    TSubstitutor = class(TAstTransformer)
    private
        FParamScope: IExecutionScope;
    protected
        function VisitIdentifier(const Node: IIdentifierNode): TDataValue; override;
    public
        constructor Create(AParamScope: IExecutionScope);
    end;

{ TMacroExpander }

constructor TMacroExpander.Create(const AExecutor: TCompileTimeExecutor);
begin
    inherited Create;
    Assert(Assigned(AExecutor), 'A TCompileTimeExecutor must be provided.');
    FMacros := TDictionary<string, IMacroDefinitionNode>.Create;
    FExecutor := AExecutor;
end;

destructor TMacroExpander.Destroy;
begin
    FMacros.Free;
    inherited;
end;

function TMacroExpander.ExpandNode(const ANode: IAstNode): IAstNode;
var
    lastNode: IAstNode;
begin
    Result := ANode;
    if not Assigned(Result) then
        exit;
    repeat
        lastNode := Result;
        Result := Self.Accept(lastNode).AsIntf<IAstNode>;
    until Result = lastNode;
end;

function TMacroExpander.Execute(const RootNode: IAstNode): IAstNode;
begin
    FMacros.Clear;
    Result := ExpandNode(RootNode);
end;

function TMacroExpander.VisitFunctionCall(const Node: IFunctionCallNode): TDataValue;
var
    macroDef: IMacroDefinitionNode;
    expandedNode: IAstNode;
    i: Integer;
begin
    if not (Node.Callee is TIdentifierNode) then
    begin
        Result := inherited VisitFunctionCall(Node);
        exit;
    end;

    var calleeIdentifier := Node.Callee as TIdentifierNode;

    if FMacros.TryGetValue(calleeIdentifier.Name, macroDef) then
    begin
        var expansionScope := TAst.CreateScope(nil); // Temporary scope for parameters, no parent needed.

        if Length(Node.Arguments) <> Length(macroDef.Parameters) then
            raise Exception.CreateFmt(
                'Macro %s expects %d arguments, but got %d',
                [macroDef.Name.Name, Length(macroDef.Parameters), Length(Node.Arguments)]);

        for i := 0 to High(macroDef.Parameters) do
            expansionScope.Define(macroDef.Parameters[i].Name, TDataValue.FromIntf<IAstNode>(Node.Arguments[i]));

        if not (macroDef.Body is TQuasiquoteNode) then
            raise Exception.CreateFmt('Macro body for "%s" must be a quasiquoted expression.', [macroDef.Name.Name]);

        expandedNode := TExpansionHelper.Expand(IQuasiquoteNode(macroDef.Body).Expression, expansionScope, Self, FExecutor);

        Result := Self.Accept(expandedNode);
    end
    else
    begin
        Result := inherited VisitFunctionCall(Node);
    end;
end;

function TMacroExpander.VisitMacroDefinition(const Node: IMacroDefinitionNode): TDataValue;
begin
    FMacros.AddOrSetValue(Node.Name.Name, Node);
    Result := TDataValue.FromIntf<IBlockExpressionNode>(TBlockExpressionNode.Create([]));
end;

{ TExpansionHelper }

constructor TExpansionHelper.Create(AExpansionScope: IExecutionScope; AExpander: TMacroExpander; AExecutor: TCompileTimeExecutor);
begin
    inherited Create;
    FExpansionScope := AExpansionScope;
    FExpander := AExpander;
    FExecutor := AExecutor;
end;

class function TExpansionHelper.Expand(
    Node: IAstNode;
    ExpansionScope: IExecutionScope;
    Expander: TMacroExpander;
    Executor: TCompileTimeExecutor
): IAstNode;
var
    helper: TExpansionHelper;
begin
    helper := TExpansionHelper.Create(ExpansionScope, Expander, Executor);
    try
        Result := helper.Accept(Node).AsIntf<IAstNode>;
    finally
        helper.Free;
    end;
end;

function TExpansionHelper.VisitFunctionCall(const Node: IFunctionCallNode): TDataValue;
var
    newArgs: TList<IAstNode>;
    transformedArgValue: TDataValue;
    transformedArgNode: IAstNode;
    transformedCallee: IAstNode;
begin
    newArgs := TList<IAstNode>.Create;
    try
        for var arg in Node.Arguments do
        begin
            transformedArgValue := Self.Accept(arg);
            if transformedArgValue.IsVoid then
                continue;
            transformedArgNode := transformedArgValue.AsIntf<IAstNode>;
            if (transformedArgNode is TBlockExpressionNode) then
                newArgs.AddRange((transformedArgNode as TBlockExpressionNode).Expressions)
            else
                newArgs.Add(transformedArgNode);
        end;
        transformedCallee := Self.Accept(Node.Callee).AsIntf<IAstNode>;
        Result := TDataValue.FromIntf<IAstNode>(TAst.FunctionCall(transformedCallee, newArgs.ToArray));
    finally
        newArgs.Free;
    end;
end;

function TExpansionHelper.VisitUnquote(const Node: IUnquoteNode): TDataValue;
var
    substitutor: TSubstitutor;
    substitutedAst, expandedAst: IAstNode;
    value: TDataValue;
begin
    // Step 1: Substitute macro parameters with their AST node values.
    substitutor := TSubstitutor.Create(FExpansionScope);
    try
        substitutedAst := substitutor.Execute(Node.Expression);
    finally
        substitutor.Free;
    end;

    // Step 2: Recursively expand any macros within the substituted expression.
    expandedAst := FExpander.ExpandNode(substitutedAst);

    // Step 3: Call the injected executor strategy to bind and evaluate the AST.
    value := FExecutor(expandedAst);

    // Step 4: Wrap the resulting value for insertion into the template.
    if value.Kind in [vkScalar, vkText] then
        Result := TDataValue.FromIntf<IAstNode>(TAst.Constant(value))
    else
        Result := value;
end;

function TExpansionHelper.VisitUnquoteSplicing(const Node: IUnquoteSplicingNode): TDataValue;
var
    unquotedValue: TDataValue;
    unquotedNode: IAstNode;
begin
    unquotedValue := Self.VisitUnquote(TAst.Unquote(Node.Expression));
    if unquotedValue.Kind = vkInterface then
        unquotedNode := unquotedValue.AsIntf<IAstNode>
    else
        unquotedNode := nil;
    if not (unquotedNode is TBlockExpressionNode) then
        raise Exception.Create('Expression inside unquote-splicing (`~@`) must evaluate to a list of nodes (a block).');
    Result := unquotedValue;
end;

{ TSubstitutor }
constructor TSubstitutor.Create(AParamScope: IExecutionScope);
begin
    inherited Create;
    FParamScope := AParamScope;
end;

function TSubstitutor.VisitIdentifier(const Node: IIdentifierNode): TDataValue;
var
    addr: TResolvedAddress;
    nodeValue: TDataValue;
begin
    addr := FParamScope.CreateDescriptor.FindSymbol(Node.Name);
    if (addr.Kind = akLocalOrParent) and (addr.ScopeDepth = 0) then
    begin
        nodeValue := FParamScope.Values[addr];
        Result := nodeValue;
        exit;
    end;
    Result := inherited VisitIdentifier(Node);
end;

end.