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.
    TAstNodeRemover = class(TAstTransformer)
    private
        FTarget: IAstNode;
        FRemoved: Boolean;

        function FilterTuple(const Source: ITupleNode): TArray<IAstNode>;
        function EnsureNode(const Node: IAstNode; const ContextNode: IAstNode): IAstNode;

    strict private
        // Structural Nodes (Replace Logic)
        function VisitIfExpression(const Node: IAstNode): IAstNode;
        function VisitCondExpression(const Node: IAstNode): IAstNode;
        function VisitAssignment(const Node: IAstNode): IAstNode;
        function VisitVariableDeclaration(const Node: IAstNode): IAstNode;
        function VisitLambdaExpression(const Node: IAstNode): IAstNode;
        function VisitFunctionCall(const Node: IAstNode): IAstNode;
        function VisitIndexer(const Node: IAstNode): IAstNode;
        function VisitMemberAccess(const Node: IAstNode): IAstNode;
        function VisitAddSeriesItem(const Node: IAstNode): IAstNode;

        // Unified Container Handler
        function VisitTuple(const Node: IAstNode): IAstNode;
        function VisitBlockExpression(const Node: IAstNode): IAstNode;
        function VisitRecurNode(const Node: IAstNode): IAstNode;

    protected
        procedure SetupHandlers; override;
        function Accept(const Node: IAstNode): IAstNode; override;

    public
        constructor Create(const ATarget: IAstNode);
        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;

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

    // Replace Logic
    Register(akIfExpression, VisitIfExpression);
    Register(akCondExpression, VisitCondExpression);
    Register(akAssignment, VisitAssignment);
    Register(akVariableDeclaration, VisitVariableDeclaration);
    Register(akLambdaExpression, VisitLambdaExpression);
    Register(akFunctionCall, VisitFunctionCall);
    Register(akIndexer, VisitIndexer);
    Register(akMemberAccess, VisitMemberAccess);
    Register(akAddSeriesItem, VisitAddSeriesItem);
    Register(akBlockExpression, VisitBlockExpression);
    Register(akRecur, VisitRecurNode);

    // List Logic
    Register(akTuple, VisitTuple);
end;

class function TAstNodeRemover.TryRemove(const RootNode, TargetNode: IAstNode; out NewRootNode: IAstNode): Boolean;
var
    Remover: TAstNodeRemover;
begin
    if (RootNode = TargetNode) then
    begin
        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
            NewRootNode := RootNode
        else if (NewRootNode = nil) then
            NewRootNode := TAst.Nop(RootNode.Identity);
    finally
        Remover.Free;
    end;
end;

function TAstNodeRemover.Accept(const Node: IAstNode): IAstNode;
begin
    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
        Result := TAst.Nop(ContextNode.Identity);
end;

function TAstNodeRemover.FilterTuple(const Source: ITupleNode): TArray<IAstNode>;
var
    i: Integer;
    item, transformed: IAstNode;
    list: TList<IAstNode>;
    elements: TArray<IAstNode>;
begin
    list := TList<IAstNode>.Create;
    try
        elements := Source.Elements;
        for i := 0 to High(elements) do
        begin
            item := elements[i];
            transformed := Accept(item);
            if Assigned(transformed) then
                list.Add(transformed);
        end;
        Result := list.ToArray;
    finally
        list.Free;
    end;
end;

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

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

    if (newCond = E.Condition) and (newThen = E.ThenBranch) and (newElse = E.ElseBranch) then
        Exit(Node);
    Result := TAst.IfExpr(Node.Identity, newCond, newThen, newElse, E.StaticType);
end;

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

    for i := 0 to High(E.Pairs) do
    begin
        newCond := EnsureNode(Accept(E.Pairs[i].Condition), Node);
        newBranch := EnsureNode(Accept(E.Pairs[i].Branch), Node);
        newPairs[i] := TCondPair.Create(newCond, newBranch);
        if (newCond <> E.Pairs[i].Condition) or (newBranch <> E.Pairs[i].Branch) then
            hasChanged := True;
    end;

    newElse := EnsureNode(Accept(E.ElseBranch), Node);
    if newElse <> E.ElseBranch then
        hasChanged := True;

    if not hasChanged then
        Exit(Node);
    Result := TAst.CondExpr(Node.Identity, newPairs, newElse, E.StaticType);
end;

function TAstNodeRemover.VisitAssignment(const Node: IAstNode): IAstNode;
var
    A: IAssignmentNode;
    newTarget, newValue: IAstNode;
begin
    A := Node.AsAssignment;
    newTarget := EnsureNode(Accept(A.Target), Node);
    newValue := EnsureNode(Accept(A.Value), Node);
    if (newTarget = A.Target) and (newValue = A.Value) then
        Exit(Node);
    Result := TAst.Assign(Node.Identity, newTarget, newValue, A.StaticType);
end;

function TAstNodeRemover.VisitVariableDeclaration(const Node: IAstNode): IAstNode;
var
    V: IVariableDeclarationNode;
    newTarget, newInit: IAstNode;
begin
    V := Node.AsVariableDeclaration;
    newTarget := EnsureNode(Accept(V.Target), Node);
    newInit := Accept(V.Initializer);
    if (newTarget = V.Target) and (newInit = V.Initializer) then
        Exit(Node);
    Result := TAst.VarDecl(Node.Identity, newTarget, newInit, V.StaticType, V.IsBoxed);
end;

function TAstNodeRemover.VisitLambdaExpression(const Node: IAstNode): IAstNode;
var
    L: ILambdaExpressionNode;
    newParams: TArray<IAstNode>;
    newBody: IAstNode;
    paramTuple: ITupleNode;
begin
    L := Node.AsLambdaExpression;

    // Filter parameters (removing a parameter is semantically risky but consistent with "Remove" logic)
    newParams := FilterTuple(L.Parameters);

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

    // Factory creates new Tuple from Array
    paramTuple := TAst.Tuple(L.Parameters.Identity, newParams);

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

function TAstNodeRemover.VisitFunctionCall(const Node: IAstNode): IAstNode;
var
    C: IFunctionCallNode;
    newCallee: IAstNode;
    newArgs: TArray<IAstNode>;
    argsTuple: ITupleNode;
begin
    C := Node.AsFunctionCall;
    newCallee := EnsureNode(Accept(C.Callee), Node);
    newArgs := FilterTuple(C.Arguments);

    argsTuple := TAst.Tuple(C.Arguments.Identity, newArgs);

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

function TAstNodeRemover.VisitIndexer(const Node: IAstNode): IAstNode;
var
    I: IIndexerNode;
    newBase, newIndex: IAstNode;
begin
    I := Node.AsIndexer;
    newBase := EnsureNode(Accept(I.Base), Node);
    newIndex := EnsureNode(Accept(I.Index), Node);
    if (newBase = I.Base) and (newIndex = I.Index) then
        Exit(Node);
    Result := TAst.Indexer(Node.Identity, newBase, newIndex, I.StaticType);
end;

function TAstNodeRemover.VisitMemberAccess(const Node: IAstNode): IAstNode;
var
    M: IMemberAccessNode;
    newBase, newMember: IAstNode;
begin
    M := Node.AsMemberAccess;
    newBase := EnsureNode(Accept(M.Base), Node);
    newMember := Accept(M.Member);
    if newMember = nil then
        Exit(TAst.Nop(Node.Identity)); // Member missing -> invalid access

    if (newBase = M.Base) and (newMember = M.Member) then
        Exit(Node);
    Result := TAst.MemberAccess(Node.Identity, newBase, newMember.AsKeyword, M.StaticType);
end;

function TAstNodeRemover.VisitAddSeriesItem(const Node: IAstNode): IAstNode;
var
    A: IAddSeriesItemNode;
    newSeries, newValue, newLookback: IAstNode;
begin
    A := Node.AsAddSeriesItem;
    newSeries := Accept(A.Series);
    if newSeries = nil then
        Exit(TAst.Nop(Node.Identity));
    newValue := EnsureNode(Accept(A.Value), Node);
    newLookback := Accept(A.Lookback);

    if (newSeries = A.Series) and (newValue = A.Value) and (newLookback = A.Lookback) then
        Exit(Node);
    Result := TAst.AddSeriesItem(Node.Identity, newSeries.AsIdentifier, newValue, newLookback, A.StaticType);
end;

function TAstNodeRemover.VisitBlockExpression(const Node: IAstNode): IAstNode;
var
    B: IBlockExpressionNode;
    newExprs: TArray<IAstNode>;
    exprsTuple: ITupleNode;
begin
    B := Node.AsBlockExpression;
    newExprs := FilterTuple(B.Expressions);

    exprsTuple := TAst.Tuple(B.Expressions.Identity, newExprs);
    Result := TAst.Block(Node.Identity, exprsTuple, B.StaticType);
end;

function TAstNodeRemover.VisitRecurNode(const Node: IAstNode): IAstNode;
var
    R: IRecurNode;
    newArgs: TArray<IAstNode>;
    argsTuple: ITupleNode;
begin
    R := Node.AsRecur;
    newArgs := FilterTuple(R.Arguments);
    argsTuple := TAst.Tuple(R.Arguments.Identity, newArgs);
    Result := TAst.Recur(Node.Identity, argsTuple, R.StaticType);
end;

function TAstNodeRemover.VisitTuple(const Node: IAstNode): IAstNode;
var
    T: ITupleNode;
    newElements: TArray<IAstNode>;
    list: TList<IAstNode>;
    item, transformed: IAstNode;
    i: Integer;
    hasChanged: Boolean;
    elements: TArray<IAstNode>;
begin
    T := Node.AsTuple;
    list := TList<IAstNode>.Create;
    hasChanged := False;
    try
        elements := T.Elements;
        for i := 0 to High(elements) do
        begin
            item := elements[i];
            transformed := Accept(item);

            if transformed <> item then
                hasChanged := True;

            if Assigned(transformed) then
                list.Add(transformed)
            else
                hasChanged := True; // Item removed
        end;

        if not hasChanged then
            Exit(Node);

        newElements := list.ToArray;
        Result := TAst.Tuple(Node.Identity, newElements, T.StaticType);
    finally
        list.Free;
    end;
end;

end.