unit Myc.Ast.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,
    Myc.Ast.Binding.Nodes;

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

    // This transformer runs *after* the TAstBinder.
    // It takes the "Bound AST" (which has addresses but mostly TTypes.Unknown)
    // and traverses it bottom-up to infer and check all static types.
    // It modifies the node.StaticType property and updates the IScopeDescriptor.
    TTypeChecker = class(TAstTransformer, IAstTypeChecker)
    private
        FCurrentDescriptor: IScopeDescriptor;
        function SetType(const NodeData: TDataValue; const AType: IStaticType): TDataValue; overload;
    protected
        // Override all visit methods to perform type checking
        function VisitIdentifier(const Node: IIdentifierNode): TDataValue; override;
        function VisitVariableDeclaration(const Node: IVariableDeclarationNode): TDataValue; override;
        function VisitAssignment(const Node: IAssignmentNode): TDataValue; override;
        function VisitLambdaExpression(const Node: ILambdaExpressionNode): TDataValue; override;
        function VisitFunctionCall(const Node: IFunctionCallNode): TDataValue; override;
        function VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue; override;
        function VisitIfExpression(const Node: IIfExpressionNode): TDataValue; override;
        function VisitTernaryExpression(const Node: ITernaryExpressionNode): TDataValue; override;
        function VisitBinaryExpression(const Node: IBinaryExpressionNode): TDataValue; override;
        function VisitUnaryExpression(const Node: IUnaryExpressionNode): TDataValue; override;
        function VisitMemberAccess(const Node: IMemberAccessNode): TDataValue; override;
        function VisitIndexer(const Node: IIndexerNode): TDataValue; override;
        function VisitRecordLiteral(const Node: IRecordLiteralNode): TDataValue; override;
        function VisitCreateSeries(const Node: ICreateSeriesNode): TDataValue; override;
        function VisitAddSeriesItem(const Node: IAddSeriesItemNode): TDataValue; override;
        function VisitSeriesLength(const Node: ISeriesLengthNode): TDataValue; override;
        function VisitRecurNode(const Node: IRecurNode): TDataValue; override;

        // Base cases (types are already known from binder)
        function VisitConstant(const Node: IConstantNode): TDataValue; override;
        function VisitKeyword(const Node: IKeywordNode): TDataValue; override;
        // Compile-time nodes (should not be present)
        function VisitMacroExpansionNode(const Node: IMacroExpansionNode): TDataValue; override;
        function VisitMacroDefinition(const Node: IMacroDefinitionNode): TDataValue; override;
    public
        constructor Create(const ADescriptor: IScopeDescriptor);
        function Execute(const RootNode: IAstNode; const ADescriptor: IScopeDescriptor): IAstNode;

        class function CheckTypes(const RootNode: IAstNode; const ADescriptor: IScopeDescriptor): IAstNode; static;
    end;

implementation

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

{ TTypeChecker }

constructor TTypeChecker.Create(const ADescriptor: IScopeDescriptor);
begin
    inherited Create;
    Assert(Assigned(ADescriptor));
    FCurrentDescriptor := ADescriptor;
end;

class function TTypeChecker.CheckTypes(const RootNode: IAstNode; const ADescriptor: IScopeDescriptor): IAstNode;
begin
    var checker := TTypeChecker.Create(ADescriptor) as IAstTypeChecker;
    Result := checker.Execute(RootNode, ADescriptor);
end;

function TTypeChecker.Execute(const RootNode: IAstNode; const ADescriptor: IScopeDescriptor): IAstNode;
begin
    FCurrentDescriptor := ADescriptor;
    var transformedValue := Accept(RootNode);
    if transformedValue.IsVoid then
        Result := TAst.Block([])
    else
        Result := transformedValue.AsIntf<IAstNode>;
    (Result as TAstNode).StaticType := (Result as TAstNode).StaticType;
end;

function TTypeChecker.SetType(const NodeData: TDataValue; const AType: IStaticType): TDataValue;
begin
    if (not NodeData.IsVoid) and (NodeData.Kind = vkInterface) then
        (NodeData.AsIntf<IAstNode> as TAstNode).StaticType := AType;
    Result := NodeData;
end;

function TTypeChecker.VisitConstant(const Node: IConstantNode): TDataValue;
begin
    // Type was set by Binder, just propagate it up.
    Result := TDataValue.FromIntf<IConstantNode>(Node);
end;

function TTypeChecker.VisitKeyword(const Node: IKeywordNode): TDataValue;
begin
    // Type was set by Binder, just propagate it up.
    Result := TDataValue.FromIntf<IKeywordNode>(Node);
end;

function TTypeChecker.VisitIdentifier(const Node: IIdentifierNode): TDataValue;
begin
    // Type was set by Binder (read from scope), just propagate it up.
    Result := TDataValue.FromIntf<IIdentifierNode>(Node);
end;

function TTypeChecker.VisitRecurNode(const Node: IRecurNode): TDataValue;
begin
    // Type was set by Binder (TTypes.Void), just propagate it up.
    Result := TDataValue.FromIntf<IRecurNode>(Node);
end;

function TTypeChecker.VisitMacroDefinition(const Node: IMacroDefinitionNode): TDataValue;
begin
    raise Exception.Create('TTypeChecker: MacroDefinition node encountered.');
end;

function TTypeChecker.VisitMacroExpansionNode(const Node: IMacroExpansionNode): TDataValue;
begin
    raise Exception.Create('TTypeChecker: MacroExpansionNode node encountered.');
end;

function TTypeChecker.VisitVariableDeclaration(const Node: IVariableDeclarationNode): TDataValue;
var
    initNode: IAstNode;
    initType: IStaticType;
    boundIdent: TBoundIdentifierNode;
    adr: TResolvedAddress;
begin
    // 1. Visit the initializer (if it exists) to get its (now-inferred) type.
    if Assigned(Node.Initializer) then
    begin
        initNode := Accept(Node.Initializer).AsIntf<IAstNode>;
        initType := (initNode as TAstNode).StaticType;
    end
    else
        initType := TTypes.Void;

    // 2. Get the address from the bound identifier.
    boundIdent := (Node.Identifier as TBoundIdentifierNode);
    adr := boundIdent.Address;

    // 3. Update the type in the scope descriptor (which was set to Unknown by the binder).
    FCurrentDescriptor.UpdateType(adr.SlotIndex, initType);

    // 4. Update the static types of the nodes themselves.
    (boundIdent as TAstNode).StaticType := initType;
    Result := SetType(TDataValue.FromIntf(Node), initType);
end;

function TTypeChecker.VisitAssignment(const Node: IAssignmentNode): TDataValue;
var
    boundIdentifier, boundValue: IAstNode;
    targetType, sourceType: IStaticType;
begin
    boundIdentifier := Accept(Node.Identifier).AsIntf<IAstNode>;
    boundValue := Accept(Node.Value).AsIntf<IAstNode>;

    targetType := (boundIdentifier as TAstNode).StaticType;
    sourceType := (boundValue as TAstNode).StaticType;

    if not TTypeRules.CanAssign(targetType, sourceType) then
        raise ETypeException.CreateFmt('Cannot assign type %s to %s', [sourceType.ToString, targetType.ToString]);

    Result := SetType(TDataValue.FromIntf<IAssignmentNode>(Node), targetType);
end;

function TTypeChecker.VisitLambdaExpression(const Node: ILambdaExpressionNode): TDataValue;
var
    boundNode: TBoundLambdaExpressionNode;
    boundBody: IAstNode;
    bodyType, methodType: IStaticType;
    paramTypes: TArray<IStaticType>;
    i: Integer;
begin
    boundNode := (Node as TBoundLambdaExpressionNode);

    // 1. Enter the lambda's scope
    FCurrentDescriptor := boundNode.ScopeDescriptor;
    try
        // 2. Set parameter types (currently Unknown, but required for signature)
        SetLength(paramTypes, Length(boundNode.Parameters));
        for i := 0 to High(boundNode.Parameters) do
            paramTypes[i] := (boundNode.Parameters[i] as TAstNode).StaticType; // Propagates Unknown

        // 3. Visit the body to infer its return type
        boundBody := Accept(boundNode.Body).AsIntf<IAstNode>;
        bodyType := (boundBody as TAstNode).StaticType;

        // 4. Create the final method type
        methodType := TTypes.CreateMethod(paramTypes, bodyType);

        // 5. Update the type for <self> (Slot 0) in the descriptor
        FCurrentDescriptor.UpdateType(0, methodType);

    finally
        // 6. Restore parent descriptor
        FCurrentDescriptor := FCurrentDescriptor.Parent;
    end;

    // 7. Set the type of the lambda node itself
    Result := SetType(TDataValue.FromIntf<ILambdaExpressionNode>(boundNode), methodType);
end;

function TTypeChecker.VisitFunctionCall(const Node: IFunctionCallNode): TDataValue;
var
    callee: IAstNode;
    args: TArray<IAstNode>;
    calleeType, retType: IStaticType;
    i: Integer;
begin
    // 1. Visit children first (bottom-up)
    callee := Accept(Node.Callee).AsIntf<IAstNode>;
    args := AcceptNodes<IAstNode>(Node.Arguments);

    // 2. Get callee type (now inferred)
    calleeType := (callee as TAstNode).StaticType;
    retType := TTypes.Unknown; // Default if not a method

    // 3. Perform type checking
    if calleeType.Kind = TStaticTypeKind.stMethod then
    begin
        var signature := calleeType.Signature;
        if Length(args) <> Length(signature.ParamTypes) then
            raise ETypeException.CreateFmt('Function expects %d arguments, but got %d', [Length(signature.ParamTypes), Length(args)]);

        retType := signature.ReturnType;

        // Check argument types
        for i := 0 to High(args) do
        begin
            var argType := (args[i] as TAstNode).StaticType;
            var paramType := signature.ParamTypes[i];
            if not TTypeRules.CanAssign(paramType, argType) then
                raise ETypeException
                    .CreateFmt('Cannot assign argument %d (type %s) to parameter (type %s)', [i, argType.ToString, paramType.ToString]);
        end;
    end
    else if calleeType.Kind <> TStaticTypeKind.stUnknown then
        raise ETypeException.CreateFmt('Cannot invoke type %s as a function.', [calleeType.ToString]);

    // 4. Set the type for this call node
    Result := SetType(TDataValue.FromIntf<IFunctionCallNode>(Node), retType);
end;

function TTypeChecker.VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue;
var
    blockType: IStaticType;
    exprs: TArray<IAstNode>;
begin
    exprs := AcceptNodes<IAstNode>(Node.Expressions);

    if Length(exprs) > 0 then
        blockType := (exprs[High(exprs)] as TAstNode).StaticType
    else
        blockType := TTypes.Void;

    Result := SetType(TDataValue.FromIntf<IBlockExpressionNode>(Node), blockType);
end;

function TTypeChecker.VisitIfExpression(const Node: IIfExpressionNode): TDataValue;
var
    condition, thenBranch, elseBranch: IAstNode;
    conditionType, thenType, elseType, resultType: IStaticType;
begin
    condition := Accept(Node.Condition).AsIntf<IAstNode>;
    thenBranch := Accept(Node.ThenBranch).AsIntf<IAstNode>;
    elseBranch := Accept(Node.ElseBranch).AsIntf<IAstNode>; // Accept(nil) returns void

    conditionType := (condition as TAstNode).StaticType;
    if (conditionType.Kind <> stUnknown) and not TTypeRules.CanAssign(TTypes.Ordinal, conditionType) then
        raise ETypeException.CreateFmt('If condition must be Ordinal, but got %s', [conditionType.ToString]);

    thenType := (thenBranch as TAstNode).StaticType;
    elseType :=
        if elseBranch <> nil then (elseBranch as TAstNode).StaticType
        else TTypes.Void;

    resultType := TTypeRules.Promote(thenType, elseType);
    Result := SetType(TDataValue.FromIntf<IIfExpressionNode>(Node), resultType);
end;

function TTypeChecker.VisitTernaryExpression(const Node: ITernaryExpressionNode): TDataValue;
var
    condition, thenBranch, elseBranch: IAstNode;
    conditionType, thenType, elseType, resultType: IStaticType;
begin
    condition := Accept(Node.Condition).AsIntf<IAstNode>;
    thenBranch := Accept(Node.ThenBranch).AsIntf<IAstNode>;
    elseBranch := Accept(Node.ElseBranch).AsIntf<IAstNode>;

    conditionType := (condition as TAstNode).StaticType;
    if (conditionType.Kind <> stUnknown) and not TTypeRules.CanAssign(TTypes.Ordinal, conditionType) then
        raise ETypeException.CreateFmt('Ternary condition must be Ordinal, but got %s', [conditionType.ToString]);

    thenType := (thenBranch as TAstNode).StaticType;
    elseType := (elseBranch as TAstNode).StaticType;
    resultType := TTypeRules.Promote(thenType, elseType);
    Result := SetType(TDataValue.FromIntf<ITernaryExpressionNode>(Node), resultType);
end;

function TTypeChecker.VisitBinaryExpression(const Node: IBinaryExpressionNode): TDataValue;
var
    left, right: IAstNode;
    leftType, rightType, resultType: IStaticType;
begin
    left := Accept(Node.Left).AsIntf<IAstNode>;
    right := Accept(Node.Right).AsIntf<IAstNode>;
    leftType := (left as TAstNode).StaticType;
    rightType := (right as TAstNode).StaticType;
    resultType := TTypeRules.ResolveBinaryOp(Node.Operator, leftType, rightType);
    Result := SetType(TDataValue.FromIntf<IBinaryExpressionNode>(Node), resultType);
end;

function TTypeChecker.VisitUnaryExpression(const Node: IUnaryExpressionNode): TDataValue;
var
    right: IAstNode;
    rightType, resultType: IStaticType;
begin
    right := Accept(Node.Right).AsIntf<IAstNode>;
    rightType := (right as TAstNode).StaticType;
    resultType := TTypeRules.ResolveUnaryOp(Node.Operator, rightType);
    Result := SetType(TDataValue.FromIntf<IUnaryExpressionNode>(Node), resultType);
end;

function TTypeChecker.VisitMemberAccess(const Node: IMemberAccessNode): TDataValue;
var
    baseNode: IAstNode;
    baseType, elemType: IStaticType;
    fieldIndex: Integer;
begin
    baseNode := Accept(Node.Base).AsIntf<IAstNode>;
    baseType := (baseNode as TAstNode).StaticType;

    elemType := TTypes.Unknown;
    if (baseType.Kind <> TStaticTypeKind.stUnknown) then
    begin
        if (baseType.Kind = TStaticTypeKind.stRecord) or (baseType.Kind = TStaticTypeKind.stRecordSeries) then
        begin
            fieldIndex := baseType.Definition.IndexOf(Node.Member.Value);
            if fieldIndex < 0 then
                raise ETypeException.CreateFmt('Member "%s" not found in type %s', [Node.Member.Value.Name, baseType.ToString]);

            var fieldType := TTypes.FromScalarKind(baseType.Definition.Fields[fieldIndex].Value);

            if baseType.Kind = TStaticTypeKind.stRecord then
                elemType := fieldType
            else // stRecordSeries
                elemType := TTypes.CreateSeries(fieldType);
        end
        else if (baseType.Kind = TStaticTypeKind.stGenericRecord) then
        begin
            var genDef := baseType.GenericDefinition;
            fieldIndex := genDef.IndexOf(Node.Member.Value);
            if fieldIndex < 0 then
                raise ETypeException.CreateFmt('Member "%s" not found in type %s', [Node.Member.Value.Name, baseType.ToString]);
            elemType := genDef.Fields[fieldIndex].Value;
        end
        else
        begin
            raise ETypeException.CreateFmt('Member access requires a record type, but got %s', [baseType.ToString]);
        end;
    end;

    Result := SetType(TDataValue.FromIntf<IMemberAccessNode>(Node), elemType);
end;

function TTypeChecker.VisitIndexer(const Node: IIndexerNode): TDataValue;
var
    baseNode, indexNode: IAstNode;
    baseType, indexType, elemType: IStaticType;
begin
    baseNode := Accept(Node.Base).AsIntf<IAstNode>;
    indexNode := Accept(Node.Index).AsIntf<IAstNode>;
    baseType := (baseNode as TAstNode).StaticType;
    indexType := (indexNode as TAstNode).StaticType;

    elemType := TTypes.Unknown;
    if (baseType.Kind <> TStaticTypeKind.stUnknown) then
    begin
        if (baseType.Kind <> TStaticTypeKind.stSeries) and (baseType.Kind <> TStaticTypeKind.stRecordSeries) then
            raise ETypeException.CreateFmt('Indexer `[]` can only be applied to series types, but got %s', [baseType.ToString]);

        if (indexType.Kind <> stUnknown) and not TTypeRules.CanAssign(TTypes.Ordinal, indexType) then
            raise ETypeException.CreateFmt('Indexer `[]` requires an Ordinal index, but got %s', [indexType.ToString]);

        if baseType.Kind = TStaticTypeKind.stSeries then
            elemType := baseType.ElementType
        else // stRecordSeries
            elemType := TTypes.CreateRecord(baseType.Definition);
    end;

    Result := SetType(TDataValue.FromIntf<IIndexerNode>(Node), elemType);
end;

function TTypeChecker.VisitRecordLiteral(const Node: IRecordLiteralNode): TDataValue;
var
    i: Integer;
    boundFields: TArray<TRecordFieldLiteral>;
    scalarDefFields: TArray<TScalarRecordField>;
    def: IScalarRecordDefinition;
    staticType: IStaticType;
    valNode: IAstNode;
    valType: IStaticType;
    scalarKind: TScalar.TKind;
    allScalar: Boolean;
begin
    SetLength(boundFields, Length(Node.Fields));
    SetLength(scalarDefFields, Length(Node.Fields));
    allScalar := True;

    // 1. Visit all child nodes first to infer their types
    for i := 0 to High(Node.Fields) do
    begin
        valNode := Accept(Node.Fields[i].Value).AsIntf<IAstNode>;
        valType := (valNode as TAstNode).StaticType;
        boundFields[i] := TRecordFieldLiteral.Create(Node.Fields[i].Key, valNode);

        // 2. Check if this field fits the scalar path
        if (valType.Kind = stOrdinal) then
            scalarKind := TScalar.TKind.Ordinal
        else if (valType.Kind = stFloat) then
            scalarKind := TScalar.TKind.Float
        else if (valType.Kind = stKeyword) then
            scalarKind := TScalar.TKind.Keyword
        else
        begin
            allScalar := False;
            scalarKind := TScalar.TKind.Ordinal; // Dummy
        end;

        if allScalar then
            scalarDefFields[i] := TScalarRecordField.Create(Node.Fields[i].Key.Value, scalarKind);
    end;

    // 3. Create the appropriate record type (Scalar or Generic)
    if allScalar then
    begin
        def := TScalarRecordRegistry.Intern(scalarDefFields);
        staticType := TTypes.CreateRecord(def);
        // We can re-use the TBoundRecordLiteralNode from the binder
        (Node as TBoundRecordLiteralNode).Definition := def;
    end
    else
    begin
        var genDefFields: TArray<TPair<IKeyword, IStaticType>>;
        SetLength(genDefFields, Length(boundFields));
        for i := 0 to High(boundFields) do
            genDefFields[i] := TPair<IKeyword, IStaticType>.Create(boundFields[i].Key.Value, (boundFields[i].Value as TAstNode).StaticType);

        var genDef := TGenericRecordRegistry.Intern(genDefFields);
        staticType := TTypes.CreateGenericRecord(genDef);
        // Re-use the TBoundGenericRecordLiteralNode from the binder
        (Node as TBoundGenericRecordLiteralNode).Definition := genDef;
    end;

    Result := SetType(TDataValue.FromIntf<IRecordLiteralNode>(Node), staticType);
end;

function TTypeChecker.VisitCreateSeries(const Node: ICreateSeriesNode): TDataValue;
begin
    // Type was set by Binder, just propagate it up.
    Result := TDataValue.FromIntf<ICreateSeriesNode>(Node);
end;

function TTypeChecker.VisitAddSeriesItem(const Node: IAddSeriesItemNode): TDataValue;
var
    seriesNode, valueNode, lookbackNode: IAstNode;
    seriesType, valueType: IStaticType;
begin
    seriesNode := Accept(Node.Series).AsIntf<IAstNode>;
    valueNode := Accept(Node.Value).AsIntf<IAstNode>;
    lookbackNode := Accept(Node.Lookback).AsIntf<IAstNode>;

    seriesType := (seriesNode as TAstNode).StaticType;
    valueType := (valueNode as TAstNode).StaticType;

    if (seriesType.Kind <> stUnknown) then
    begin
        if (seriesType.Kind <> TStaticTypeKind.stSeries) then
            raise ETypeException.CreateFmt('"add" requires a series as its first argument, but got %s', [seriesType.ToString]);

        if not TTypeRules.CanAssign(seriesType.ElementType, valueType) then
            raise ETypeException
                .CreateFmt('Cannot add item of type %s to series of type %s', [valueType.ToString, seriesType.ElementType.ToString]);
    end;

    if (lookbackNode <> nil) then
    begin
        var lookbackType := (lookbackNode as TAstNode).StaticType;
        if (lookbackType.Kind <> stUnknown) and not (lookbackType.Kind = TStaticTypeKind.stOrdinal) then
            raise ETypeException.Create('Lookback parameter for "add" must be an ordinal value.');
    end;

    Result := SetType(TDataValue.FromIntf<IAddSeriesItemNode>(Node), TTypes.Void);
end;

function TTypeChecker.VisitSeriesLength(const Node: ISeriesLengthNode): TDataValue;
var
    seriesNode: IAstNode;
    seriesType: IStaticType;
begin
    seriesNode := Accept(Node.Series).AsIntf<IAstNode>;
    seriesType := (seriesNode as TAstNode).StaticType;

    if (seriesType.Kind <> stUnknown)
        and (seriesType.Kind <> TStaticTypeKind.stSeries)
        and (seriesType.Kind <> TStaticTypeKind.stRecordSeries) then
        raise ETypeException.CreateFmt('"length" requires a series, but got %s', [seriesType.ToString]);

    Result := SetType(TDataValue.FromIntf<ISeriesLengthNode>(Node), TTypes.Ordinal);
end;

end.