unit Myc.Ast.Visitor;

interface

uses
    System.SysUtils,
    System.Generics.Collections,
    Myc.Data.Value,
    Myc.Ast,
    Myc.Ast.Types,
    Myc.Ast.Nodes,
    Myc.Ast.Identities;

type
    TAstVisitor<T> = class;

    TNodeHandler<T> = reference to function(const Node: IAstNode): T;

    TAstVisitor<T> = class abstract(TInterfacedObject, IAstVisitor)
    strict private
        FHandlers: array[TAstNodeKind] of TNodeHandler<T>;

        // IAstVisitor bridge
        function IAstVisitor.Visit = VisitBridge;
        function VisitBridge(const Node: IAstNode): TDataValue;

        // Sentinel to fill gaps (Fail Fast protection)
        function DefaultSentinel(const Node: IAstNode): T;

    protected
        // --- Registration ---

        // Register a handler.
        procedure Register(Kind: TAstNodeKind; const Handler: TNodeHandler<T>); virtual;

        // Subclasses implement this to call Register()
        procedure SetupHandlers; virtual;

        // --- Dispatch ---

        function Visit(const Node: IAstNode): T; virtual;
        function Accept(const Node: IAstNode): T; virtual;

    public
        procedure AfterConstruction; override;
        procedure BeforeDestruction; override;
    end;

    // Helper for transformation visitors (returns IAstNode)
    // Implements recursive reconstruction: If a child changes, a new node is returned, but node.Identity stays the same.
    TAstTransformer = class abstract(TAstVisitor<IAstNode>)
    protected
        // Leaves
        function VisitLeaf(const N: IAstNode): IAstNode;

        // Elements
        function VisitRecordField(const N: IAstNode): IAstNode;

        // Structures
        function VisitIfExpression(const N: IAstNode): IAstNode;
        function VisitCondExpression(const N: IAstNode): IAstNode;
        function VisitLambdaExpression(const N: IAstNode): IAstNode;
        function VisitFunctionCall(const N: IAstNode): IAstNode;
        function VisitMacroExpansionNode(const N: IAstNode): IAstNode;
        function VisitBlockExpression(const N: IAstNode): IAstNode;
        function VisitVariableDeclaration(const N: IAstNode): IAstNode;
        function VisitAssignment(const N: IAstNode): IAstNode;
        function VisitMacroDefinition(const N: IAstNode): IAstNode;

        // Wrappers
        function VisitQuasiquote(const N: IAstNode): IAstNode;
        function VisitUnquote(const N: IAstNode): IAstNode;
        function VisitUnquoteSplicing(const N: IAstNode): IAstNode;

        function VisitIndexer(const N: IAstNode): IAstNode;
        function VisitMemberAccess(const N: IAstNode): IAstNode;
        function VisitRecordLiteral(const N: IAstNode): IAstNode;
        function VisitAddSeriesItem(const N: IAstNode): IAstNode;
        function VisitRecurNode(const N: IAstNode): IAstNode;

        // Unified List Type
        function VisitTuple(const N: IAstNode): IAstNode;

        // Pipes
        function VisitPipe(const N: IAstNode): IAstNode;

        procedure SetupHandlers; override;
    end;

    TVoid = record
    end;

    // A concrete "Walker" that traverses all children by default.
    // Useful for Analyzers, Linters, or Searchers.
    TAstVisitor = class(TAstVisitor<TVoid>)
    protected
        // Walker implementations
        function WalkLeaf(const N: IAstNode): TVoid;

        function WalkRecordField(const N: IAstNode): TVoid;

        function WalkIfExpression(const N: IAstNode): TVoid;
        function WalkCondExpression(const N: IAstNode): TVoid;
        function WalkLambdaExpression(const N: IAstNode): TVoid;
        function WalkFunctionCall(const N: IAstNode): TVoid;
        function WalkMacroExpansionNode(const N: IAstNode): TVoid;
        function WalkBlockExpression(const N: IAstNode): TVoid;
        function WalkVariableDeclaration(const N: IAstNode): TVoid;
        function WalkAssignment(const N: IAstNode): TVoid;
        function WalkMacroDefinition(const N: IAstNode): TVoid;
        function WalkWrapper(const N: IAstNode): TVoid; // Quasiquote, Unquote, etc.
        function WalkIndexer(const N: IAstNode): TVoid;
        function WalkMemberAccess(const N: IAstNode): TVoid;
        function WalkRecordLiteral(const N: IAstNode): TVoid;
        function WalkAddSeriesItem(const N: IAstNode): TVoid;
        function WalkRecurNode(const N: IAstNode): TVoid;

        function WalkTuple(const N: IAstNode): TVoid;

        function WalkPipe(const N: IAstNode): TVoid;

        procedure SetupHandlers; override;
    end;

implementation

{ TAstVisitor<T> }

procedure TAstVisitor<T>.AfterConstruction;
var
    k: TAstNodeKind;
begin
    inherited;

    // 1. Initialize with Sentinel (Fail Fast / Safety)
    for k := Low(TAstNodeKind) to High(TAstNodeKind) do
        FHandlers[k] := DefaultSentinel;

    // 2. Allow subclass to overwrite with real handlers
    SetupHandlers;
end;

procedure TAstVisitor<T>.BeforeDestruction;
begin
    for var i := Low(TAstNodeKind) to High(TAstNodeKind) do
        FHandlers[i] := nil;
    inherited;
end;

function TAstVisitor<T>.DefaultSentinel(const Node: IAstNode): T;
begin
    Assert(false, 'No visitor handler registered for AST node kind: ' + Node.Kind.ToString + ' in class ' + ClassName);
    exit(Default(T));
end;

procedure TAstVisitor<T>.SetupHandlers;
begin
    // Base implementation does nothing.
    // Subclasses must call Register() here.
end;

procedure TAstVisitor<T>.Register(Kind: TAstNodeKind; const Handler: TNodeHandler<T>);
begin
    if Assigned(Handler) then
        FHandlers[Kind] := Handler
    else
        FHandlers[Kind] := DefaultSentinel;
end;

function TAstVisitor<T>.Visit(const Node: IAstNode): T;
begin
    if Assigned(Node) then
        Result := FHandlers[Node.Kind](Node)
    else
        Result := Default(T);
end;

function TAstVisitor<T>.Accept(const Node: IAstNode): T;
begin
    Result := Visit(Node);
end;

function TAstVisitor<T>.VisitBridge(const Node: IAstNode): TDataValue;
begin
    Result := TDataValue.FromGeneric<T>(Visit(Node));
end;

{ TAstTransformer }

procedure TAstTransformer.SetupHandlers;
begin
    // Leaves (Identity)
    Register(akConstant, VisitLeaf);
    Register(akIdentifier, VisitLeaf);
    Register(akKeyword, VisitLeaf);
    Register(akNop, VisitLeaf);
    Register(akCreateSeries, VisitLeaf);

    // Elements
    Register(akRecordField, VisitRecordField);

    // Structures
    Register(akIfExpression, VisitIfExpression);
    Register(akCondExpression, VisitCondExpression);
    Register(akLambdaExpression, VisitLambdaExpression);
    Register(akFunctionCall, VisitFunctionCall);
    Register(akMacroExpansion, VisitMacroExpansionNode);
    Register(akBlockExpression, VisitBlockExpression);
    Register(akVariableDeclaration, VisitVariableDeclaration);
    Register(akAssignment, VisitAssignment);
    Register(akMacroDefinition, VisitMacroDefinition);

    Register(akQuasiquote, VisitQuasiquote);
    Register(akUnquote, VisitUnquote);
    Register(akUnquoteSplicing, VisitUnquoteSplicing);

    Register(akIndexer, VisitIndexer);
    Register(akMemberAccess, VisitMemberAccess);
    Register(akRecordLiteral, VisitRecordLiteral);
    Register(akAddSeriesItem, VisitAddSeriesItem);
    Register(akRecur, VisitRecurNode);

    // Unified List Type
    Register(akTuple, VisitTuple);

    // Pipes
    Register(akPipe, VisitPipe);
end;

function TAstTransformer.VisitLeaf(const N: IAstNode): IAstNode;
begin
    Result := N;
end;

function TAstTransformer.VisitRecordField(const N: IAstNode): IAstNode;
var
    F: IRecordFieldNode;
    newKey: IKeywordNode;
    newValue: IAstNode;
begin
    F := N.AsRecordField;
    newKey := Visit(F.Key).AsKeyword;
    newValue := Visit(F.Value);

    if (newKey = F.Key) and (newValue = F.Value) then
        Result := N
    else
        Result := TAst.RecordField(N.Identity, newKey, newValue);
end;

// --- Structures ---

function TAstTransformer.VisitIfExpression(const N: IAstNode): IAstNode;
var
    E: IIfExpressionNode;
    c, t, eBr: IAstNode;
begin
    E := N.AsIfExpression;
    c := Visit(E.Condition);
    t := Visit(E.ThenBranch);
    eBr := Visit(E.ElseBranch);

    if (c = E.Condition) and (t = E.ThenBranch) and (eBr = E.ElseBranch) then
        Result := N
    else
        Result := TAst.IfExpr(N.Identity, c, t, eBr, E.StaticType);
end;

function TAstTransformer.VisitCondExpression(const N: IAstNode): IAstNode;
var
    E: ICondExpressionNode;
    hasChanged: Boolean;
    newPairs: TArray<TCondPair>;
    newElse, c, b: IAstNode;
    i: Integer;
begin
    E := N.AsCondExpression;
    hasChanged := False;
    SetLength(newPairs, Length(E.Pairs));
    for i := 0 to High(E.Pairs) do
    begin
        c := Visit(E.Pairs[i].Condition);
        b := Visit(E.Pairs[i].Branch);
        newPairs[i] := TCondPair.Create(c, b);
        if (c <> E.Pairs[i].Condition) or (b <> E.Pairs[i].Branch) then
            hasChanged := True;
    end;

    newElse := Visit(E.ElseBranch);
    if newElse <> E.ElseBranch then
        hasChanged := True;

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

function TAstTransformer.VisitLambdaExpression(const N: IAstNode): IAstNode;
var
    E: ILambdaExpressionNode;
    p: ITupleNode;
    b: IAstNode;
begin
    E := N.AsLambdaExpression;
    p := Visit(E.Parameters).AsTuple;
    b := Visit(E.Body);

    if (p = E.Parameters) and (b = E.Body) then
        Result := N
    else
        Result :=
            TLambdaExpressionNode.Create(p, b, E.StaticType, E.Layout, E.Descriptor, E.Upvalues, E.HasNestedLambdas, E.IsPure, N.Identity);
end;

function TAstTransformer.VisitFunctionCall(const N: IAstNode): IAstNode;
var
    C: IFunctionCallNode;
    callee: IAstNode;
    args: ITupleNode;
begin
    C := N.AsFunctionCall;
    callee := Visit(C.Callee);
    args := Visit(C.Arguments).AsTuple;

    if (callee = C.Callee) and (args = C.Arguments) then
        Result := N
    else
        Result := TAst.FunctionCall(N.Identity, callee, args, C.StaticType, C.IsTailCall, C.StaticTarget, C.IsTargetPure);
end;

function TAstTransformer.VisitMacroExpansionNode(const N: IAstNode): IAstNode;
var
    M: IMacroExpansionNode;
    exp: IAstNode;
begin
    M := N.AsMacroExpansion;
    exp := Visit(M.ExpandedBody);
    if exp = M.ExpandedBody then
        Result := N
    else
        Result := TAst.MacroExpansionNode(N.Identity, M.CallNode, exp);
end;

function TAstTransformer.VisitBlockExpression(const N: IAstNode): IAstNode;
var
    B: IBlockExpressionNode;
    e: ITupleNode;
begin
    B := N.AsBlockExpression;
    e := Visit(B.Expressions).AsTuple;
    if e = B.Expressions then
        Result := N
    else
        Result := TAst.Block(N.Identity, e, B.StaticType);
end;

function TAstTransformer.VisitVariableDeclaration(const N: IAstNode): IAstNode;
var
    V: IVariableDeclarationNode;
    t, i: IAstNode;
begin
    V := N.AsVariableDeclaration;
    t := Visit(V.Target);
    i := Visit(V.Initializer);
    if (t = V.Target) and (i = V.Initializer) then
        Result := N
    else
        Result := TAst.VarDecl(N.Identity, t, i, V.StaticType, V.IsBoxed);
end;

function TAstTransformer.VisitAssignment(const N: IAstNode): IAstNode;
var
    A: IAssignmentNode;
    t, v: IAstNode;
begin
    A := N.AsAssignment;
    t := Visit(A.Target);
    v := Visit(A.Value);
    if (t = A.Target) and (v = A.Value) then
        Result := N
    else
        Result := TAst.Assign(N.Identity, t, v, A.StaticType);
end;

function TAstTransformer.VisitMacroDefinition(const N: IAstNode): IAstNode;
begin
    // Macro definitions are typically constant during runtime transformation, return as is.
    Result := N;
end;

function TAstTransformer.VisitQuasiquote(const N: IAstNode): IAstNode;
var
    Q: IQuasiquoteNode;
    e: IAstNode;
begin
    Q := N.AsQuasiquote;
    e := Visit(Q.Expression);
    if e = Q.Expression then
        Result := N
    else
        Result := TAst.Quasiquote(N.Identity, e);
end;

function TAstTransformer.VisitUnquote(const N: IAstNode): IAstNode;
var
    U: IUnquoteNode;
    e: IAstNode;
begin
    U := N.AsUnquote;
    e := Visit(U.Expression);
    if e = U.Expression then
        Result := N
    else
        Result := TAst.Unquote(N.Identity, e);
end;

function TAstTransformer.VisitUnquoteSplicing(const N: IAstNode): IAstNode;
var
    U: IUnquoteSplicingNode;
    e: IAstNode;
begin
    U := N.AsUnquoteSplicing;
    e := Visit(U.Expression);
    if e = U.Expression then
        Result := N
    else
        Result := TAst.UnquoteSplicing(N.Identity, e.AsQuasiquote);
end;

function TAstTransformer.VisitIndexer(const N: IAstNode): IAstNode;
var
    I: IIndexerNode;
    b, idx: IAstNode;
begin
    I := N.AsIndexer;
    b := Visit(I.Base);
    idx := Visit(I.Index);
    if (b = I.Base) and (idx = I.Index) then
        Result := N
    else
        Result := TAst.Indexer(N.Identity, b, idx, I.StaticType);
end;

function TAstTransformer.VisitMemberAccess(const N: IAstNode): IAstNode;
var
    M: IMemberAccessNode;
    b: IAstNode;
begin
    M := N.AsMemberAccess;
    b := Visit(M.Base);
    if b = M.Base then
        Result := N
    else
        Result := TAst.MemberAccess(N.Identity, b, M.Member, M.StaticType);
end;

function TAstTransformer.VisitRecordLiteral(const N: IAstNode): IAstNode;
var
    R: IRecordLiteralNode;
    f: ITupleNode;
begin
    R := N.AsRecordLiteral;
    f := Visit(R.Fields).AsTuple;
    if f = R.Fields then
        Result := N
    else
        Result := TAst.RecordLiteral(N.Identity, f, R.ScalarDefinition, R.GenericDefinition, R.StaticType);
end;

function TAstTransformer.VisitAddSeriesItem(const N: IAstNode): IAstNode;
var
    A: IAddSeriesItemNode;
    s, v, l: IAstNode;
begin
    A := N.AsAddSeriesItem;
    s := Visit(A.Series);
    v := Visit(A.Value);
    l := Visit(A.Lookback);

    if (s = A.Series) and (v = A.Value) and (l = A.Lookback) then
        Result := N
    else
        Result := TAst.AddSeriesItem(N.Identity, s.AsIdentifier, v, l, A.StaticType);
end;

function TAstTransformer.VisitRecurNode(const N: IAstNode): IAstNode;
var
    R: IRecurNode;
    a: ITupleNode;
begin
    R := N.AsRecur;
    a := Visit(R.Arguments).AsTuple;
    if a = R.Arguments then
        Result := N
    else
        Result := TAst.Recur(N.Identity, a, R.StaticType);
end;

function TAstTransformer.VisitPipe(const N: IAstNode): IAstNode;
var
    P: IPipeNode;
    inp: ITupleNode;
    trans: ILambdaExpressionNode;
begin
    P := N.AsPipe;
    inp := Visit(P.Inputs).AsTuple;
    trans := Visit(P.Transformation).AsLambdaExpression;

    if (inp = P.Inputs) and (trans = P.Transformation) then
        Result := N
    else
        Result := TAst.Pipe(N.Identity, inp, trans, P.StaticType);
end;

function TAstTransformer.VisitTuple(const N: IAstNode): IAstNode;
var
    T: ITupleNode;
    hasChanged: Boolean;
    newElements: TArray<IAstNode>;
    item, newItem: IAstNode;
    i: Integer;
begin
    T := N.AsTuple;
    hasChanged := False;

    // Updated: Use Elements directly
    var oldElements := T.Elements;
    SetLength(newElements, Length(oldElements));

    for i := 0 to High(oldElements) do
    begin
        item := oldElements[i];
        newItem := Visit(item);
        newElements[i] := newItem;
        if item <> newItem then
            hasChanged := True;
    end;

    if not hasChanged then
        Result := N
    else
        Result := TAst.Tuple(N.Identity, newElements, T.AsTypedNode.StaticType);
end;

{ TAstVisitor (Walker) }

procedure TAstVisitor.SetupHandlers;
begin
    // Leaves (Do nothing)
    Register(akConstant, WalkLeaf);
    Register(akIdentifier, WalkLeaf);
    Register(akKeyword, WalkLeaf);
    Register(akNop, WalkLeaf);
    Register(akCreateSeries, WalkLeaf);

    // Elements
    Register(akRecordField, WalkRecordField);

    // Structures
    Register(akIfExpression, WalkIfExpression);
    Register(akCondExpression, WalkCondExpression);
    Register(akLambdaExpression, WalkLambdaExpression);
    Register(akFunctionCall, WalkFunctionCall);
    Register(akMacroExpansion, WalkMacroExpansionNode);
    Register(akBlockExpression, WalkBlockExpression);
    Register(akVariableDeclaration, WalkVariableDeclaration);
    Register(akAssignment, WalkAssignment);
    Register(akMacroDefinition, WalkMacroDefinition);

    // Wrappers (Generic Node Handler)
    Register(akQuasiquote, WalkWrapper);
    Register(akUnquote, WalkWrapper);
    Register(akUnquoteSplicing, WalkWrapper);

    Register(akIndexer, WalkIndexer);
    Register(akMemberAccess, WalkMemberAccess);
    Register(akRecordLiteral, WalkRecordLiteral);
    Register(akAddSeriesItem, WalkAddSeriesItem);
    Register(akRecur, WalkRecurNode);

    Register(akTuple, WalkTuple);

    // Pipes
    Register(akPipe, WalkPipe);
end;

function TAstVisitor.WalkLeaf(const N: IAstNode): TVoid;
begin
    // Do nothing
end;

function TAstVisitor.WalkWrapper(const N: IAstNode): TVoid;
begin
    if N.Kind = akQuasiquote then
        Visit(N.AsQuasiquote.Expression)
    else if N.Kind = akUnquote then
        Visit(N.AsUnquote.Expression)
    else if N.Kind = akUnquoteSplicing then
        Visit(N.AsUnquoteSplicing.Expression);
end;

function TAstVisitor.WalkRecordField(const N: IAstNode): TVoid;
begin
    var F := N.AsRecordField;
    Visit(F.Key);
    Visit(F.Value);
end;

function TAstVisitor.WalkIfExpression(const N: IAstNode): TVoid;
begin
    var E := N.AsIfExpression;
    Visit(E.Condition);
    Visit(E.ThenBranch);
    Visit(E.ElseBranch);
end;

function TAstVisitor.WalkCondExpression(const N: IAstNode): TVoid;
begin
    var E := N.AsCondExpression;
    for var pair in E.Pairs do
    begin
        Visit(pair.Condition);
        Visit(pair.Branch);
    end;
    Visit(E.ElseBranch);
end;

function TAstVisitor.WalkLambdaExpression(const N: IAstNode): TVoid;
begin
    var E := N.AsLambdaExpression;
    Visit(E.Parameters);
    Visit(E.Body);
end;

function TAstVisitor.WalkFunctionCall(const N: IAstNode): TVoid;
begin
    var C := N.AsFunctionCall;
    Visit(C.Callee);
    Visit(C.Arguments);
end;

function TAstVisitor.WalkMacroExpansionNode(const N: IAstNode): TVoid;
begin
    Visit(N.AsMacroExpansion.ExpandedBody);
end;

function TAstVisitor.WalkBlockExpression(const N: IAstNode): TVoid;
begin
    Visit(N.AsBlockExpression.Expressions);
end;

function TAstVisitor.WalkVariableDeclaration(const N: IAstNode): TVoid;
begin
    var V := N.AsVariableDeclaration;
    Visit(V.Target);
    Visit(V.Initializer);
end;

function TAstVisitor.WalkAssignment(const N: IAstNode): TVoid;
begin
    var A := N.AsAssignment;
    Visit(A.Target);
    Visit(A.Value);
end;

function TAstVisitor.WalkMacroDefinition(const N: IAstNode): TVoid;
begin
    var M := N.AsMacroDefinition;
    Visit(M.Name);
    Visit(M.Parameters);
    Visit(M.Body);
end;

function TAstVisitor.WalkIndexer(const N: IAstNode): TVoid;
begin
    var I := N.AsIndexer;
    Visit(I.Base);
    Visit(I.Index);
end;

function TAstVisitor.WalkMemberAccess(const N: IAstNode): TVoid;
begin
    var M := N.AsMemberAccess;
    Visit(M.Base);
    Visit(M.Member);
end;

function TAstVisitor.WalkRecordLiteral(const N: IAstNode): TVoid;
begin
    Visit(N.AsRecordLiteral.Fields);
end;

function TAstVisitor.WalkAddSeriesItem(const N: IAstNode): TVoid;
begin
    var A := N.AsAddSeriesItem;
    Visit(A.Series);
    Visit(A.Value);
    Visit(A.Lookback);
end;

function TAstVisitor.WalkRecurNode(const N: IAstNode): TVoid;
begin
    Visit(N.AsRecur.Arguments);
end;

function TAstVisitor.WalkPipe(const N: IAstNode): TVoid;
begin
    var P := N.AsPipe;
    Visit(P.Inputs);
    Visit(P.Transformation);
end;

function TAstVisitor.WalkTuple(const N: IAstNode): TVoid;
begin
    // Updated: Iterate directly via Elements
    var t := N.AsTuple;
    for var item in t.Elements do
        Visit(item);
end;

end.