unit Myc.Ast.Compiler.Macros;

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.Environment; // Added for IMacroRegistry

type
    IAstMacroExpander = interface(IAstVisitor)
        // Signatur geaendert: Descriptor wird nicht mehr zurueckgegeben
        function Execute(const RootNode: IAstNode): 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;
        FRenameMap: TDictionary<string, string>; // [*] Hygiene Rename Map (Rule 1)
        FDefinitionDescriptor: IScopeDescriptor; // [*] Template Scope (Rule 1)
        class var
            FGensymCounter: Int64;

        function TransformAndSpliceNodes(const ANodes: TArray<IAstNode>): TArray<IAstNode>;

        // [*] Gensym helper
        function Gensym(const ABaseName: string): string;
        // [*] Scope management for template
        procedure EnterScope;
        procedure ExitScope;
    protected
        // Override the new IAstNode-returning methods
        function VisitUnquote(const Node: IUnquoteNode): IAstNode; override;
        function VisitUnquoteSplicing(const Node: IUnquoteSplicingNode): IAstNode; override;
        function VisitFunctionCall(const Node: IFunctionCallNode): IAstNode; override;
        function VisitBlockExpression(const Node: IBlockExpressionNode): IAstNode; override;
        function VisitRecordLiteral(const Node: IRecordLiteralNode): IAstNode; override;

        // [*] Hygiene implementation (Rule 1)
        function VisitVariableDeclaration(const Node: IVariableDeclarationNode): IAstNode; override;
        function VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode; override;

        // [*] Hygiene implementation (Rule 1, 2, 3)
        function VisitIdentifier(const Node: IIdentifierNode): IAstNode; override;
    public
        // [*] Updated constructor
        constructor Create(const AMacroScope: IExecutionScope; const AEvaluate: TEvaluateProc);
        destructor Destroy; override; // [*] Added

        // [*] Updated static helper
        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;

        FCurrentMacroRegistry: IMacroRegistry; // Changed from TMacroRegistry

        FEvaluatorFactory: TEvaluatorFactory;

        procedure EnterMacroScope;
        procedure ExitMacroScope;

    protected
        // Finds macro definitions
        function VisitMacroDefinition(const Node: IMacroDefinitionNode): IAstNode; override;
        // Finds and expands macro calls
        function VisitFunctionCall(const Node: IFunctionCallNode): IAstNode; override;

        // [*] Do not traverse into quasiquotes (they are handled by TExpansionVisitor)
        function VisitQuasiquote(const Node: IQuasiquoteNode): IAstNode; override;
        function VisitUnquote(const Node: IUnquoteNode): IAstNode; override;
        function VisitUnquoteSplicing(const Node: IUnquoteSplicingNode): IAstNode; override;

        function VisitBlockExpression(const Node: IBlockExpressionNode): IAstNode; override;
        function VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode; override;

    public
        // Updated Constructor
        constructor Create(
            const ARootRegistry: IMacroRegistry;
            const AInitialScope: IExecutionScope;
            const AEvaluatorFactory: TEvaluatorFactory
        );
        destructor Destroy; override; // Modified

        function Execute(const RootNode: IAstNode): IAstNode;

        // Updated static helper
        class function ExpandMacros(
            const ARootRegistry: IMacroRegistry; // Added
            const AInitialScope: IExecutionScope;
            const RootNode: IAstNode;
            const EvaluatorFactory: TEvaluatorFactory
        ): IAstNode; static;
    end;

implementation

uses
    System.Generics.Defaults,
    System.SyncObjs,
    Myc.Ast.Compiler.Binder,
    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;

    // [*] Initialize hygiene tools
    FRenameMap := TDictionary<string, string>.Create;

    // Create a descriptor for the *template* scope, parented to nil.
    // It only tracks template-internal definitions for shadowing.
    FDefinitionDescriptor := TScope.CreateDescriptor(nil);
end;

destructor TExpansionVisitor.Destroy;
begin
    FRenameMap.Free;
    // FDefinitionDescriptor is an interface
    inherited;
end;

procedure TExpansionVisitor.EnterScope;
begin
    // [*] Create a new nested definition scope
    FDefinitionDescriptor := TScope.CreateDescriptor(FDefinitionDescriptor);
end;

procedure TExpansionVisitor.ExitScope;
begin
    // [*] Revert to parent definition scope
    FDefinitionDescriptor := FDefinitionDescriptor.Parent;
end;

function TExpansionVisitor.Gensym(const ABaseName: string): string;
begin
    if not FRenameMap.TryGetValue(ABaseName, Result) then
    begin
        var count := TInterlocked.Add(FGensymCounter, 1);
        inc(count);
        Result := ABaseName + '_' + count.ToString;
        FRenameMap.Add(ABaseName, Result);
    end;
end;

class function TExpansionVisitor.Expand(
    const MacroScope: IExecutionScope;
    const RootNode: IAstNode;
    const AEvaluate: TEvaluateProc
): IAstNode;
begin
    var expander := TExpansionVisitor.Create(MacroScope, AEvaluate);
    try
        Result := expander.Accept(RootNode); // Use IAstNode-returning Accept
    finally
        expander.Free;
    end;
end;

function TExpansionVisitor.TransformAndSpliceNodes(const ANodes: TArray<IAstNode>): TArray<IAstNode>;
var
    newList: TList<IAstNode>;
    nodeToSplice: IAstNode;
    transformedNode: IAstNode;
begin
    newList := TList<IAstNode>.Create;
    try
        for var node in ANodes do
        begin
            if node.Kind = akUnquoteSplicing then
            begin
                var spliceExpr := node.AsUnquoteSplicing.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.Kind = akBlockExpression then
                        newList.AddRange(nodeToSplice.AsBlockExpression.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
                // [*] Use the IAstNode-returning Accept helper (hygienic transform)
                transformedNode := Self.Accept(node);
                if Assigned(transformedNode) then
                    newList.Add(transformedNode);
            end;
        end;
        Result := newList.ToArray;
    finally
        newList.Free;
    end;
end;

function TExpansionVisitor.VisitIdentifier(const Node: IIdentifierNode): IAstNode;
var
    newName: string;
begin
    // [*] Hygiene Rule 1: Check for local rename
    // Is this symbol defined *within the template* (and thus renamed)?
    if FRenameMap.TryGetValue(Node.Name, newName) then
    begin
        // Yes, replace it with the hygienic name (e.g. 'start-time$1')
        // We use TTypes.Unknown because the Binder/TypeChecker run *after* this.
        Result := TAst.Identifier(newName, TTypes.Unknown);
        exit;
    end;

    // [*] Hygiene Rule 2 & 3: Free Symbol (Hygienic 'print' or Unhygienic '*debug-mode*')
    // No, it's a free symbol. Leave it as-is.
    // The TAstBinder will resolve it later, either in the
    // macro's definition scope (Rule 2) or the caller's scope (Rule 3).
    Result := Node;
end;

function TExpansionVisitor.VisitVariableDeclaration(const Node: IVariableDeclarationNode): IAstNode;
var
    newName: string;
    newIdent: IIdentifierNode;
    newInit: IAstNode;
begin
    // [*] Hygiene Rule 1: This is a definition within the template.

    // 1. Visit initializer *before* defining the var
    // (so initializer can't see the var it's defining)
    newInit := Accept(Node.Initializer);

    // 2. Generate hygienic name
    newName := Gensym(Node.Identifier.Name);
    FDefinitionDescriptor.Define(newName);

    // 3. Create new identifier node (type is unknown)
    newIdent := TAst.Identifier(newName, TTypes.Unknown);

    // 4. Create new VarDecl node
    Result := TAst.VarDecl(newIdent, newInit, TTypes.Unknown);
end;

function TExpansionVisitor.VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode;
var
    newParams: TArray<IIdentifierNode>;
    newBody: IAstNode;
    newName: string;
    i: Integer;
begin
    // [*] Hygiene Rule 1: Lambda parameters are definitions.

    // 1. Enter a new definition scope for the template
    EnterScope;
    try
        // 2. Rename parameters
        SetLength(newParams, Length(Node.Parameters));
        for i := 0 to High(Node.Parameters) do
        begin
            newName := Gensym(Node.Parameters[i].Name);
            FDefinitionDescriptor.Define(newName);
            newParams[i] := TAst.Identifier(newName, TTypes.Unknown);
        end;

        // 3. Visit body recursively in the new scope
        newBody := Accept(Node.Body);

    finally
        // 4. Exit definition scope
        ExitScope;
    end;

    // 5. Create new Lambda node (type is unknown)
    Result := TAst.LambdaExpr(newParams, newBody, nil, nil, False, TTypes.Unknown);
end;

function TExpansionVisitor.VisitUnquote(const Node: IUnquoteNode): IAstNode;
var
    value: TDataValue;
    expr: IAstNode;
    symbol: TResolvedSymbol;
begin
    // [*] This node *stops* the hygienic transformation.
    // It reverts to the previous compile-time evaluation logic.

    expr := Node.Expression;

    // Check if we are unquoting a macro parameter (simple identifier)
    if expr.Kind = akIdentifier then
    begin
        symbol := FMacroScope.Descriptor.FindSymbol(expr.AsIdentifier.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.AsIntf<IAstNode>; // Return the node directly
                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.AsIntf<IAstNode>;
        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 := 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): IAstNode;
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): IAstNode;
var
    newArgs: TArray<IAstNode>;
    transformedCallee: IAstNode;
begin
    // [*] Callee must be transformed hygienically
    transformedCallee := Self.Accept(Node.Callee); // Calls IAstNode helper
    // [*] Arguments must be transformed and allow splicing
    newArgs := TransformAndSpliceNodes(Node.Arguments);
    Result := TAst.FunctionCall(transformedCallee, newArgs);
end;

function TExpansionVisitor.VisitBlockExpression(const Node: IBlockExpressionNode): IAstNode;
var
    newExprs: TArray<IAstNode>;
begin
    // [*] Expressions must be transformed and allow splicing
    newExprs := TransformAndSpliceNodes(Node.Expressions);
    Result := TAst.Block(newExprs);
end;

function TExpansionVisitor.VisitRecordLiteral(const Node: IRecordLiteralNode): IAstNode;
begin
    // Record literals do not support splicing.
    // We just use the default TAstTransformer implementation which visits child values.
    // [*] This will correctly call VisitIdentifier for keys/values.
    Result := inherited VisitRecordLiteral(Node);
end;

{ TMacroExpander }

constructor TMacroExpander.Create(
    const ARootRegistry: IMacroRegistry;
    const AInitialScope: IExecutionScope;
    const AEvaluatorFactory: TEvaluatorFactory
);
begin
    inherited Create;
    Assert(Assigned(ARootRegistry));
    Assert(Assigned(AInitialScope));
    Assert(Assigned(AEvaluatorFactory));

    FInitialScope := AInitialScope;
    FEvaluatorFactory := AEvaluatorFactory;

    // Creates the root registry for this specific compilation run,
    // parented to the global registry from the environment.
    FCurrentMacroRegistry := ARootRegistry.CreateChildRegistry;
end;

destructor TMacroExpander.Destroy;
begin
    // FCurrentMacroRegistry is an interface, managed by ARC
    inherited Destroy;
end;

procedure TMacroExpander.EnterMacroScope;
begin
    // Creates a new registry with the current as Parent
    FCurrentMacroRegistry := FCurrentMacroRegistry.CreateChildRegistry;
end;

procedure TMacroExpander.ExitMacroScope;
begin
    // Revert to parent registry. ARC will free the old child.
    FCurrentMacroRegistry := FCurrentMacroRegistry.Parent;
end;

class function TMacroExpander.ExpandMacros(
    const ARootRegistry: IMacroRegistry;
    const AInitialScope: IExecutionScope;
    const RootNode: IAstNode;
    const EvaluatorFactory: TEvaluatorFactory
): IAstNode;
begin
    // IAstMacroExpander.Execute gibt keinen Descriptor mehr zurueck
    var expander := TMacroExpander.Create(ARootRegistry, AInitialScope, EvaluatorFactory) as IAstMacroExpander;
    Result := expander.Execute(RootNode);
end;

function TMacroExpander.Execute(const RootNode: IAstNode): IAstNode;
begin
    // This executes the transformation (Accept)
    Result := Accept(RootNode); // Calls the IAstNode-returning helper

    if not Assigned(Result) then
        Result := TAst.Block([]); // e.g. if root was (defmacro ...)
end;

function TMacroExpander.VisitBlockExpression(const Node: IBlockExpressionNode): IAstNode;
begin
    EnterMacroScope;
    try
        Result := inherited VisitBlockExpression(Node);
    finally
        ExitMacroScope;
    end;
end;

function TMacroExpander.VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode;
begin
    EnterMacroScope;
    try
        Result := inherited VisitLambdaExpression(Node);
    finally
        ExitMacroScope;
    end;
end;

function TMacroExpander.VisitMacroDefinition(const Node: IMacroDefinitionNode): IAstNode;
begin
    // Register the macro in the *current* macro registry
    FCurrentMacroRegistry.Define(Node);

    // Keep the node in the AST (for the visualizer).
    Result := Node;
end;

function TMacroExpander.VisitFunctionCall(const Node: IFunctionCallNode): IAstNode;
var
    calleeIdentifier: IIdentifierNode;
    macroDef: IMacroDefinitionNode;
    i: Integer;
begin
    // Check if this is an identifier
    if Node.Callee.Kind <> akIdentifier then
        exit(inherited VisitFunctionCall(Node));

    calleeIdentifier := Node.Callee.AsIdentifier;

    // Check if this identifier is a macro in the current registry
    macroDef := FCurrentMacroRegistry.Find(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
    //    It must be parented to the *initial* scope to find other globals.
    var expansionScope := TScope.CreateScope(FInitialScope, nil, 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])); // Use FromIntf

    // 3. Create the evaluation callback (TEvaluateProc)
    //    It must use the FInitialScope (for globals) and the expansionScope (for args)
    var evaluatorProc: TEvaluateProc;
    evaluatorProc :=
        function(const ANodeToEvaluate: IAstNode): TDataValue
        var
            subDescriptor: IScopeDescriptor;
            evalScope: IExecutionScope;
            evaluator: IEvaluatorVisitor;
            boundSubAst: IAstNode;
        begin
            // This is the compile-time evaluation (Binder + Evaluator)
            // [*] It must bind against the *expansion* scope, not the initial scope
            boundSubAst := TAstBinder.Bind(expansionScope.Descriptor, ANodeToEvaluate, subDescriptor);

            // Create eval scope from the new descriptor, parented to the expansion scope
            evalScope := subDescriptor.CreateScope(expansionScope);

            evaluator := FEvaluatorFactory(evalScope);
            Result := evaluator.Execute(boundSubAst);
        end;

    // 4. Expand the macro body using the TExpansionVisitor
    // [*] FInitialScope removed from call
    var expandedBody := TExpansionVisitor.Expand(expansionScope, macroDef.Body.AsQuasiquote.Expression, evaluatorProc);

    // 5. Wrap the result in a MacroExpansionNode (for debugging/tracing)
    var macroNode := TAst.MacroExpansionNode(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); // Calls the IAstNode helper, returns IAstNode
end;

function TMacroExpander.VisitQuasiquote(const Node: IQuasiquoteNode): IAstNode;
begin
    // [*] Do not traverse into nested quasiquotes.
    // TExpansionVisitor handles them when (if) they are unquoted.
    Result := Node;
end;

function TMacroExpander.VisitUnquote(const Node: IUnquoteNode): IAstNode;
begin
    // [*] Standalone unquote (outside a quasiquote) is a syntax error.
    raise Exception.Create('Unquote (`~`) can only be used inside a quasiquote (` `) template.');
end;

function TMacroExpander.VisitUnquoteSplicing(const Node: IUnquoteSplicingNode): IAstNode;
begin
    // [*] Standalone unquote-splicing (outside a quasiquote) is a syntax error.
    raise Exception.Create('Unquote-splicing (`~@`) can only be used inside a quasiquote (` `) template.');
end;

end.