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;

type
    EMacroException = class(EAstException);

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

    TMacroEvaluatorProc = reference to function(const Scope: IExecutionScope; const Node: IAstNode): TDataValue;

    // Handles body expansion (Template instantiation)
    TExpansionVisitor = class(TAstTransformer)
    strict private
        FMacroEvaluator: TMacroEvaluatorProc;
        FMacroScope: IExecutionScope;
        FRenameMap: TDictionary<string, string>;
        class var
            FGensymCounter: Int64;

        function TransformAndSpliceNodes(const ANodes: ITupleNode): TArray<IAstNode>;
        function Gensym(const ABaseName: string): string;

        // Custom Logic Handlers (IAstNode signature)
        function VisitUnquote(const Node: IAstNode): IAstNode;
        function VisitUnquoteSplicing(const Node: IAstNode): IAstNode;
        function VisitFunctionCall(const Node: IAstNode): IAstNode;
        function VisitBlockExpression(const Node: IAstNode): IAstNode;
        function VisitRecordLiteral(const Node: IAstNode): IAstNode;
        function VisitVariableDeclaration(const Node: IAstNode): IAstNode;
        function VisitLambdaExpression(const Node: IAstNode): IAstNode;
        function VisitIdentifier(const Node: IAstNode): IAstNode;

    protected
        procedure SetupHandlers; override;

    public
        constructor Create(const AMacroScope: IExecutionScope; const AMacroEvaluator: TMacroEvaluatorProc);
        destructor Destroy; override;

        class function Expand(
            const MacroScope: IExecutionScope;
            const RootNode: IAstNode;
            const MacroEvaluator: TMacroEvaluatorProc
        ): IAstNode;
    end;

    IMacroRegistry = interface
        {$region 'private'}
        function GetParent: IMacroRegistry;
        {$endregion}
        procedure Define(const Node: IMacroDefinitionNode);
        function Find(const Name: string): IMacroDefinitionNode;
        function CreateChildRegistry: IMacroRegistry;
        property Parent: IMacroRegistry read GetParent;
    end;

    // Handles macro calls in AST
    TMacroExpander = class(TAstTransformer, IAstMacroExpander)
    strict private
        FInitialScope: IExecutionScope;
        FCurrentMacroRegistry: IMacroRegistry;
        FMacroEvaluator: TMacroEvaluatorProc;

        procedure EnterMacroScope;
        procedure ExitMacroScope;

        function VisitMacroDefinition(const Node: IAstNode): IAstNode;
        function VisitFunctionCall(const Node: IAstNode): IAstNode;
        function VisitQuasiquote(const Node: IAstNode): IAstNode;
        function VisitUnquote(const Node: IAstNode): IAstNode;
        function VisitUnquoteSplicing(const Node: IAstNode): IAstNode;
        function VisitBlockExpression(const Node: IAstNode): IAstNode;
        function VisitLambdaExpression(const Node: IAstNode): IAstNode;

    protected
        procedure SetupHandlers; override;

    public
        constructor Create(
            const ARootRegistry: IMacroRegistry;
            const AInitialScope: IExecutionScope;
            const AMacroEvaluator: TMacroEvaluatorProc
        );
        destructor Destroy; override;

        function Execute(const RootNode: IAstNode): IAstNode;

        class function ExpandMacros(
            const ARootRegistry: IMacroRegistry;
            const AInitialScope: IExecutionScope;
            const RootNode: IAstNode;
            const MacroEvaluator: TMacroEvaluatorProc
        ): IAstNode; static;
    end;

    TMacroRegistry = class(TInterfacedObject, IMacroRegistry)
    private
        FParent: IMacroRegistry;
        FMacros: TDictionary<string, IMacroDefinitionNode>;
        function GetParent: IMacroRegistry;
    public
        constructor Create(AParent: IMacroRegistry);
        destructor Destroy; override;
        procedure Define(const Node: IMacroDefinitionNode);
        function Find(const Name: string): IMacroDefinitionNode;
        function CreateChildRegistry: IMacroRegistry;
    end;

implementation

uses
    System.Generics.Defaults,
    System.SyncObjs,
    Myc.Data.Keyword;

{ TExpansionVisitor }

constructor TExpansionVisitor.Create(const AMacroScope: IExecutionScope; const AMacroEvaluator: TMacroEvaluatorProc);
begin
    inherited Create;
    FMacroScope := AMacroScope;
    FMacroEvaluator := AMacroEvaluator;
    FRenameMap := TDictionary<string, string>.Create;
end;

destructor TExpansionVisitor.Destroy;
begin
    FRenameMap.Free;
    inherited;
end;

procedure TExpansionVisitor.SetupHandlers;
begin
    inherited SetupHandlers; // Defaults

    // Register custom handlers
    Register(akUnquote, VisitUnquote);
    Register(akUnquoteSplicing, VisitUnquoteSplicing);

    Register(akFunctionCall, VisitFunctionCall);
    Register(akBlockExpression, VisitBlockExpression);
    Register(akRecordLiteral, VisitRecordLiteral);
    Register(akVariableDeclaration, VisitVariableDeclaration);
    Register(akLambdaExpression, VisitLambdaExpression);
    Register(akIdentifier, VisitIdentifier);
end;

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

class function TExpansionVisitor.Expand(
    const MacroScope: IExecutionScope;
    const RootNode: IAstNode;
    const MacroEvaluator: TMacroEvaluatorProc
): IAstNode;
begin
    var expander := TExpansionVisitor.Create(MacroScope, MacroEvaluator);
    try
        Result := expander.Accept(RootNode);
    finally
        expander.Free;
    end;
end;

function TExpansionVisitor.TransformAndSpliceNodes(const ANodes: ITupleNode): TArray<IAstNode>;
var
    newList: TList<IAstNode>;
    nodeToSplice: IAstNode;
    transformedNode: IAstNode;
    node: IAstNode;
    i: Integer;
begin
    newList := TList<IAstNode>.Create;
    try
        for i := 0 to ANodes.Count - 1 do
        begin
            node := ANodes.Items[i];
            if node.Kind = akUnquoteSplicing then
            begin
                var spliceExpr := node.AsUnquoteSplicing.Expression;
                var evaluatedSpliceValue: TDataValue;
                try
                    evaluatedSpliceValue := FMacroEvaluator(FMacroScope, spliceExpr);
                except
                    on E: EAstException do
                        raise;
                    on E: Exception do
                        raise EMacroException.Create('Error during macro splice evaluation: ' + E.Message);
                end;

                if (not evaluatedSpliceValue.IsVoid) and (evaluatedSpliceValue.Kind = vkInterface) then
                begin
                    nodeToSplice := evaluatedSpliceValue.AsIntf<IAstNode>;
                    // Unification: Lists are now Tuples.
                    if nodeToSplice.Kind = akTuple then
                    begin
                        var tuple := nodeToSplice.AsTuple;
                        for var k := 0 to tuple.Count - 1 do
                            newList.Add(tuple.Items[k]);
                    end
                    else if nodeToSplice.Kind = akBlockExpression then
                    begin
                        var blkTuple := nodeToSplice.AsBlockExpression.Expressions;
                        for var k := 0 to blkTuple.Count - 1 do
                            newList.Add(blkTuple.Items[k]);
                    end
                    else
                    begin
                        // Single node splicing
                        newList.Add(nodeToSplice);
                    end;
                end
                else if (not evaluatedSpliceValue.IsVoid) then
                    newList.Add(TAst.Constant(evaluatedSpliceValue, node.Identity.Location));
            end
            else
            begin
                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: IAstNode): IAstNode;
var
    I: IIdentifierNode;
    newName: string;
begin
    I := Node.AsIdentifier;
    if FRenameMap.TryGetValue(I.Name, newName) then
    begin
        Result := TAst.Identifier(newName, I.Identity.Location);
        exit;
    end;
    Result := Node;
end;

function TExpansionVisitor.VisitVariableDeclaration(const Node: IAstNode): IAstNode;
var
    V: IVariableDeclarationNode;
    newTarget, newInit: IAstNode;
    newName: string;
begin
    V := Node.AsVariableDeclaration;
    newInit := Accept(V.Initializer);
    if V.Target.Kind = akIdentifier then
    begin
        newName := Gensym(V.Target.AsIdentifier.Name);
        newTarget := TAst.Identifier(newName, V.Target.Identity.Location);
    end
    else
        newTarget := Accept(V.Target);

    Result := TAst.VarDecl(Node.Identity, newTarget, newInit, TTypes.Unknown, False);
end;

function TExpansionVisitor.VisitLambdaExpression(const Node: IAstNode): IAstNode;
var
    L: ILambdaExpressionNode;
    newParams: TArray<IAstNode>;
    newBody: IAstNode;
    newName: string;
    i: Integer;
    paramsTuple: ITupleNode;
begin
    L := Node.AsLambdaExpression;
    paramsTuple := L.Parameters;
    SetLength(newParams, paramsTuple.Count);

    for i := 0 to paramsTuple.Count - 1 do
    begin
        var param := paramsTuple.Items[i].AsIdentifier; // Binder ensures params are Identifiers
        newName := Gensym(param.Name);
        newParams[i] := TAst.Identifier(newName, param.Identity.Location);
    end;

    newBody := Accept(L.Body);

    Result := TAst.LambdaExpr(Node.Identity, TAst.Tuple(L.Parameters.Identity, newParams), newBody);
end;

function TExpansionVisitor.VisitUnquote(const Node: IAstNode): IAstNode;
var
    U: IUnquoteNode;
    value: TDataValue;
    expr: IAstNode;
    addr: TResolvedAddress;
begin
    U := Node.AsUnquote;
    expr := U.Expression;
    if expr.Kind = akIdentifier then
    begin
        addr := FMacroScope.Resolve(expr.AsIdentifier.Name);
        if (addr.Kind = akLocalOrParent) and (addr.ScopeDepth = 0) then
        begin
            var argValue := FMacroScope.Values[addr];
            if argValue.Kind = vkInterface then
                exit(argValue.AsIntf<IAstNode>);
        end;
    end;

    try
        value := FMacroEvaluator(FMacroScope, expr);
    except
        on E: EAstException do
            raise;
        on E: Exception do
            raise EMacroException.Create('Error during macro evaluation: ' + E.Message);
    end;

    if value.Kind = vkInterface then
        exit(value.AsIntf<IAstNode>);

    if value.Kind in [vkScalar, vkText, vkVoid] then
        Result := TAst.Constant(value, Node.Identity.Location)
    else
        raise EMacroException.CreateFmt('Cannot unquote complex runtime value of type %s at compile time.', [value.Kind.ToString]);
end;

function TExpansionVisitor.VisitUnquoteSplicing(const Node: IAstNode): IAstNode;
begin
    raise EMacroException.Create('Unquote-splicing (`~@`) can only appear inside a list/tuple form.');
end;

function TExpansionVisitor.VisitFunctionCall(const Node: IAstNode): IAstNode;
var
    C: IFunctionCallNode;
    newArgs: TArray<IAstNode>;
    transformedCallee: IAstNode;
    argsTuple: ITupleNode;
begin
    C := Node.AsFunctionCall;
    transformedCallee := Self.Accept(C.Callee);
    // Arguments is a Tuple, so we can splice into it
    newArgs := TransformAndSpliceNodes(C.Arguments);

    // Wrap result array in a TupleNode before creating FunctionCall
    argsTuple := TAst.Tuple(C.Arguments.Identity, newArgs);

    Result := TAst.FunctionCall(Node.Identity, transformedCallee, argsTuple);
end;

function TExpansionVisitor.VisitBlockExpression(const Node: IAstNode): IAstNode;
var
    B: IBlockExpressionNode;
    newExprs: TArray<IAstNode>;
    exprsTuple: ITupleNode;
begin
    B := Node.AsBlockExpression;
    // Expressions is a Tuple
    newExprs := TransformAndSpliceNodes(B.Expressions);

    exprsTuple := TAst.Tuple(B.Expressions.Identity, newExprs);

    Result := TAst.Block(Node.Identity, exprsTuple);
end;

function TExpansionVisitor.VisitRecordLiteral(const Node: IAstNode): IAstNode;
begin
    // Delegate to standard transformer to process fields recursively
    // This calls inherited VisitRecordLiteral(IAstNode) which visits the Fields Tuple
    Result := inherited VisitRecordLiteral(Node);
end;

{ TMacroExpander }

constructor TMacroExpander.Create(
    const ARootRegistry: IMacroRegistry;
    const AInitialScope: IExecutionScope;
    const AMacroEvaluator: TMacroEvaluatorProc
);
begin
    inherited Create;
    FInitialScope := AInitialScope;
    FMacroEvaluator := AMacroEvaluator;
    FCurrentMacroRegistry := ARootRegistry.CreateChildRegistry;
end;

destructor TMacroExpander.Destroy;
begin
    inherited Destroy;
end;

procedure TMacroExpander.SetupHandlers;
begin
    inherited SetupHandlers; // Defaults

    Register(akMacroDefinition, VisitMacroDefinition);
    Register(akFunctionCall, VisitFunctionCall);
    Register(akQuasiquote, VisitQuasiquote);
    Register(akUnquote, VisitUnquote);
    Register(akUnquoteSplicing, VisitUnquoteSplicing);

    // Wrappers for scope management
    Register(akBlockExpression, VisitBlockExpression);
    Register(akLambdaExpression, VisitLambdaExpression);
end;

procedure TMacroExpander.EnterMacroScope;
begin
    FCurrentMacroRegistry := FCurrentMacroRegistry.CreateChildRegistry;
end;

procedure TMacroExpander.ExitMacroScope;
begin
    FCurrentMacroRegistry := FCurrentMacroRegistry.Parent;
end;

class function TMacroExpander.ExpandMacros(
    const ARootRegistry: IMacroRegistry;
    const AInitialScope: IExecutionScope;
    const RootNode: IAstNode;
    const MacroEvaluator: TMacroEvaluatorProc
): IAstNode;
begin
    var expander := TMacroExpander.Create(ARootRegistry, AInitialScope, MacroEvaluator) as IAstMacroExpander;
    Result := expander.Execute(RootNode);
end;

function TMacroExpander.Execute(const RootNode: IAstNode): IAstNode;
begin
    Result := Accept(RootNode);
    if not Assigned(Result) then
        Result := TAst.Block([], nil);
end;

function TMacroExpander.VisitBlockExpression(const Node: IAstNode): IAstNode;
begin
    EnterMacroScope;
    try
        // Delegate to standard recursive transform
        Result := inherited VisitBlockExpression(Node);
    finally
        ExitMacroScope;
    end;
end;

function TMacroExpander.VisitLambdaExpression(const Node: IAstNode): IAstNode;
begin
    EnterMacroScope;
    try
        // Delegate to standard recursive transform
        Result := inherited VisitLambdaExpression(Node);
    finally
        ExitMacroScope;
    end;
end;

function TMacroExpander.VisitMacroDefinition(const Node: IAstNode): IAstNode;
begin
    FCurrentMacroRegistry.Define(Node.AsMacroDefinition);
    Result := Node;
end;

function TMacroExpander.VisitFunctionCall(const Node: IAstNode): IAstNode;
var
    C: IFunctionCallNode;
    calleeIdentifier: IIdentifierNode;
    macroDef: IMacroDefinitionNode;
    i: Integer;
begin
    C := Node.AsFunctionCall;
    if C.Callee.Kind = akIdentifier then
    begin
        calleeIdentifier := C.Callee.AsIdentifier;
        macroDef := FCurrentMacroRegistry.Find(calleeIdentifier.Name);

        if macroDef <> nil then
        begin
            var expansionScope := TScope.CreateScope(FInitialScope, nil, nil);
            var paramsTuple := macroDef.Parameters;
            var argsTuple := C.Arguments;

            if argsTuple.Count <> paramsTuple.Count then
                raise EMacroException.CreateFmt('Macro %s expects %d arguments.', [calleeIdentifier.Name, paramsTuple.Count]);

            for i := 0 to paramsTuple.Count - 1 do
            begin
                var paramName := paramsTuple.Items[i].AsIdentifier.Name;
                expansionScope.Define(paramName, TDataValue.FromIntf<IAstNode>(argsTuple.Items[i]));
            end;

            var expandedBody := TExpansionVisitor.Expand(expansionScope, macroDef.Body.AsQuasiquote.Expression, FMacroEvaluator);
            var macroNode := TAst.MacroExpansionNode(Node.Identity, C, expandedBody);

            Result := Self.Accept(macroNode);
            exit;
        end;
    end;

    // Standard recursion for non-macro calls
    Result := inherited VisitFunctionCall(Node);
end;

function TMacroExpander.VisitQuasiquote(const Node: IAstNode): IAstNode;
begin
    Result := Node;
end;

function TMacroExpander.VisitUnquote(const Node: IAstNode): IAstNode;
begin
    raise EMacroException.Create('Unquote (`~`) can only be used inside a quasiquote.');
end;

function TMacroExpander.VisitUnquoteSplicing(const Node: IAstNode): IAstNode;
begin
    raise EMacroException.Create('Unquote-splicing (`~@`) can only be used inside a quasiquote.');
end;

{ TMacroRegistry }

constructor TMacroRegistry.Create(AParent: IMacroRegistry);
begin
    inherited Create;
    FParent := AParent;
    FMacros := TDictionary<string, IMacroDefinitionNode>.Create;
end;

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

function TMacroRegistry.GetParent: IMacroRegistry;
begin
    Result := FParent;
end;

procedure TMacroRegistry.Define(const Node: IMacroDefinitionNode);
begin
    FMacros.AddOrSetValue(Node.Name.Name, Node);
end;

function TMacroRegistry.Find(const Name: string): IMacroDefinitionNode;
var
    current: IMacroRegistry;
begin
    current := Self;
    while Assigned(current) do
    begin
        if (current as TMacroRegistry).FMacros.TryGetValue(Name, Result) then
            exit;
        current := current.Parent;
    end;
    Result := nil;
end;

function TMacroRegistry.CreateChildRegistry: IMacroRegistry;
begin
    Result := TMacroRegistry.Create(Self);
end;

end.