unit Myc.Ast.Compiler.TypeChecker;

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

type
    IAstTypeChecker = interface(IAstVisitor)
        function Execute(const RootNode: IAstNode): IAstNode;
    end;

    TTypeChecker = class(TAstTransformer, IAstTypeChecker)
    private
        type
            TTypeContext = class
            private
                FParent: TTypeContext;
                FLayout: IScopeLayout;
                FSlotTypes: TArray<IStaticType>;
                FUpvalueTypes: TArray<IStaticType>;
            public
                constructor Create(
                    AParent: TTypeContext;
                    ALayout: IScopeLayout;
                    const AUpvalueTypes: TArray<IStaticType>;
                    ADescriptor: IScopeDescriptor
                );
                function LookupType(const Address: TResolvedAddress): IStaticType;
                procedure SetType(SlotIndex: Integer; AType: IStaticType);
                property Types: TArray<IStaticType> read FSlotTypes;
            end;

    private
        FCurrentContext: TTypeContext;
        FLog: ICompilerLog;
        FRootScope: IExecutionScope;

        function CreateContextChain(L: IScopeLayout): TTypeContext;

        // Helpers for Null Propagation
        function PrepareBaseType(const BaseNode: IAstNode; out IsOptional: Boolean): IStaticType;
        function ApplyOptionality(const AType: IStaticType; IsOptional: Boolean): IStaticType;

    strict private
        // Typed Handlers (non-virtual, 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;
        function VisitBlockExpression(const Node: IAstNode): IAstNode;
        function VisitIfExpression(const Node: IAstNode): IAstNode;
        function VisitMemberAccess(const Node: IAstNode): IAstNode;
        function VisitIndexer(const Node: IAstNode): IAstNode;
        function VisitCreateSeries(const Node: IAstNode): IAstNode;
        function VisitSeriesLength(const Node: IAstNode): IAstNode;
        function VisitRecurNode(const Node: IAstNode): IAstNode;
        function VisitNop(const Node: IAstNode): IAstNode;
        function VisitRecordLiteral(const Node: IAstNode): IAstNode;
        function VisitConstant(const Node: IAstNode): IAstNode;
        function VisitKeyword(const Node: IAstNode): IAstNode;
        function VisitTuple(const Node: IAstNode): IAstNode;

        // Pipe Support
        function VisitPipe(const Node: IAstNode): IAstNode;

    protected
        procedure SetupHandlers; override;

    public
        constructor Create(const RootLayout: IScopeLayout; const RootScope: IExecutionScope; const ALog: ICompilerLog);
        destructor Destroy; override;

        function Execute(const RootNode: IAstNode): IAstNode;

        class function CheckTypes(
            const RootNode: IAstNode;
            const Layout: IScopeLayout;
            const RootScope: IExecutionScope;
            const ALog: ICompilerLog
        ): IAstNode; static;
    end;

implementation

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

{ TTypeChecker.TTypeContext }

constructor TTypeChecker.TTypeContext.Create(
    AParent: TTypeContext;
    ALayout: IScopeLayout;
    const AUpvalueTypes: TArray<IStaticType>;
    ADescriptor: IScopeDescriptor
);
var
    i: Integer;
begin
    inherited Create;
    FParent := AParent;
    FLayout := ALayout;
    FUpvalueTypes := AUpvalueTypes;

    if Assigned(FLayout) then
    begin
        SetLength(FSlotTypes, FLayout.SlotCount);

        if Assigned(ADescriptor) then
        begin
            // Load known types from descriptor (e.g. for Root Scope / RTL)
            for i := 0 to High(FSlotTypes) do
                FSlotTypes[i] := ADescriptor.GetSymbolType(i);
        end
        else
        begin
            // Initialize with Unknown for new scopes
            for i := 0 to High(FSlotTypes) do
                FSlotTypes[i] := TTypes.Unknown;
        end;
    end;
end;

function TTypeChecker.TTypeContext.LookupType(const Address: TResolvedAddress): IStaticType;
var
    ctx: TTypeContext;
    i: Integer;
begin
    case Address.Kind of
        akLocalOrParent:
        begin
            ctx := Self;
            for i := 1 to Address.ScopeDepth do
            begin
                if not Assigned(ctx.FParent) then
                begin
                    Result := TTypes.Unknown;
                    Exit;
                end;
                ctx := ctx.FParent;
            end;

            if (Address.SlotIndex >= 0) and (Address.SlotIndex < Length(ctx.FSlotTypes)) then
                Result := ctx.FSlotTypes[Address.SlotIndex]
            else
                Result := TTypes.Unknown;
        end;
        akUpvalue:
        begin
            if (Address.SlotIndex >= 0) and (Address.SlotIndex < Length(FUpvalueTypes)) then
                Result := FUpvalueTypes[Address.SlotIndex]
            else
                Result := TTypes.Unknown;
        end;
    else
        Result := TTypes.Unknown;
    end;
end;

procedure TTypeChecker.TTypeContext.SetType(SlotIndex: Integer; AType: IStaticType);
begin
    if (SlotIndex >= 0) and (SlotIndex < Length(FSlotTypes)) then
        FSlotTypes[SlotIndex] := AType;
end;

{ TTypeChecker }

function TTypeChecker.CreateContextChain(L: IScopeLayout): TTypeContext;
var
    p: TTypeContext;
    desc: IScopeDescriptor;
begin
    if L = nil then
        exit(nil);

    p := CreateContextChain(L.Parent);

    desc := nil;
    if (L.Parent = nil) and Assigned(FRootScope) then
    begin
        desc := FRootScope.Descriptor;
    end;

    Result := TTypeContext.Create(p, L, [], desc);
end;

constructor TTypeChecker.Create(const RootLayout: IScopeLayout; const RootScope: IExecutionScope; const ALog: ICompilerLog);
begin
    inherited Create;
    FLog := ALog;
    FRootScope := RootScope;

    FCurrentContext := CreateContextChain(RootLayout);
    if FCurrentContext = nil then
        FCurrentContext := TTypeContext.Create(nil, nil, [], nil);
end;

destructor TTypeChecker.Destroy;
begin
    while Assigned(FCurrentContext) do
    begin
        var temp := FCurrentContext;
        FCurrentContext := FCurrentContext.FParent;
        temp.Free;
    end;
    inherited;
end;

procedure TTypeChecker.SetupHandlers;
begin
    inherited SetupHandlers; // Load Defaults

    Register(akIdentifier, VisitIdentifier);
    Register(akVariableDeclaration, VisitVariableDeclaration);
    Register(akAssignment, VisitAssignment);
    Register(akLambdaExpression, VisitLambdaExpression);
    Register(akFunctionCall, VisitFunctionCall);
    Register(akBlockExpression, VisitBlockExpression);
    Register(akIfExpression, VisitIfExpression);
    Register(akMemberAccess, VisitMemberAccess);
    Register(akIndexer, VisitIndexer);
    Register(akCreateSeries, VisitCreateSeries);
    Register(akSeriesLength, VisitSeriesLength);
    Register(akRecur, VisitRecurNode);
    Register(akNop, VisitNop);
    Register(akRecordLiteral, VisitRecordLiteral);
    Register(akConstant, VisitConstant);
    Register(akKeyword, VisitKeyword);
    Register(akTuple, VisitTuple);

    // Pipe Support
    Register(akPipe, VisitPipe);
end;

class function TTypeChecker.CheckTypes(
    const RootNode: IAstNode;
    const Layout: IScopeLayout;
    const RootScope: IExecutionScope;
    const ALog: ICompilerLog
): IAstNode;
var
    startLayout: IScopeLayout;
begin
    if Assigned(Layout) then
        startLayout := Layout.Parent
    else
        startLayout := nil;

    var checker := TTypeChecker.Create(startLayout, RootScope, ALog) as IAstTypeChecker;
    Result := checker.Execute(RootNode);
end;

function TTypeChecker.Execute(const RootNode: IAstNode): IAstNode;
begin
    Result := Accept(RootNode);
    if not Assigned(Result) then
        Result := TAst.Block([], nil);
end;

// --- Helper ---

function TTypeChecker.PrepareBaseType(const BaseNode: IAstNode; out IsOptional: Boolean): IStaticType;
var
    baseType: IStaticType;
begin
    baseType := BaseNode.AsTypedNode.StaticType;
    IsOptional := baseType.IsOptional;
    Result := TTypes.Unwrap(baseType);
end;

function TTypeChecker.ApplyOptionality(const AType: IStaticType; IsOptional: Boolean): IStaticType;
begin
    if IsOptional and (AType.Kind <> stUnknown) then
        Result := TTypes.MakeOptional(AType)
    else
        Result := AType;
end;

// --- Visits ---

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

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

function TTypeChecker.VisitTuple(const Node: IAstNode): IAstNode;
var
    T: ITupleNode;
    newElements: TArray<IAstNode>;
    elementTypes: TArray<IStaticType>;
    i: Integer;

    // Inference variables
    firstType: IStaticType;
    isHomogeneous: Boolean;
    commonDim: TArray<Integer>;
    newDim: TArray<Integer>;

    finalType: IStaticType;
    elementsArray: TArray<IAstNode>;
begin
    T := Node.AsTuple;
    elementsArray := T.Elements;
    var count := Length(elementsArray);

    SetLength(newElements, count);
    SetLength(elementTypes, count);

    // 1. Visit Children
    // Recursively type-check all elements first to determine their static types.
    for i := 0 to count - 1 do
    begin
        newElements[i] := Accept(elementsArray[i]);
        if newElements[i].IsTyped then
            elementTypes[i] := newElements[i].AsTypedNode.StaticType
        else
            elementTypes[i] := TTypes.Unknown;
    end;

    // 2. Inference Logic: Tuple vs. Vector vs. Matrix
    if count = 0 then
    begin
        // Empty Tuple -> stTuple (safest fallback, effectively Void)
        finalType := TTypes.CreateTuple([]);
    end
    else
    begin
        firstType := elementTypes[0];
        isHomogeneous := True;

        for i := 1 to count - 1 do
        begin
            if not firstType.IsEqual(elementTypes[i]) then
            begin
                isHomogeneous := False;
                break;
            end;
        end;

        if isHomogeneous and (firstType.Kind <> stUnknown) then
        begin
            if firstType.Kind = stVector then
            begin
                // Vector of Vectors -> Matrix (2D)
                newDim := [count, firstType.AsVector.Count];
                finalType := TTypes.CreateMatrix(firstType.AsVector.ElementType, newDim);
            end
            else if firstType.Kind = stMatrix then
            begin
                // Vector of Matrices -> Higher dimensional Matrix
                commonDim := firstType.AsMatrix.Dimensions;
                SetLength(newDim, Length(commonDim) + 1);

                newDim[0] := count;
                for i := 0 to High(commonDim) do
                    newDim[i + 1] := commonDim[i];

                finalType := TTypes.CreateMatrix(firstType.AsMatrix.ElementType, newDim);
            end
            else
            begin
                // Base case: Homogeneous Scalars/Records -> Vector (1D)
                finalType := TTypes.CreateVector(firstType, count);
            end;
        end
        else
        begin
            // Heterogeneous types -> Standard Tuple
            finalType := TTypes.CreateTuple(elementTypes);
        end;
    end;

    // 3. Return the new node with the inferred type definition
    Result := TAst.Tuple(Node.Identity, newElements, finalType);
end;

function TTypeChecker.VisitIdentifier(const Node: IAstNode): IAstNode;
var
    I: IIdentifierNode;
    typ: IStaticType;
    adr: TResolvedAddress;
    identity: INamedIdentity;
begin
    I := Node.AsIdentifier;
    adr := I.Address;
    identity := I.Identity.AsNamed;

    if adr.Kind = akUnresolved then
    begin
        Result := TAst.Identifier(identity, adr, TTypes.Unknown);
        Exit;
    end;

    typ := FCurrentContext.LookupType(adr);
    Result := TAst.Identifier(identity, adr, typ);
end;

function TTypeChecker.VisitRecurNode(const Node: IAstNode): IAstNode;
var
    R: IRecurNode;
    args: ITupleNode;
begin
    R := Node.AsRecur;
    args := Accept(R.Arguments).AsTuple;
    Result := TAst.Recur(Node.Identity, args, TTypes.Void);
end;

function TTypeChecker.VisitVariableDeclaration(const Node: IAstNode): IAstNode;
var
    V: IVariableDeclarationNode;
    initType: IStaticType;
    newInitializer, newIdent: IAstNode;
    adr: TResolvedAddress;
    identNode: IIdentifierNode;
begin
    V := Node.AsVariableDeclaration;
    identNode := V.Target.AsIdentifier;
    adr := identNode.Address;
    initType := TTypes.Unknown;

    if adr.Kind = akUnresolved then
    begin
        if Assigned(V.Initializer) then
            Accept(V.Initializer);
        Result := Node;
        Exit;
    end;

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

    if Assigned(newInitializer) then
        initType := newInitializer.AsTypedNode.StaticType;

    if initType.Kind <> stUnknown then
        FCurrentContext.SetType(adr.SlotIndex, initType);

    newIdent := TAst.Identifier(identNode.Identity.AsNamed, adr, initType);
    Result := TAst.VarDecl(Node.Identity, newIdent, newInitializer, initType, V.IsBoxed);
end;

function TTypeChecker.VisitAssignment(const Node: IAstNode): IAstNode;
var
    A: IAssignmentNode;
    targetType, sourceType: IStaticType;
    newIdent, newValue: IAstNode;
    adr: TResolvedAddress;
    identNode: IIdentifierNode;
begin
    A := Node.AsAssignment;
    identNode := A.Target.AsIdentifier;
    newIdent := Accept(A.Target);
    targetType := newIdent.AsTypedNode.StaticType;
    adr := identNode.Address;

    if adr.Kind = akUnresolved then
    begin
        Accept(A.Value);
        Result := Node;
        Exit;
    end;

    newValue := Accept(A.Value);
    sourceType := newValue.AsTypedNode.StaticType;

    if (targetType.Kind <> stUnknown) and (sourceType.Kind <> stUnknown) then
    begin
        if not TTypeRules.CanAssign(targetType, sourceType) then
        begin
            if Assigned(FLog) then
                FLog.AddError(Format('Cannot assign type %s to %s', [sourceType.ToString, targetType.ToString]), Node);
        end;
    end;

    if (targetType.Kind = stUnknown) and (sourceType.Kind <> stUnknown) then
    begin
        FCurrentContext.SetType(adr.SlotIndex, sourceType);
        newIdent := TAst.Identifier(identNode.Identity.AsNamed, adr, sourceType);
        targetType := sourceType;
    end;

    Result := TAst.Assign(Node.Identity, newIdent, newValue, targetType);
end;

function TTypeChecker.VisitBlockExpression(const Node: IAstNode): IAstNode;
var
    newBlock: IBlockExpressionNode;
    blockType: IStaticType;
    exprs: TArray<IAstNode>;
begin
    var transformedNode := inherited VisitBlockExpression(Node);

    if not Assigned(transformedNode) then
        raise ECompilationFailed.Create([TCompilerError.Create(elError, 'Internal Error: Block transformation returned nil.', Node)]);

    newBlock := transformedNode.AsBlockExpression;
    exprs := newBlock.Expressions.Elements;

    if Length(exprs) > 0 then
    begin
        // The type of the block is the type of the last expression
        blockType := exprs[High(exprs)].AsTypedNode.StaticType;
    end
    else
    begin
        blockType := TTypes.Void;
    end;

    Result := TAst.Block(Node.Identity, newBlock.Expressions, blockType);
end;

function TTypeChecker.VisitLambdaExpression(const Node: IAstNode): IAstNode;
var
    L: ILambdaExpressionNode;
    newParams: TArray<IAstNode>;
    newBody: IAstNode;
    bodyType, methodType: IStaticType;
    paramTypes: TArray<IStaticType>;
    upvalueTypes: TArray<IStaticType>;
    i: Integer;
    finalDescriptor: IScopeDescriptor;
    paramIdent: IIdentifierNode;
    injectedType: IStaticType;
    paramsTuple: ITupleNode;
    paramsElements: TArray<IAstNode>;
begin
    L := Node.AsLambdaExpression;

    var upvalueAddrs := L.Upvalues;
    SetLength(upvalueTypes, Length(upvalueAddrs));

    for i := 0 to High(upvalueAddrs) do
    begin
        var lookupAddr := upvalueAddrs[i];
        if lookupAddr.Kind = akLocalOrParent then
        begin
            if lookupAddr.ScopeDepth > 0 then
                Dec(lookupAddr.ScopeDepth);
        end;
        upvalueTypes[i] := FCurrentContext.LookupType(lookupAddr);
    end;

    FCurrentContext := TTypeContext.Create(FCurrentContext, L.Layout, upvalueTypes, nil);
    try
        paramsTuple := L.Parameters;
        paramsElements := paramsTuple.Elements;
        SetLength(newParams, Length(paramsElements));
        SetLength(paramTypes, Length(paramsElements));

        for i := 0 to High(paramsElements) do
        begin
            paramIdent := paramsElements[i].AsIdentifier;
            injectedType := paramIdent.AsTypedNode.StaticType;

            if injectedType.Kind = stUnknown then
                paramTypes[i] := TTypes.Unknown
            else
                paramTypes[i] := injectedType;

            if paramIdent.Address.Kind <> akUnresolved then
                FCurrentContext.SetType(paramIdent.Address.SlotIndex, paramTypes[i]);

            newParams[i] := TAst.Identifier(paramIdent.Identity.AsNamed, paramIdent.Address, paramTypes[i]);
        end;

        newBody := Accept(L.Body);

        bodyType := newBody.AsTypedNode.StaticType;
        methodType := TTypes.CreateMethod(paramTypes, bodyType);

        finalDescriptor := TScope.CreateDescriptor(L.Layout, FCurrentContext.Types);
    finally
        var temp := FCurrentContext;
        FCurrentContext := FCurrentContext.FParent;
        temp.Free;
    end;

    var paramList := TAst.Tuple(L.Parameters.Identity, newParams);

    Result :=
        TAst.LambdaExpr(Node.Identity, paramList, newBody, L.Layout, finalDescriptor, L.Upvalues, L.HasNestedLambdas, L.IsPure, methodType);
end;

function TTypeChecker.VisitFunctionCall(const Node: IAstNode): IAstNode;
var
    newCall: IFunctionCallNode;
    calleeType, retType: IStaticType;
    i, j: Integer;
    argTypes: TArray<IStaticType>;
    hasUnknownArgs: Boolean;
    bestSig: IMethodSignature;
    match: Boolean;
    newArgs: ITupleNode;
    argsElements: TArray<IAstNode>;
begin
    newCall := inherited VisitFunctionCall(Node).AsFunctionCall;
    var newCallee := newCall.Callee;
    newArgs := newCall.Arguments;
    argsElements := newArgs.Elements;

    SetLength(argTypes, Length(argsElements));
    hasUnknownArgs := False;

    for i := 0 to High(argsElements) do
    begin
        argTypes[i] := argsElements[i].AsTypedNode.StaticType;
        if argTypes[i].Kind = stUnknown then
            hasUnknownArgs := True;
    end;

    calleeType := newCallee.AsTypedNode.StaticType;
    retType := TTypes.Unknown;

    if calleeType.Kind = stMethod then
    begin
        if not hasUnknownArgs then
        begin
            bestSig := nil;
            for var sig in calleeType.AsMethod.Signatures do
            begin
                if Length(sig.ParamTypes) <> Length(argTypes) then
                    continue;
                match := True;
                for j := 0 to High(argTypes) do
                    if not TTypeRules.CanAssign(sig.ParamTypes[j], argTypes[j]) then
                    begin
                        match := False;
                        break;
                    end;
                if match then
                begin
                    bestSig := sig;
                    break;
                end;
            end;

            if Assigned(bestSig) then
                retType := bestSig.ReturnType
            else if Assigned(FLog) then
                FLog.AddError(Format('No matching signature found for method call on %s', [calleeType.ToString]), Node);
        end;
    end
    else if calleeType.Kind <> stUnknown then
        if Assigned(FLog) then
            FLog.AddError(Format('Cannot invoke type %s as a function.', [calleeType.ToString]), Node);

    Result := TAst.FunctionCall(Node.Identity, newCallee, newArgs, retType, newCall.IsTailCall, nil, False);
end;

function TTypeChecker.VisitIfExpression(const Node: IAstNode): IAstNode;
var
    newIf: IIfExpressionNode;
    resType: IStaticType;
begin
    newIf := inherited VisitIfExpression(Node).AsIfExpression;

    if (newIf.ElseBranch <> nil) then
        resType := TTypeRules.Promote(newIf.ThenBranch.AsTypedNode.StaticType, newIf.ElseBranch.AsTypedNode.StaticType)
    else
        resType := TTypes.MakeOptional(newIf.ThenBranch.AsTypedNode.StaticType);

    Result := TAst.IfExpr(Node.Identity, newIf.Condition, newIf.ThenBranch, newIf.ElseBranch, resType);
end;

function TTypeChecker.VisitIndexer(const Node: IAstNode): IAstNode;
var
    I: IIndexerNode;
    newBase, newIndex: IAstNode;
    baseType, elemType: IStaticType;
    isOpt: Boolean;
begin
    I := Node.AsIndexer;
    newBase := Accept(I.Base);
    newIndex := Accept(I.Index);
    baseType := PrepareBaseType(newBase, isOpt);

    elemType := TTypes.Unknown;
    if baseType.Kind = stSeries then
        elemType := baseType.AsSeries.ElementType
    else if baseType.Kind = stRecordSeries then
        elemType := TTypes.CreateRecord(baseType.AsRecord.Definition)
    else if baseType.Kind = stTuple then
    begin
        // Tuple Indexing (requires constant index for strong typing!)
        if (newIndex.Kind = akConstant) and (newIndex.AsConstant.Value.Kind = vkScalar) then
        begin
            var idx := newIndex.AsConstant.Value.AsScalar.Value.AsInt64;
            var tpl := baseType.AsTuple;
            if (idx >= 0) and (idx < tpl.Count) then
                elemType := tpl.Elements[Integer(idx)];
        end;
    end;

    elemType := ApplyOptionality(elemType, isOpt);
    Result := TAst.Indexer(Node.Identity, newBase, newIndex, elemType);
end;

function TTypeChecker.VisitMemberAccess(const Node: IAstNode): IAstNode;
var
    M: IMemberAccessNode;
    newBase: IAstNode;
    baseType, resType: IStaticType;
    idx: Integer;
    isOpt: Boolean;
begin
    M := Node.AsMemberAccess;
    newBase := Accept(M.Base);
    baseType := PrepareBaseType(newBase, isOpt);
    resType := TTypes.Unknown;

    if (baseType.Kind = stRecord) or (baseType.Kind = stRecordSeries) then
    begin
        idx := baseType.AsRecord.Definition.IndexOf(M.Member.Value);
        if idx >= 0 then
        begin
            var fieldType := TTypes.FromScalarKind(baseType.AsRecord.Definition[idx]);
            if baseType.Kind = stRecordSeries then
                resType := TTypes.CreateSeries(fieldType)
            else
                resType := fieldType;
        end
        else if Assigned(FLog) then
            FLog.AddError('Member not found', Node);
    end;

    resType := ApplyOptionality(resType, isOpt);
    Result := TAst.MemberAccess(Node.Identity, newBase, M.Member, resType);
end;

function TTypeChecker.VisitRecordLiteral(const Node: IAstNode): IAstNode;
var
    R: IRecordLiteralNode;
    i: Integer;
    fieldTypes: TArray<TPair<IKeyword, IStaticType>>;
    scalarFieldTypes: TArray<TPair<IKeyword, TScalar.TKind>>;
    isScalar: Boolean;
    valType: IStaticType;
    key: IKeyword;
    newFields: ITupleNode;
    fieldsElements: TArray<IAstNode>;
begin
    R := Node.AsRecordLiteral;
    newFields := Visit(R.Fields).AsTuple;
    fieldsElements := newFields.Elements;

    var count := Length(fieldsElements);
    SetLength(fieldTypes, count);
    SetLength(scalarFieldTypes, count);
    isScalar := True;

    for i := 0 to count - 1 do
    begin
        var field := fieldsElements[i].AsRecordField;
        key := field.Key.Value;
        valType := field.Value.AsTypedNode.StaticType;

        if (valType.Kind in [stOrdinal, stFloat, stBoolean, stDateTime, stKeyword]) and (not valType.IsOptional) then
        begin
            var kind: TScalar.TKind;
            case valType.Kind of
                stOrdinal: kind := TScalar.TKind.Ordinal;
                stFloat: kind := TScalar.TKind.Float;
                stBoolean: kind := TScalar.TKind.Boolean;
                stDateTime: kind := TScalar.TKind.DateTime;
                stKeyword: kind := TScalar.TKind.Keyword;
            else
                kind := TScalar.TKind.Ordinal;
            end;
            scalarFieldTypes[i] := TPair<IKeyword, TScalar.TKind>.Create(key, kind);
        end
        else
        begin
            isScalar := False;
        end;

        fieldTypes[i] := TPair<IKeyword, IStaticType>.Create(key, valType);
    end;

    var scalarDef: IScalarRecordDefinition := nil;
    var genericDef: IGenericRecordDefinition := nil;
    var resultType: IStaticType;

    if isScalar and (count > 0) then
    begin
        scalarDef := TKeywordMappingRegistry<TScalar.TKind>.Intern(scalarFieldTypes);
        resultType := TTypes.CreateRecord(scalarDef);
    end
    else
    begin
        genericDef := TGenericRecordRegistry.Intern(fieldTypes);
        resultType := TTypes.CreateGenericRecord(genericDef);
    end;

    Result := TAst.RecordLiteral(Node.Identity, newFields, scalarDef, genericDef, resultType);
end;

function TTypeChecker.VisitCreateSeries(const Node: IAstNode): IAstNode;
var
    C: ICreateSeriesNode;
    defNode: IAstNode;
    recDef: IScalarRecordDefinition;
    resType: IStaticType;
    elemKind: TScalar.TKind;
    fields: TArray<TPair<IKeyword, TScalar.TKind>>;
    tuple: ITupleNode;
    elements: TArray<IAstNode>;
    fieldTuple: ITupleNode;
    keyNode: IKeywordNode;
    typeNode: IKeywordNode;
    i: Integer;
begin
    C := Node.AsCreateSeries;
    // The definition node comes from the parser and is likely raw (untyped keywords/tuples).
    // We don't necessarily need to "Visit" it for type checking children,
    // but we need to analyze its structure.
    defNode := C.DefinitionNode;

    recDef := nil;
    resType := TTypes.Unknown;

    case defNode.Kind of
        akKeyword:
        begin
            // Case 1: Simple Series, e.g. (new-series :Float)
            // DefinitionNode is a single Keyword
            elemKind := TScalar.StringToKind(defNode.AsKeyword.Value.Name);
            resType := TTypes.CreateSeries(TTypes.FromScalarKind(elemKind));
        end;
        akTuple:
        begin
            // Case 2: Record Series, e.g. (new-series [[:Price :Float] [:Vol :Ordinal]])
            tuple := defNode.AsTuple;
            elements := tuple.Elements;
            SetLength(fields, Length(elements));

            for i := 0 to High(elements) do
            begin
                if elements[i].Kind <> akTuple then
                begin
                    if Assigned(FLog) then
                        FLog.AddError('Record definition elements must be vectors [Key Type]', elements[i]);
                    continue;
                end;

                fieldTuple := elements[i].AsTuple;
                if Length(fieldTuple.Elements) <> 2 then
                begin
                    if Assigned(FLog) then
                        FLog.AddError('Record definition entry must have exactly 2 elements [Key Type]', fieldTuple);
                    continue;
                end;

                if (fieldTuple.Elements[0].Kind <> akKeyword) or (fieldTuple.Elements[1].Kind <> akKeyword) then
                begin
                    if Assigned(FLog) then
                        FLog.AddError('Record definition keys and types must be keywords', fieldTuple);
                    continue;
                end;

                keyNode := fieldTuple.Elements[0].AsKeyword;
                typeNode := fieldTuple.Elements[1].AsKeyword;

                elemKind := TScalar.StringToKind(typeNode.Value.Name);
                fields[i] := TPair<IKeyword, TScalar.TKind>.Create(keyNode.Value, elemKind);
            end;

            if Length(fields) > 0 then
            begin
                recDef := TKeywordMappingRegistry<TScalar.TKind>.Intern(fields);
                resType := TTypes.CreateRecordSeries(recDef);
            end;
        end
    else
        begin
            if Assigned(FLog) then
                FLog.AddError('Invalid series definition format. Expected :Type or [[:Key :Type] ...]', defNode);
        end;
    end;

    // Return updated node with RecordDefinition (if any) and StaticType
    Result := TAst.CreateSeries(Node.Identity, defNode, recDef, resType);
end;

function TTypeChecker.VisitSeriesLength(const Node: IAstNode): IAstNode;
begin
    Result := TAst.SeriesLength(Node.Identity, Accept(Node.AsSeriesLength.Series).AsIdentifier, TTypes.Ordinal);
end;

function TTypeChecker.VisitNop(const Node: IAstNode): IAstNode;
begin
    Result := TAst.Nop(Node.Identity, TTypes.Void);
end;

// =============================================================================
// PIPE IMPLEMENTATION (Type Checking)
// =============================================================================

function TTypeChecker.VisitPipe(const Node: IAstNode): IAstNode;
var
    P: IPipeNode;
    rawInputs: ITupleNode;
    i, k: Integer;
    newInputs: TList<IAstNode>;
    paramTypes: TList<IStaticType>;
    lambda: ILambdaExpressionNode;
    newParams: TArray<IIdentifierNode>;
    newParamsAsNodes: TArray<IAstNode>;
    inferredType: IStaticType;

    // Iteration vars
    inputPair: IAstNode;
    rawTuple: ITupleNode;
    streamNode: IIdentifierNode;
    selectorsNode: ITupleNode;
    newSelectors: TList<IAstNode>;
    streamType: IStaticType;
    rawInputsElements: TArray<IAstNode>;
    rawTupleElements: TArray<IAstNode>;
    selectorsElements: TArray<IAstNode>;
    lambdaParamsElements: TArray<IAstNode>;
begin
    P := Node.AsPipe;
    rawInputs := P.Inputs; // This is the Tuple of Tuples [[id [sel]] ...]
    rawInputsElements := rawInputs.Elements;

    newInputs := TList<IAstNode>.Create;
    paramTypes := TList<IStaticType>.Create;
    try
        // 1. Iterate over input definitions
        for i := 0 to High(rawInputsElements) do
        begin
            inputPair := rawInputsElements[i];

            // Validate Structure: [StreamSource, [Selectors]]
            if inputPair.Kind <> akTuple then
            begin
                if Assigned(FLog) then
                    FLog.AddError('Pipe input must be a vector: [Source [Selectors]].', inputPair);
                continue;
            end;

            rawTuple := inputPair.AsTuple;
            rawTupleElements := rawTuple.Elements;

            if Length(rawTupleElements) <> 2 then
            begin
                if Assigned(FLog) then
                    FLog.AddError('Pipe input vector must have exactly 2 elements.', inputPair);
                continue;
            end;

            // Element 0: Stream Identifier
            if rawTupleElements[0].Kind <> akIdentifier then
            begin
                if Assigned(FLog) then
                    FLog.AddError('Pipe source must be an identifier.', rawTupleElements[0]);
                continue;
            end;

            // Resolve Stream Identifier (Type Binding)
            streamNode := Accept(rawTupleElements[0]).AsIdentifier;
            streamType := streamNode.StaticType;

            // Validate Stream Type
            if (streamType.Kind <> stUnknown) then
            begin
                if not ((streamType.Kind = stSeries) or (streamType.Kind = stRecordSeries)) then
                begin
                    if Assigned(FLog) then
                        FLog.AddError(
                            Format('Pipe input "%s" must be a Series or RecordSeries, but got %s.', [streamNode.Name, streamType.ToString]),
                            streamNode
                        );
                end;
            end;

            // Element 1: Selector Vector
            if rawTupleElements[1].Kind <> akTuple then
            begin
                if Assigned(FLog) then
                    FLog.AddError('Pipe selectors must be a vector of keywords.', rawTupleElements[1]);
                continue;
            end;
            selectorsNode := rawTupleElements[1].AsTuple;
            selectorsElements := selectorsNode.Elements;

            newSelectors := TList<IAstNode>.Create;
            try
                // Iterate Selectors
                for k := 0 to High(selectorsElements) do
                begin
                    var selItem := selectorsElements[k];
                    if selItem.Kind <> akKeyword then
                    begin
                        if Assigned(FLog) then
                            FLog.AddError('Selector must be a keyword.', selItem);
                        continue;
                    end;

                    var kw := selItem.AsKeyword;
                    newSelectors.Add(kw); // Keywords don't need 'Accept' as they are constant, but we keep them structurally

                    // Infer Type for Lambda Parameter
                    inferredType := TTypes.Unknown;
                    if streamType.Kind = stRecordSeries then
                    begin
                        var def := streamType.AsRecord.Definition;
                        var idx := def.IndexOf(kw.Value);
                        if idx >= 0 then
                            inferredType := TTypes.FromScalarKind(def[idx])
                        else if Assigned(FLog) then
                            FLog.AddError(Format('Field ":%s" not found in stream "%s".', [kw.Value.Name, streamNode.Name]), selItem);
                    end
                    else if streamType.Kind = stSeries then
                    begin
                        if Assigned(streamType.AsSeries.ElementType) then
                            inferredType := streamType.AsSeries.ElementType
                        else
                            inferredType := TTypes.Ordinal;
                    end;
                    paramTypes.Add(inferredType);
                end;

                // Reconstruct the typed Input Tuple [StreamIdent, [Selectors]]
                var typedSelectorTuple := TAst.Tuple(selectorsNode.Identity, newSelectors.ToArray, TTypes.Unknown);
                var typedInputPair := TAst.Tuple(rawTuple.Identity, [streamNode, typedSelectorTuple], TTypes.Unknown);
                newInputs.Add(typedInputPair);
            finally
                newSelectors.Free;
            end;
        end;

        // 2. Process Transformation Lambda
        lambda := P.Transformation;
        lambdaParamsElements := lambda.Parameters.Elements;

        if Length(lambdaParamsElements) <> paramTypes.Count then
        begin
            if Assigned(FLog) then
                FLog.AddError(
                    Format(
                        'Pipe lambda expects %d parameters (one per selector), but got %d.',
                        [paramTypes.Count, Length(lambdaParamsElements)]
                    ),
                    lambda
                );
        end;

        SetLength(newParams, Length(lambdaParamsElements));
        for k := 0 to High(lambdaParamsElements) do
        begin
            var oldP := lambdaParamsElements[k].AsIdentifier;
            var typeToInject :=
                if k < paramTypes.Count then paramTypes[k]
                else TTypes.Unknown;

            newParams[k] := TAst.Identifier(oldP.Identity.AsNamed, oldP.Address, typeToInject);
        end;

        // Convert for Tuple Factory
        SetLength(newParamsAsNodes, Length(newParams));
        for k := 0 to High(newParams) do
            newParamsAsNodes[k] := newParams[k];

        var preTypedLambda :=
            TAst.LambdaExpr(
                lambda.Identity,
                TAst.Tuple(lambda.Parameters.Identity, newParamsAsNodes),
                lambda.Body,
                lambda.Layout,
                lambda.Descriptor,
                lambda.Upvalues,
                lambda.HasNestedLambdas,
                lambda.IsPure,
                TTypes.Unknown
            );

        var typedLambda := Accept(preTypedLambda).AsLambdaExpression;

        // 3. Determine Result Type
        var lambdaRetType := typedLambda.AsTypedNode.StaticType.AsMethod.Signatures[0].ReturnType;
        var pipeType: IStaticType;

        if lambdaRetType.Kind = stRecord then
        begin
            pipeType := TTypes.CreateRecordSeries(lambdaRetType.AsRecord.Definition);
        end
        else
        begin
            if (lambdaRetType.Kind <> stUnknown) and (lambdaRetType.Kind <> stVoid) then
            begin
                if Assigned(FLog) then
                    FLog.AddError(
                        Format('Pipe function must return a Record (e.g. {:res ...}). Type "%s" is invalid.', [lambdaRetType.ToString]),
                        lambda
                    );
            end;
            pipeType := TTypes.Unknown;
        end;

        // Reconstruct the Pipe Node with typed Inputs tuple and typed Lambda
        var typedInputsTuple := TAst.Tuple(rawInputs.Identity, newInputs.ToArray, TTypes.Unknown);
        Result := TAst.Pipe(Node.Identity, typedInputsTuple, typedLambda, pipeType);

    finally
        newInputs.Free;
        paramTypes.Free;
    end;
end;

end.
