unit Myc.Ast.Compiler.Binder;

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.Compiler.Binder.Upvalues,
    Myc.Ast,
    Myc.Ast.Environment;

type
    IAstBinder = interface(IAstVisitor)
        // AArgTypes is now passed to the implementation via its constructor
        function Execute(const RootNode: IAstNode; out Descriptor: IScopeDescriptor): IAstNode;
    end;

    TAstBinder = class; // Forward declaration

    TAstBinder = class(TAstTransformer, IAstBinder)
    private
        type
            TUpvalueMapping = TDictionary<TResolvedAddress, Integer>;
    private
        FCurrentDescriptor: IScopeDescriptor;
        FUpvalueStack: TStack<TUpvalueMapping>;
        FNestedLambdaCount: Integer;
        FBoxedDeclarations: THashSet<IVariableDeclarationNode>;
        FArgTypes: TArray<IStaticType>;
        FFunctionRegistry: IFunctionDefinitionRegistry;

        procedure EnterScope;
        procedure ExitScope;
        function IsValidIdentifier(const Name: string): Boolean;

    protected
        // --- Core Binding Logic (Mutators) ---
        function VisitIdentifier(const Node: IIdentifierNode): IAstNode; override;
        function VisitVariableDeclaration(const Node: IVariableDeclarationNode): IAstNode; override;
        function VisitAssignment(const Node: IAssignmentNode): IAstNode; override;
        function VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode; override;

        // --- Transformation Logic (Restored) ---
        function VisitFunctionCall(const Node: IFunctionCallNode): IAstNode; override;

        // --- Standard Traversal (Use inherited) ---
        function VisitKeyword(const Node: IKeywordNode): IAstNode; override;
        function VisitConstant(const Node: IConstantNode): IAstNode; override;
        function VisitCreateSeries(const Node: ICreateSeriesNode): IAstNode; override;
    public
        constructor Create(
            const AInitialDescriptor: IScopeDescriptor;
            const AFunctionRegistry: IFunctionDefinitionRegistry;
            const AArgTypes: TArray<IStaticType>
        );
        destructor Destroy; override;
        function Execute(const RootNode: IAstNode; out Descriptor: IScopeDescriptor): IAstNode;

        class function Bind(
            const InitialDescriptor: IScopeDescriptor;
            const RootNode: IAstNode;
            out Descriptor: IScopeDescriptor;
            const AFunctionRegistry: IFunctionDefinitionRegistry = nil;
            const AArgTypes: TArray<IStaticType> = nil
        ): IAstNode; static;
    end;

implementation

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

type
    TResolvedAddressComparer = class(TEqualityComparer<TResolvedAddress>)
    public
        function Equals(const Left, Right: TResolvedAddress): Boolean; override;
        function GetHashCode(const Value: TResolvedAddress): Integer; override;
    end;

{ TResolvedAddressComparer }
function TResolvedAddressComparer.Equals(const Left, Right: TResolvedAddress): Boolean;
begin
    Result := (Left = Right);
end;

function TResolvedAddressComparer.GetHashCode(const Value: TResolvedAddress): Integer;
begin
    Result := 17;
    Result := Result * 23 + Ord(Value.Kind);
    Result := Result * 23 + Value.ScopeDepth;
    Result := Result * 23 + Value.SlotIndex;
end;

{ TAstBinder }

constructor TAstBinder.Create(
    const AInitialDescriptor: IScopeDescriptor;
    const AFunctionRegistry: IFunctionDefinitionRegistry;
    const AArgTypes: TArray<IStaticType>
);
begin
    inherited Create;
    Assert(Assigned(AInitialDescriptor));

    FCurrentDescriptor := AInitialDescriptor;
    FFunctionRegistry := AFunctionRegistry;
    FArgTypes := AArgTypes;
    FUpvalueStack := TObjectStack<TUpvalueMapping>.Create(True); // Use Comparer
    FNestedLambdaCount := 0;
    FBoxedDeclarations := nil;
end;

destructor TAstBinder.Destroy;
begin
    FUpvalueStack.Free;
    FBoxedDeclarations.Free;
    inherited;
end;

class function TAstBinder.Bind(
    const InitialDescriptor: IScopeDescriptor;
    const RootNode: IAstNode;
    out Descriptor: IScopeDescriptor;
    const AFunctionRegistry: IFunctionDefinitionRegistry = nil;
    const AArgTypes: TArray<IStaticType> = nil
): IAstNode;
begin
    var binder := TAstBinder.Create(InitialDescriptor, AFunctionRegistry, AArgTypes) as IAstBinder;
    Result := binder.Execute(RootNode, Descriptor);
end;

procedure TAstBinder.EnterScope;
begin
    FCurrentDescriptor := TScope.CreateDescriptor(FCurrentDescriptor);
end;

procedure TAstBinder.ExitScope;
begin
    FCurrentDescriptor := FCurrentDescriptor.Parent;
end;

function TAstBinder.IsValidIdentifier(const Name: string): Boolean;
var
    c: Char;
begin
    if Name.IsEmpty then
        exit(False);
    c := Name[1];
    if not (c.IsLetter or (c = '_')) then
        exit(False);
    for c in Name do
    begin
        if not (c.IsLetterOrDigit or (c = '_') or (c = '-')) then
            exit(False);
    end;
    Result := True;
end;

function TAstBinder.Execute(const RootNode: IAstNode; out Descriptor: IScopeDescriptor): IAstNode;
begin
    // Pre-pass: Find all variables that need boxing
    FBoxedDeclarations := TUpvalueAnalyzer.Analyze(RootNode, FCurrentDescriptor.Parent);
    try
        EnterScope;
        try
            Result := Accept(RootNode); // Accept returns IAstNode
            if not Assigned(Result) then
                Result := TAst.Block([]);

            // Set the type of the root expression (e.g., the final 'do' block)
            Descriptor := FCurrentDescriptor;
        finally
            ExitScope;
        end;
    finally
        FBoxedDeclarations.Free; // Free the set
        FBoxedDeclarations := nil;
    end;
end;

function TAstBinder.VisitFunctionCall(const Node: IFunctionCallNode): IAstNode;
begin
    // --- Transformation: Keyword-as-Function ---
    // This transformation must happen *before* type checking.
    if Node.Callee.Kind = akKeyword then
    begin
        var keywordNode := Node.Callee.AsKeyword;
        if Length(Node.Arguments) <> 1 then
            raise EArgumentException.CreateFmt(
                'Keyword :%s expects exactly one argument (the record/map), but got %d',
                [keywordNode.Value.Name, Length(Node.Arguments)]);

        // Manually visit the argument (which is the base)
        var baseNode := Accept(Node.Arguments[0]);
        // Manually visit the keyword (which is the member)
        var memberNode := Accept(keywordNode).AsKeyword;

        // Create the new MemberAccess node
        var memberAccessNode := TAst.MemberAccess(baseNode, memberNode);

        // Visit the *new* node to bind it (though it has no bindable symbols)
        Result := Accept(memberAccessNode);
        exit;
    end;

    // --- Default: Bind as a standard function call ---
    // Use the inherited implementation to visit children (Callee, Arguments)
    Result := inherited VisitFunctionCall(Node);

    // Set metadata for *this* node (IsTailCall is set later by TCO)
end;

function TAstBinder.VisitConstant(const Node: IConstantNode): IAstNode;
begin
    // Binding pass does not assign types.
    Result := Node;
end;

function TAstBinder.VisitKeyword(const Node: IKeywordNode): IAstNode;
begin
    // Binding pass does not assign types.
    Result := Node;
end;

function TAstBinder.VisitCreateSeries(const Node: ICreateSeriesNode): IAstNode;
begin
    // Binding pass does not assign types.
    Result := Node;
end;

function TAstBinder.VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode;
var
    i: integer;
    adr: TResolvedAddress;
    scopeDescriptor: IScopeDescriptor;
    upvalues: TArray<TResolvedAddress>;
    hasNestedLambdas: Boolean;
    lastNestedLambdaCount: Integer;
    newParams: TArray<IIdentifierNode>;
    newBody: IAstNode;
    paramType: IStaticType;
begin
    // We do *not* call inherited, as we must manage the scope manually.

    FUpvalueStack.Push(TUpvalueMapping.Create(TResolvedAddressComparer.Create));
    try
        EnterScope;
        try
            // Define <self> (slot 0)
            FCurrentDescriptor.Define('<self>', TTypes.Unknown);

            // Define parameters
            SetLength(newParams, Length(Node.Parameters));
            for i := 0 to High(Node.Parameters) do
            begin
                var paramNode := Node.Parameters[i]; // This is a TIdentifierNode

                // Check if this is the root lambda (stack count = 1)
                // and we have pre-defined types for it (from specialization)
                paramType := TTypes.Unknown;
                if (FUpvalueStack.Count = 1) and (FArgTypes <> nil) and (i < Length(FArgTypes)) then
                    paramType := FArgTypes[i]; // Use the specialized type

                // Define the symbol in the descriptor WITH the type
                var slotIndex := FCurrentDescriptor.Define(paramNode.Name, paramType);
                adr := TResolvedAddress.Create(akLocalOrParent, 0, slotIndex);

                // Create a new TIdentifierNode, now with the correct type
                newParams[i] := TAst.Identifier(paramNode.Name, adr, paramType);
            end;

            // Visit the body *within the new scope*
            lastNestedLambdaCount := FNestedLambdaCount;
            newBody := Accept(Node.Body);
            hasNestedLambdas := FNestedLambdaCount > lastNestedLambdaCount;

            // Save the descriptor for the evaluator
            scopeDescriptor := FCurrentDescriptor;
        finally
            ExitScope;
        end;

        // --- Extract Upvalues ---
        var upvalueMapping := FUpvalueStack.Peek;
        var sortedPairs := upvalueMapping.ToArray;
        // Sort by index (Value)
        TArray.Sort<TPair<TResolvedAddress, Integer>>(
            sortedPairs,
            TComparer<TPair<TResolvedAddress, Integer>>.Construct(
                function(const Left, Right: TPair<TResolvedAddress, Integer>): Integer begin Result := Left.Value - Right.Value; end
            )
        );

        SetLength(upvalues, Length(sortedPairs));
        for i := 0 to High(upvalues) do
            upvalues[i] := sortedPairs[i].Key;

    finally
        FUpvalueStack.Pop;
    end;

    inc(FNestedLambdaCount);

    // Create new immutable node using the factory
    Result := TAst.LambdaExpr(newParams, newBody, scopeDescriptor, upvalues, hasNestedLambdas);
end;

function TAstBinder.VisitIdentifier(const Node: IIdentifierNode): IAstNode;
var
    symbol: TResolvedSymbol;
    adr: TResolvedAddress;
begin
    // We only bind nodes that are the base parser type
    // (i.e., address is unresolved).
    if Node.Address.Kind = akUnresolved then
    begin
        symbol := FCurrentDescriptor.FindSymbol(Node.Name);
        adr := symbol.Address;

        if adr.Kind = akUnresolved then
            raise Exception.CreateFmt('Undefined identifier: "%s"', [Node.Name]);

        if adr.Kind = akLocalOrParent then
        begin
            if (adr.ScopeDepth > 0) and (FUpvalueStack.Count > 0) then
            begin
                // --- Handle Upvalue ---
                var upvalueMap := FUpvalueStack.Peek;
                // Adjust address to be relative to the captured scope
                dec(adr.ScopeDepth);

                var upvalueIndex: Integer;
                if not upvalueMap.TryGetValue(adr, upvalueIndex) then
                begin
                    // This is a new upvalue for this lambda
                    upvalueIndex := upvalueMap.Count;
                    upvalueMap.Add(adr, upvalueIndex);
                end;

                // Create final upvalue address
                adr := TResolvedAddress.Create(akUpvalue, 0, upvalueIndex);
            end;
            // else: It's a local (ScopeDepth=0), adr is already correct.
        end;
        // else: It was already an upvalue (e.g. nested lambda), adr is correct.

        // --- Replace Node ---
        Result := TAst.Identifier(Node.Name, adr, symbol.StaticType);
    end
    else
    begin
        // It's already bound (e.g., a parameter), just return it.
        Result := Node;
    end;
end;

function TAstBinder.VisitVariableDeclaration(const Node: IVariableDeclarationNode): IAstNode;
var
    slotIndex: Integer;
    address: TResolvedAddress;
    isBoxed: Boolean;
    newInitializer: IAstNode;
    newIdentifier: IIdentifierNode;
begin
    // 1. Validate name
    if not IsValidIdentifier(Node.Identifier.Name) then
        raise Exception.CreateFmt('Invalid identifier name: "%s".', [Node.Identifier.Name]);

    // 2. Define the symbol in the scope *first*.
    // This allows the initializer (e.g., a lambda) to recursively refer to this symbol.
    slotIndex := FCurrentDescriptor.Define(Node.Identifier.Name, TTypes.Unknown);
    address := TResolvedAddress.Create(akLocalOrParent, 0, slotIndex);

    // 3. Visit the Initializer *after* the symbol is defined
    if Assigned(Node.Initializer) then
        newInitializer := Accept(Node.Initializer)
    else
        newInitializer := nil;

    // 4. Create the new (bound) identifier node
    // The TypeChecker will set the final type
    newIdentifier := TAst.Identifier(Node.Identifier.Name, address, TTypes.Unknown);

    // 5. Check if this declaration needs boxing (from UpvalueAnalyzer pass)
    isBoxed := (FBoxedDeclarations <> nil) and FBoxedDeclarations.Contains(Node);

    // 6. Create the new VariableDeclaration node using the factory
    Result := TAst.VarDecl(newIdentifier, newInitializer, TTypes.Unknown, isBoxed);

    // 7. Store in registry *only if the registry was provided*
    // (This must happen *after* visiting the initializer)
    if Assigned(FFunctionRegistry) and (newInitializer <> nil) and (newInitializer.Kind = akLambdaExpression) then
        FFunctionRegistry.Register(address, newInitializer.AsLambdaExpression);
end;

function TAstBinder.VisitAssignment(const Node: IAssignmentNode): IAstNode;
var
    newIdent: IIdentifierNode;
    newValue: IAstNode;
begin
    Result := inherited VisitAssignment(Node);

    newValue := Result.AsAssignment.Value;
    newIdent := Result.AsAssignment.Identifier;

    if Assigned(FFunctionRegistry) and (newValue <> nil) and (newValue.Kind = akLambdaExpression) then
    begin
        FFunctionRegistry.Register(newIdent.Address, newValue.AsLambdaExpression);
    end;
end;

end.