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;

type
    IAstBinder = interface(IAstVisitor)
        function Execute(const RootNode: IAstNode; out Layout: IScopeLayout): IAstNode;
    end;

    IFunctionDefinitionRegistry = interface
        procedure Register(const Address: TResolvedAddress; const ADef: IFunctionDefinition);
        function Resolve(const Address: TResolvedAddress): IFunctionDefinition;
    end;

    TAstBinder = class(TAstTransformer, IAstBinder)
    private
        type
            TUpvalueMap = TDictionary<TResolvedAddress, Integer>;

    private
        FCurrentBuilder: IScopeBuilder;
        FUpvalueStack: TObjectStack<TUpvalueMap>;

        FArgTypes: TArray<IStaticType>;
        FFunctionRegistry: IFunctionDefinitionRegistry;
        FBoxedDeclarations: THashSet<IVariableDeclarationNode>;
        FLog: ICompilerLog;

        // Counts lambdas encountered to determine if a lambda has nested children
        FLambdaCounter: Integer;

        function IsValidIdentifier(const Name: string): Boolean;
        function ResolveSymbol(const Name: string; out Address: TResolvedAddress): Boolean;
        function CaptureUpvalue(const PhysicalAddress: TResolvedAddress): Integer;

    protected
        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;
        function VisitFunctionCall(const Node: IFunctionCallNode): IAstNode; override;

        function VisitConstant(const Node: IConstantNode): IAstNode; override;
        function VisitKeyword(const Node: IKeywordNode): IAstNode; override;
        function VisitCreateSeries(const Node: ICreateSeriesNode): IAstNode; override;

    public
        constructor Create(
            const AParentLayout: IScopeLayout;
            const AFunctionRegistry: IFunctionDefinitionRegistry;
            const AArgTypes: TArray<IStaticType>;
            const ALog: ICompilerLog
        );
        destructor Destroy; override;

        function Execute(const RootNode: IAstNode; out Layout: IScopeLayout): IAstNode;

        class function Bind(
            const ParentLayout: IScopeLayout;
            const RootNode: IAstNode;
            out Layout: IScopeLayout;
            const ALog: ICompilerLog;
            const AFunctionRegistry: IFunctionDefinitionRegistry = nil;
            const AArgTypes: TArray<IStaticType> = nil
        ): IAstNode; static;
    end;

implementation

uses
    System.Generics.Defaults,
    System.Character,
    System.Hash,
    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;

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

function TResolvedAddressComparer.GetHashCode(const Value: TResolvedAddress): Integer;
begin
    Result := THashBobJenkins.GetHashValue(Value.Kind, SizeOf(TAddressKind), 0);
    Result := THashBobJenkins.GetHashValue(Value.ScopeDepth, SizeOf(Integer), Result);
    Result := THashBobJenkins.GetHashValue(Value.SlotIndex, SizeOf(Integer), Result);
end;

{ TAstBinder }

constructor TAstBinder.Create(
    const AParentLayout: IScopeLayout;
    const AFunctionRegistry: IFunctionDefinitionRegistry;
    const AArgTypes: TArray<IStaticType>;
    const ALog: ICompilerLog
);
begin
    inherited Create;
    FCurrentBuilder := TScope.CreateBuilder(AParentLayout);
    FUpvalueStack := TObjectStack<TUpvalueMap>.Create(True);
    FUpvalueStack.Push(TUpvalueMap.Create(TResolvedAddressComparer.Create));

    FFunctionRegistry := AFunctionRegistry;
    FArgTypes := AArgTypes;
    FLog := ALog;
    FBoxedDeclarations := nil;
    FLambdaCounter := 0;
end;

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

class function TAstBinder.Bind(
    const ParentLayout: IScopeLayout;
    const RootNode: IAstNode;
    out Layout: IScopeLayout;
    const ALog: ICompilerLog;
    const AFunctionRegistry: IFunctionDefinitionRegistry = nil;
    const AArgTypes: TArray<IStaticType> = nil
): IAstNode;
begin
    var binder := TAstBinder.Create(ParentLayout, AFunctionRegistry, AArgTypes, ALog);
    try
        Result := binder.Execute(RootNode, Layout);
    finally
        binder.Free;
    end;
end;

function TAstBinder.Execute(const RootNode: IAstNode; out Layout: IScopeLayout): IAstNode;
begin
    FBoxedDeclarations := TUpvalueAnalyzer.Analyze(RootNode);

    Result := Accept(RootNode);
    if not Assigned(Result) then
        Result := TAst.Block([]);

    Layout := FCurrentBuilder.Build;
end;

function TAstBinder.IsValidIdentifier(const Name: string): Boolean;
var
    c: Char;
begin
    if Name.IsEmpty then
        exit(False);

    // Support standard identifiers and hygienic macro names (e.g. loop#1)
    for c in Name do
        if not (c.IsLetterOrDigit or (c = '_') or (c = '-') or (c = '#')) then
            exit(False);

    // First char check (cannot be digit unless it's a pure number which is tokenized differently)
    c := Name[1];
    if not (c.IsLetter or (c = '_') or (c = '#')) then
        exit(False);

    Result := True;
end;

function TAstBinder.ResolveSymbol(const Name: string; out Address: TResolvedAddress): Boolean;
var
    layout: IScopeLayout;
    depth: Integer;
    slot: Integer;
begin
    // Start search in current builder
    layout := FCurrentBuilder;
    depth := 0;

    while Assigned(layout) do
    begin
        slot := layout.FindSlot(Name);
        if slot >= 0 then
        begin
            Address := TResolvedAddress.Create(akLocalOrParent, depth, slot);
            Result := True;
            exit;
        end;
        layout := layout.Parent;
        Inc(depth);
    end;

    Result := False;
end;

function TAstBinder.CaptureUpvalue(const PhysicalAddress: TResolvedAddress): Integer;
var
    currentMap: TUpvalueMap;
begin
    currentMap := FUpvalueStack.Peek;
    if not currentMap.TryGetValue(PhysicalAddress, Result) then
    begin
        Result := currentMap.Count;
        currentMap.Add(PhysicalAddress, Result);
    end;
end;

function TAstBinder.VisitConstant(const Node: IConstantNode): IAstNode;
begin
    Result := Node;
end;

function TAstBinder.VisitKeyword(const Node: IKeywordNode): IAstNode;
begin
    Result := Node;
end;

function TAstBinder.VisitCreateSeries(const Node: ICreateSeriesNode): IAstNode;
begin
    Result := Node;
end;

function TAstBinder.VisitIdentifier(const Node: IIdentifierNode): IAstNode;
var
    physAddr: TResolvedAddress;
    upvalueIndex: Integer;
begin
    if Node.Address.Kind = akUnresolved then
    begin
        if not ResolveSymbol(Node.Name, physAddr) then
        begin
            // Log error but don't crash. Return unresolved node (Poison Pill).
            if Assigned(FLog) then
                FLog.AddError(Format('Undefined identifier: "%s"', [Node.Name]), Node);
            Result := Node;
            Exit;
        end;

        if physAddr.ScopeDepth > 0 then
        begin
            // It's an upvalue (from a parent scope). Capture it.
            upvalueIndex := CaptureUpvalue(physAddr);
            Result := TAst.Identifier(Node.Name, TResolvedAddress.Create(akUpvalue, 0, upvalueIndex), TTypes.Unknown);
        end
        else
            // It's local.
            Result := TAst.Identifier(Node.Name, physAddr, TTypes.Unknown);
    end
    else
        Result := Node;
end;

function TAstBinder.VisitVariableDeclaration(const Node: IVariableDeclarationNode): IAstNode;
var
    slot: Integer;
    addr: TResolvedAddress;
    newInit: IAstNode;
    newIdent: IIdentifierNode;
    isBoxed: Boolean;
begin
    var identifier := Node.Target.AsIdentifier;

    if not IsValidIdentifier(identifier.Name) then
    begin
        if Assigned(FLog) then
            FLog.AddError(Format('Invalid identifier name: "%s".', [identifier.Name]), Node);
        // Continue to try processing the initializer
    end;

    // Check for duplicates before defining
    if FCurrentBuilder.FindSlot(identifier.Name) >= 0 then
    begin
        if Assigned(FLog) then
            FLog.AddError(Format('Variable "%s" is already defined in this scope.', [identifier.Name]), Node);

        // Do NOT define the slot to avoid scope corruption.
        // Visit initializer to catch errors there, then return original node or dummy.
        if Assigned(Node.Initializer) then
            Accept(Node.Initializer);

        Result := Node;
        Exit;
    end;

    slot := FCurrentBuilder.Define(identifier.Name);
    addr := TResolvedAddress.Create(akLocalOrParent, 0, slot);

    if Assigned(Node.Initializer) then
        newInit := Accept(Node.Initializer)
    else
        newInit := nil;

    newIdent := TAst.Identifier(identifier.Name, addr, TTypes.Unknown);
    isBoxed := (FBoxedDeclarations <> nil) and FBoxedDeclarations.Contains(Node);

    Result := TAst.VarDecl(newIdent, newInit, TTypes.Unknown, isBoxed);

    if Assigned(FFunctionRegistry) and (newInit <> nil) and (newInit.Kind = akLambdaExpression) then
        FFunctionRegistry.Register(addr, newInit.AsLambdaExpression);
end;

function TAstBinder.VisitAssignment(const Node: IAssignmentNode): IAstNode;
var
    newIdent: IAstNode;
    newValue: IAstNode;
begin
    newIdent := Accept(Node.Target);
    newValue := Accept(Node.Value);

    Result := TAst.Assign(newIdent.AsIdentifier, newValue, TTypes.Unknown);

    if Assigned(FFunctionRegistry) and (newValue <> nil) and (newValue.Kind = akLambdaExpression) then
    begin
        // Only register if it resolved correctly
        if newIdent.AsIdentifier.Address.Kind = akLocalOrParent then
            FFunctionRegistry.Register(newIdent.AsIdentifier.Address, newValue.AsLambdaExpression);
    end;
end;

function TAstBinder.VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode;
var
    parentBuilder: IScopeBuilder;
    paramType: IStaticType;
    newParams: TArray<IIdentifierNode>;
    newBody: IAstNode;
    finalLayout: IScopeLayout;
    i, slot: Integer;
    addr: TResolvedAddress;
    capturedMap: TUpvalueMap;
    sortedPairs: TArray<TPair<TResolvedAddress, Integer>>;
    upvaluesList: TArray<TResolvedAddress>;

    startCount: Integer;
    hasNested: Boolean;
begin
    // 1. Count logic to determine if this lambda has nested lambdas
    startCount := FLambdaCounter;
    Inc(FLambdaCounter); // Count myself

    // 2. Push Scope
    parentBuilder := FCurrentBuilder;
    FCurrentBuilder := TScope.CreateBuilder(parentBuilder);

    capturedMap := TUpvalueMap.Create(TResolvedAddressComparer.Create);
    FUpvalueStack.Push(capturedMap);

    try
        FCurrentBuilder.Define('<self>');

        SetLength(newParams, Length(Node.Parameters));
        for i := 0 to High(Node.Parameters) do
        begin
            var paramName := Node.Parameters[i].Name;

            // Check duplicate parameters
            if FCurrentBuilder.FindSlot(paramName) >= 0 then
            begin
                if Assigned(FLog) then
                    FLog.AddError(Format('Duplicate parameter name "%s".', [paramName]), Node);
                // Skip defining this parameter to avoid scope crash, but create a dummy binding
                // so index alignment isn't completely broken if we were to continue harder.
                // Here we just skip definition.
            end
            else
                FCurrentBuilder.Define(paramName);

            // Note: We still create a valid (or attempted) identifier node mapping
            // even if definition failed, to keep AST shape consistent.
            slot := FCurrentBuilder.FindSlot(paramName); // might return <0 or previous slot if duplicate
            if slot < 0 then
                slot := 0; // Fallback for duplicates

            addr := TResolvedAddress.Create(akLocalOrParent, 0, slot);

            paramType := TTypes.Unknown;
            if (parentBuilder.Parent = nil) and (FArgTypes <> nil) and (i < Length(FArgTypes)) then
                paramType := FArgTypes[i];

            newParams[i] := TAst.Identifier(paramName, addr, paramType);
        end;

        newBody := Accept(Node.Body);
        finalLayout := FCurrentBuilder.Build;

        sortedPairs := capturedMap.ToArray;
        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(upvaluesList, Length(sortedPairs));
        for i := 0 to High(sortedPairs) do
        begin
            var rawAddr := sortedPairs[i].Key;
            if rawAddr.Kind = akLocalOrParent then
            begin
                Assert(rawAddr.ScopeDepth > 0, 'Logic Error: Trying to capture a local variable.');
                Dec(rawAddr.ScopeDepth);
            end;
            upvaluesList[i] := rawAddr;
        end;

    finally
        FCurrentBuilder := parentBuilder;
        FUpvalueStack.Pop;
    end;

    // 3. Calculate hasNested
    // If FLambdaCounter increased by more than just 1 (myself), children were visited.
    hasNested := FLambdaCounter > (startCount + 1);

    Result := TAst.LambdaExpr(newParams, newBody, finalLayout, nil, upvaluesList, hasNested, Node.IsPure);
end;

function TAstBinder.VisitFunctionCall(const Node: IFunctionCallNode): IAstNode;
begin
    if Node.Callee.Kind = akKeyword then
    begin
        var memberAccess := TAst.MemberAccess(Accept(Node.Arguments[0]), Node.Callee.AsKeyword);
        Result := memberAccess;
        exit;
    end;
    Result := inherited VisitFunctionCall(Node);
end;

end.