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.Identities,
    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;

        FLambdaCounter: Integer;

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

    strict private
        // Specific Handlers (IAstNode signature)
        function VisitIdentifier(const Node: IAstNode): IAstNode;
        function VisitVariableDeclaration(const Node: IAstNode): IAstNode;
        function VisitAssignment(const Node: IAstNode): IAstNode;
        function VisitLambdaExpression(const Node: IAstNode): IAstNode;
        function VisitFunctionCall(const Node: IAstNode): IAstNode;

    protected
        procedure SetupHandlers; 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; // Calls SetupHandlers
    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;

procedure TAstBinder.SetupHandlers;
begin
    inherited SetupHandlers; // Loads default transformers (Identity, Lists, etc.)

    // Override specific binding logic
    Register(akIdentifier, VisitIdentifier);
    Register(akVariableDeclaration, VisitVariableDeclaration);
    Register(akAssignment, VisitAssignment);
    Register(akLambdaExpression, VisitLambdaExpression);
    Register(akFunctionCall, VisitFunctionCall);
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) as IAstBinder;
    Result := binder.Execute(RootNode, Layout);
end;

function TAstBinder.Execute(const RootNode: IAstNode; out Layout: IScopeLayout): IAstNode;
begin
    // Pre-pass to find captured variables
    FBoxedDeclarations := TUpvalueAnalyzer.Analyze(RootNode);

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

    Layout := FCurrentBuilder.Build;
end;

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

    for c in Name do
        if not (c.IsLetterOrDigit or (c = '_') or (c = '-') or (c = '#')) then
            exit(False);

    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
    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.VisitIdentifier(const Node: IAstNode): IAstNode;
var
    I: IIdentifierNode;
    physAddr: TResolvedAddress;
    upvalueIndex: Integer;
    identity: INamedIdentity;
begin
    I := Node.AsIdentifier;
    identity := I.Identity.AsNamed;

    if I.Address.Kind = akUnresolved then
    begin
        if not ResolveSymbol(I.Name, physAddr) then
        begin
            if Assigned(FLog) then
                FLog.AddError(Format('Undefined identifier: "%s"', [I.Name]), Node);
            Result := Node;
            Exit;
        end;

        if physAddr.ScopeDepth > 0 then
        begin
            upvalueIndex := CaptureUpvalue(physAddr);
            Result := TAst.Identifier(identity, TResolvedAddress.Create(akUpvalue, 0, upvalueIndex), TTypes.Unknown);
        end
        else
            Result := TAst.Identifier(identity, physAddr, TTypes.Unknown);
    end
    else
        Result := Node;
end;

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

    if not IsValidIdentifier(identifier.Name) then
    begin
        if Assigned(FLog) then
            FLog.AddError(Format('Invalid identifier name: "%s".', [identifier.Name]), Node);
    end;

    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);

        if Assigned(V.Initializer) then
            Accept(V.Initializer);

        Result := Node;
        Exit;
    end;

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

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

    newIdent := TAst.Identifier(identIdentity, addr, TTypes.Unknown);
    isBoxed := (FBoxedDeclarations <> nil) and FBoxedDeclarations.Contains(V);
    Result := TAst.VarDecl(Node.Identity, 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: IAstNode): IAstNode;
var
    A: IAssignmentNode;
    newIdent: IAstNode;
    newValue: IAstNode;
begin
    A := Node.AsAssignment;
    // Manually accept children because we need the results for registry logic
    newIdent := Accept(A.Target);
    newValue := Accept(A.Value);

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

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

function TAstBinder.VisitLambdaExpression(const Node: IAstNode): IAstNode;
var
    L: ILambdaExpressionNode;
    parentBuilder: IScopeBuilder;
    paramType: IStaticType;
    newParams: TArray<IIdentifierNode>;
    newParamsAsNodes: TArray<IAstNode>;
    newBody: IAstNode;
    finalLayout: IScopeLayout;
    i, slot: Integer;
    addr: TResolvedAddress;
    capturedMap: TUpvalueMap;
    sortedPairs: TArray<TPair<TResolvedAddress, Integer>>;
    upvaluesList: TArray<TResolvedAddress>;
    startCount: Integer;
    hasNested: Boolean;
    paramIdentity: INamedIdentity;
    paramNode: IAstNode;
begin
    L := Node.AsLambdaExpression;
    startCount := FLambdaCounter;
    Inc(FLambdaCounter);

    parentBuilder := FCurrentBuilder;
    FCurrentBuilder := TScope.CreateBuilder(parentBuilder);

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

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

        var paramElements := L.Parameters.Elements;
        SetLength(newParams, Length(paramElements));

        for i := 0 to High(paramElements) do
        begin
            paramNode := paramElements[i];

            if paramNode.Kind <> akIdentifier then
            begin
                if Assigned(FLog) then
                    FLog.AddError('Parameter must be an identifier.', paramNode);
                continue;
            end;

            var paramIdent := paramNode.AsIdentifier;
            var paramName := paramIdent.Name;
            paramIdentity := paramIdent.Identity.AsNamed;

            if FCurrentBuilder.FindSlot(paramName) >= 0 then
            begin
                if Assigned(FLog) then
                    FLog.AddError(Format('Duplicate parameter name "%s".', [paramName]), Node);
            end
            else
                FCurrentBuilder.Define(paramName);

            slot := FCurrentBuilder.FindSlot(paramName);
            if slot < 0 then
                slot := 0;

            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(paramIdentity, addr, paramType);
        end;

        newBody := Accept(L.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;

    hasNested := FLambdaCounter > (startCount + 1);

    // Convert TArray<IIdentifierNode> -> TArray<IAstNode> for Tuple factory
    SetLength(newParamsAsNodes, Length(newParams));
    for i := 0 to High(newParams) do
        newParamsAsNodes[i] := newParams[i];

    Result :=
        TAst.LambdaExpr(
            Node.Identity,
            TAst.Tuple(L.Parameters.Identity, newParamsAsNodes),
            newBody,
            finalLayout,
            nil,
            upvaluesList,
            hasNested,
            L.IsPure
        );
end;

function TAstBinder.VisitFunctionCall(const Node: IAstNode): IAstNode;
var
    C: IFunctionCallNode;
begin
    C := Node.AsFunctionCall;
    // 1. Keyword as Function Logic
    if C.Callee.Kind = akKeyword then
    begin
        var args := C.Arguments.Elements;
        if Length(args) <> 1 then
        begin
            if Assigned(FLog) then
                FLog.AddError(Format('Keyword access :%s requires exactly one argument.', [C.Callee.AsKeyword.Value.Name]), Node);
        end
        else
        begin
            var memberAccess := TAst.MemberAccess(Node.Identity, Accept(args[0]), C.Callee.AsKeyword);
            Result := memberAccess;
            exit;
        end;
    end;

    // 2. Delegate to standard AST transformation (recurse on callee and args)
    // This reuses logic from TAstTransformer!
    Result := inherited VisitFunctionCall(Node);
end;

end.