unit Myc.Ast.Refactoring.Remove;

interface

uses
    System.SysUtils,
    System.Classes,
    System.Generics.Collections,
    Myc.Ast,
    Myc.Ast.Nodes,
    Myc.Ast.Visitor,
    Myc.Ast.Types;

type
    // Removes a specific node from the AST.
    // - If the node is part of a list, it is removed from the list.
    // - If the node is a mandatory child of a structure (e.g. If-Condition),
    //   it is replaced by a Nop node to maintain tree integrity.
    TAstNodeRemover = class(TAstTransformer)
    private
        FTarget: IAstNode;
        FRemoved: Boolean;

        // Generic helper to filter nil items from lists during transformation
        function FilterList<T: IAstNode>(const Source: INodeList<T>): TArray<T>;

        // Helper to ensure a mandatory slot is never nil (replaces with Nop)
        function EnsureNode(const Node: IAstNode; const ContextNode: IAstNode): IAstNode;
    protected
        function Accept(const Node: IAstNode): IAstNode; override;

        // --- List Containers (Filter Logic) ---
        function VisitBlockExpression(const Node: IBlockExpressionNode): IAstNode; override;
        function VisitExpressionList(const Node: IExpressionList): IAstNode; override;
        function VisitArgumentList(const Node: IArgumentList): IAstNode; override;
        function VisitParameterList(const Node: IParameterList): IAstNode; override;
        function VisitRecordFieldList(const Node: IRecordFieldList): IAstNode; override;

        // --- Structural Nodes (Replace Logic) ---
        function VisitIfExpression(const Node: IIfExpressionNode): IAstNode; override;
        function VisitCondExpression(const Node: ICondExpressionNode): IAstNode; override; // Replaces Ternary
        function VisitAssignment(const Node: IAssignmentNode): IAstNode; override;
        function VisitVariableDeclaration(const Node: IVariableDeclarationNode): IAstNode; override;
        function VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode; override;
        function VisitFunctionCall(const Node: IFunctionCallNode): IAstNode; override;
        function VisitIndexer(const Node: IIndexerNode): IAstNode; override;
        function VisitMemberAccess(const Node: IMemberAccessNode): IAstNode; override;
        function VisitAddSeriesItem(const Node: IAddSeriesItemNode): IAstNode; override;

    public
        constructor Create(const ATarget: IAstNode);

        // Tries to remove TargetNode from RootNode.
        // Returns True if the node was found and removed/replaced.
        // If successful, NewRootNode contains the modified AST.
        // If not found, NewRootNode returns the original RootNode.
        class function TryRemove(const RootNode, TargetNode: IAstNode; out NewRootNode: IAstNode): Boolean; static;
    end;

implementation

{ TAstNodeRemover }

constructor TAstNodeRemover.Create(const ATarget: IAstNode);
begin
    inherited Create;
    FTarget := ATarget;
    FRemoved := False;
end;

class function TAstNodeRemover.TryRemove(const RootNode, TargetNode: IAstNode; out NewRootNode: IAstNode): Boolean;
var
    Remover: TAstNodeRemover;
begin
    // Edge case: Removing the root itself
    if (RootNode = TargetNode) then
    begin
        // We replace the root with a Nop, as we cannot return 'nil' for a valid AST root expectation usually
        NewRootNode := TAst.Nop(RootNode.Identity);
        Result := True;
        Exit;
    end;

    Remover := TAstNodeRemover.Create(TargetNode);
    try
        NewRootNode := Remover.Accept(RootNode);
        Result := Remover.FRemoved;

        if not Result then
        begin
            // If nothing was removed, return the original node to ensure reference stability
            NewRootNode := RootNode;
        end
        else if (NewRootNode = nil) then
        begin
            // Safety fallback: if traversal somehow resulted in nil (should be caught by EnsureNode logic usually)
            NewRootNode := TAst.Nop(RootNode.Identity);
        end;
    finally
        Remover.Free;
    end;
end;

function TAstNodeRemover.Accept(const Node: IAstNode): IAstNode;
begin
    // If we matched the target, flag it as removed and return nil.
    // The parent visitor is responsible for handling the nil result (Filter or Replace).
    if Node = FTarget then
    begin
        FRemoved := True;
        Exit(nil);
    end;

    Result := inherited Accept(Node);
end;

function TAstNodeRemover.EnsureNode(const Node: IAstNode; const ContextNode: IAstNode): IAstNode;
begin
    if Assigned(Node) then
        Result := Node
    else
        // Mandatory child was deleted -> Replace with Nop using parent's identity location
        Result := TAst.Nop(ContextNode.Identity);
end;

function TAstNodeRemover.FilterList<T>(const Source: INodeList<T>): TArray<T>;
var
    i: Integer;
    item: T;
    transformed: IAstNode;
    list: TList<T>;
begin
    list := TList<T>.Create;
    try
        for i := 0 to Source.Count - 1 do
        begin
            item := Source[i];
            transformed := Accept(item);

            // If transformed is not nil, keep it.
            // If it is nil, it was the target, so we skip (delete) it.
            if Assigned(transformed) then
                list.Add(T(transformed));
        end;
        Result := list.ToArray;
    finally
        list.Free;
    end;
end;

// -----------------------------------------------------------------------------
// List Visitors (Filtering)
// -----------------------------------------------------------------------------

function TAstNodeRemover.VisitBlockExpression(const Node: IBlockExpressionNode): IAstNode;
var
    newExprs: TArray<IAstNode>;
    newList: IExpressionList;
begin
    newExprs := FilterList<IAstNode>(Node.Expressions);

    // If counts differ, something was removed
    if Length(newExprs) = Node.Expressions.Count then
    begin
        // If not marked as removed, we assume no change in this branch
        if not FRemoved then
            Exit(Node);

        // Deep check for replacement
        var changed := False;
        for var i := 0 to High(newExprs) do
            if newExprs[i] <> Node.Expressions[i] then
                changed := True;

        if not changed then
            Exit(Node);
    end;

    newList := TExpressionList.Create(newExprs, Node.Expressions.Identity);
    Result := TAst.Block(Node.Identity, newList, Node.StaticType);
end;

function TAstNodeRemover.VisitExpressionList(const Node: IExpressionList): IAstNode;
var
    newExprs: TArray<IAstNode>;
begin
    newExprs := FilterList<IAstNode>(Node);

    if Length(newExprs) = Node.Count then
    begin
        var changed := False;
        for var i := 0 to High(newExprs) do
            if newExprs[i] <> Node[i] then
                changed := True;
        if not changed then
            Exit(Node);
    end;

    Result := TExpressionList.Create(newExprs, Node.Identity);
end;

function TAstNodeRemover.VisitArgumentList(const Node: IArgumentList): IAstNode;
var
    newArgs: TArray<IAstNode>;
begin
    newArgs := FilterList<IAstNode>(Node);

    if Length(newArgs) = Node.Count then
    begin
        var changed := False;
        for var i := 0 to High(newArgs) do
            if newArgs[i] <> Node[i] then
                changed := True;
        if not changed then
            Exit(Node);
    end;

    Result := TArgumentList.Create(newArgs, Node.Identity);
end;

function TAstNodeRemover.VisitParameterList(const Node: IParameterList): IAstNode;
var
    newParams: TArray<IIdentifierNode>;
begin
    newParams := FilterList<IIdentifierNode>(Node);

    if Length(newParams) = Node.Count then
    begin
        var changed := False;
        for var i := 0 to High(newParams) do
            if newParams[i] <> Node[i] then
                changed := True;
        if not changed then
            Exit(Node);
    end;

    Result := TParameterList.Create(newParams, Node.Identity);
end;

function TAstNodeRemover.VisitRecordFieldList(const Node: IRecordFieldList): IAstNode;
var
    newFields: TArray<IRecordFieldNode>;
begin
    newFields := FilterList<IRecordFieldNode>(Node);

    if Length(newFields) = Node.Count then
    begin
        var changed := False;
        for var i := 0 to High(newFields) do
            if newFields[i] <> Node[i] then
                changed := True;
        if not changed then
            Exit(Node);
    end;

    Result := TRecordFieldList.Create(newFields, Node.Identity);
end;

// -----------------------------------------------------------------------------
// Structural Visitors (Replacement with Nop/Void)
// -----------------------------------------------------------------------------

function TAstNodeRemover.VisitIfExpression(const Node: IIfExpressionNode): IAstNode;
var
    newCond, newThen, newElse: IAstNode;
begin
    newCond := Accept(Node.Condition);
    newThen := Accept(Node.ThenBranch);
    newElse := Accept(Node.ElseBranch);

    // If condition or Then-branch is deleted, replace with Nop
    newCond := EnsureNode(newCond, Node);
    newThen := EnsureNode(newThen, Node);

    // Else branch is optional
    if (newCond = Node.Condition) and (newThen = Node.ThenBranch) and (newElse = Node.ElseBranch) then
        Exit(Node);

    Result := TAst.IfExpr(Node.Identity, newCond, newThen, newElse, Node.StaticType);
end;

function TAstNodeRemover.VisitCondExpression(const Node: ICondExpressionNode): IAstNode;
var
    i: Integer;
    newPairs: TArray<TCondPair>;
    newElse: IAstNode;
    newCond, newBranch: IAstNode;
    hasChanged: Boolean;
begin
    SetLength(newPairs, Length(Node.Pairs));
    hasChanged := False;

    for i := 0 to High(Node.Pairs) do
    begin
        // Ensure mandatory children exist (replace with Nop if removed)
        newCond := EnsureNode(Accept(Node.Pairs[i].Condition), Node);
        newBranch := EnsureNode(Accept(Node.Pairs[i].Branch), Node);

        newPairs[i] := TCondPair.Create(newCond, newBranch);

        if (newCond <> Node.Pairs[i].Condition) or (newBranch <> Node.Pairs[i].Branch) then
            hasChanged := True;
    end;

    // Else branch is effectively mandatory in the structure, even if it's Nop/Void
    newElse := EnsureNode(Accept(Node.ElseBranch), Node);
    if newElse <> Node.ElseBranch then
        hasChanged := True;

    if not hasChanged then
        Exit(Node);

    Result := TAst.CondExpr(Node.Identity, newPairs, newElse, Node.StaticType);
end;

function TAstNodeRemover.VisitAssignment(const Node: IAssignmentNode): IAstNode;
var
    newTarget, newValue: IAstNode;
begin
    newTarget := EnsureNode(Accept(Node.Target), Node);
    newValue := EnsureNode(Accept(Node.Value), Node);

    if (newTarget = Node.Target) and (newValue = Node.Value) then
        Exit(Node);

    Result := TAst.Assign(Node.Identity, newTarget, newValue, Node.StaticType);
end;

function TAstNodeRemover.VisitVariableDeclaration(const Node: IVariableDeclarationNode): IAstNode;
var
    newTarget, newInit: IAstNode;
begin
    newTarget := EnsureNode(Accept(Node.Target), Node);
    newInit := Accept(Node.Initializer); // Initializer is optional

    if (newTarget = Node.Target) and (newInit = Node.Initializer) then
        Exit(Node);

    Result := TAst.VarDecl(Node.Identity, newTarget, newInit, Node.StaticType, Node.IsBoxed);
end;

function TAstNodeRemover.VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode;
var
    newParams: IAstNode;
    newBody: IAstNode;
begin
    newParams := Accept(Node.Parameters);
    // If the parameter list itself was the target, we replace it with an empty list
    if newParams = nil then
        newParams := TParameterList.Create([], Node.Parameters.Identity);

    newBody := EnsureNode(Accept(Node.Body), Node);

    if (newParams = Node.Parameters) and (newBody = Node.Body) then
        Exit(Node);

    Result :=
        TAst.LambdaExpr(
            Node.Identity,
            newParams.AsParameterList,
            newBody,
            Node.Layout,
            Node.Descriptor,
            Node.Upvalues,
            Node.HasNestedLambdas,
            Node.IsPure,
            Node.StaticType
        );
end;

function TAstNodeRemover.VisitFunctionCall(const Node: IFunctionCallNode): IAstNode;
var
    newCallee: IAstNode;
    newArgs: IAstNode;
begin
    newCallee := EnsureNode(Accept(Node.Callee), Node);

    newArgs := Accept(Node.Arguments);
    if newArgs = nil then
        newArgs := TArgumentList.Create([], Node.Arguments.Identity);

    if (newCallee = Node.Callee) and (newArgs = Node.Arguments) then
        Exit(Node);

    Result :=
        TAst.FunctionCall(
            Node.Identity,
            newCallee,
            newArgs.AsArgumentList,
            Node.StaticType,
            Node.IsTailCall,
            Node.StaticTarget,
            Node.IsTargetPure
        );
end;

function TAstNodeRemover.VisitIndexer(const Node: IIndexerNode): IAstNode;
var
    newBase, newIndex: IAstNode;
begin
    newBase := EnsureNode(Accept(Node.Base), Node);
    newIndex := EnsureNode(Accept(Node.Index), Node);

    if (newBase = Node.Base) and (newIndex = Node.Index) then
        Exit(Node);

    Result := TAst.Indexer(Node.Identity, newBase, newIndex, Node.StaticType);
end;

function TAstNodeRemover.VisitMemberAccess(const Node: IMemberAccessNode): IAstNode;
var
    newBase: IAstNode;
    newMember: IAstNode;
begin
    newBase := EnsureNode(Accept(Node.Base), Node);

    newMember := Accept(Node.Member);
    if newMember = nil then
    begin
        // If member name is deleted, the access is invalid. Replace whole node with Nop.
        Exit(TAst.Nop(Node.Identity));
    end;

    if (newBase = Node.Base) and (newMember = Node.Member) then
        Exit(Node);

    Result := TAst.MemberAccess(Node.Identity, newBase, newMember.AsKeyword, Node.StaticType);
end;

function TAstNodeRemover.VisitAddSeriesItem(const Node: IAddSeriesItemNode): IAstNode;
var
    newSeries: IAstNode;
    newValue: IAstNode;
    newLookback: IAstNode;
begin
    newSeries := Accept(Node.Series);
    if newSeries = nil then
        Exit(TAst.Nop(Node.Identity)); // Cannot exist without series identifier

    newValue := EnsureNode(Accept(Node.Value), Node);
    newLookback := Accept(Node.Lookback); // Optional

    if (newSeries = Node.Series) and (newValue = Node.Value) and (newLookback = Node.Lookback) then
        Exit(Node);

    Result := TAst.AddSeriesItem(Node.Identity, newSeries.AsIdentifier, newValue, newLookback, Node.StaticType);
end;

end.