unit Myc.Ast.Environment;

interface

uses
    System.SysUtils,
    System.Classes,
    System.Generics.Collections,
    System.Generics.Defaults,
    Myc.Data.Value,
    Myc.Ast,
    Myc.Ast.Nodes,
    Myc.Ast.Scope,
    Myc.Ast.RTL,
    Myc.Ast.Types;

type
    IMacroRegistry = interface;
    IEnvironment = interface;
    IExecutionStrategy = interface;
    IFunctionDefinitionRegistry = interface;

    // --- Monomorphization Cache Definitions ---

    TMonoCacheKey = record
    public
        Address: TResolvedAddress;
        ArgTypes: TArray<IStaticType>;
        constructor Create(const AAddress: TResolvedAddress; const AArgTypes: TArray<IStaticType>);
    end;

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

    TMonomorphCache = TDictionary<TMonoCacheKey, TSpecializedMethod>;

    TCompiledFunction = record
    public
        Func: TDataValue.TFunc;
        StaticType: IStaticType;
        IsPure: Boolean;
        constructor Create(const AFunc: TDataValue.TFunc; const AStaticType: IStaticType; AIsPure: Boolean);
    end;

    IExecutionStrategy = interface
        function CreateVisitor(const AScope: IExecutionScope): IEvaluatorVisitor;
    end;

    IMacroRegistry = interface
        {$region 'private'}
        function GetParent: IMacroRegistry;
        {$endregion}
        procedure Define(const Node: IMacroDefinitionNode);
        function Find(const Name: string): IMacroDefinitionNode;
        function CreateChildRegistry: IMacroRegistry;
        property Parent: IMacroRegistry read GetParent;
    end;

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

    IEnvironment = interface
        {$region 'private'}
        function GetRootScope: IExecutionScope;
        function GetMacroRegistry: IMacroRegistry;
        function GetMonomorphCache: TMonomorphCache;
        function GetFunctionRegistry: IFunctionDefinitionRegistry;
        {$endregion}

        procedure SetExecutionStrategy(const AStrategy: IExecutionStrategy);

        function ExpandMacros(const Node: IAstNode): IAstNode;

        // Updated: Returns Layout, not Descriptor. Performs Binding AND TypeChecking.
        function Bind(const Node: IAstNode; out Layout: IScopeLayout; const AArgTypes: TArray<IStaticType> = []): IAstNode;

        // Updated: Removed Descriptor parameter.
        function Specialize(const Node: IAstNode): IAstNode;

        function Compile(
            const ANode: IAstNode;
            const Params: TArray<IIdentifierNode> = [];
            const AArgTypes: TArray<IStaticType> = []
        ): TCompiledFunction; overload;

        function Compile(const Node: IFunctionDefinition; const ArgTypes: TArray<IStaticType> = []): TCompiledFunction; overload;

        function CreateEnvironment: IEnvironment;

        property RootScope: IExecutionScope read GetRootScope;
        property MacroRegistry: IMacroRegistry read GetMacroRegistry;
        property MonomorphCache: TMonomorphCache read GetMonomorphCache;
        property FunctionRegistry: IFunctionDefinitionRegistry read GetFunctionRegistry;
    end;

    // Interface Helper for IEnvironment
    TAstEnvironment = record
    private
        FEnvironment: IEnvironment;
        function GetRootScope: IExecutionScope; inline;
        function GetMacroRegistry: IMacroRegistry; inline;
    public
        constructor Create(const AEnvironment: IEnvironment);

        class function Construct(const Scope: IExecutionScope): TAstEnvironment; static;

        class operator Implicit(const A: IEnvironment): TAstEnvironment;
        class operator Implicit(const A: TAstEnvironment): IEnvironment;

        function CreateEnvironment: TAstEnvironment;

        procedure SetStandardMode;
        procedure SetDebugMode(ALog: TStrings; AShowScope: Boolean);

        function Run(const ANode: IAstNode; const Params: TArray<IIdentifierNode> = []; const Args: TArray<TDataValue> = []): TDataValue;

        function Compile(
            const Node: IAstNode;
            const Params: TArray<IIdentifierNode> = [];
            const ArgTypes: TArray<IStaticType> = []
        ): TCompiledFunction; overload;

        function Compile(const Node: IFunctionDefinition; const ArgTypes: TArray<IStaticType> = []): TCompiledFunction; overload;

        procedure Define(const Name: String; const AScript: IAstNode);

        property Environment: IEnvironment read FEnvironment;
        property RootScope: IExecutionScope read GetRootScope;
        property MacroRegistry: IMacroRegistry read GetMacroRegistry;
    end;

implementation

uses
    System.Hash,
    Myc.Ast.Evaluator,
    Myc.Ast.Debugger,
    Myc.Ast.Compiler.Macros,
    Myc.Ast.Compiler.Binder,
    Myc.Ast.Compiler.TypeChecker,
    Myc.Ast.Compiler.Specializer,
    Myc.Ast.Compiler.TCO,
    Myc.Ast.Analysis.Purity;

{ TMonoCacheKey }

constructor TMonoCacheKey.Create(const AAddress: TResolvedAddress; const AArgTypes: TArray<IStaticType>);
begin
    Address := AAddress;
    ArgTypes := AArgTypes;
end;

{ TMonoCacheKeyComparer }

function TMonoCacheKeyComparer.Equals(const Left, Right: TMonoCacheKey): Boolean;
var
    i: Integer;
begin
    if not (Left.Address = Right.Address) then
        exit(False);

    if Length(Left.ArgTypes) <> Length(Right.ArgTypes) then
        exit(False);

    for i := 0 to High(Left.ArgTypes) do
    begin
        if not Left.ArgTypes[i].IsEqual(Right.ArgTypes[i]) then
            exit(False);
    end;

    Result := True;
end;

function TMonoCacheKeyComparer.GetHashCode(const Value: TMonoCacheKey): Integer;
var
    i: Integer;
    hash: Integer;
    typeHash: Integer;
    adr: TResolvedAddress;
begin
    adr := Value.Address;

    hash := THashBobJenkins.GetHashValue(adr.Kind, SizeOf(TAddressKind), 0);
    hash := THashBobJenkins.GetHashValue(adr.ScopeDepth, SizeOf(Integer), hash);
    hash := THashBobJenkins.GetHashValue(adr.SlotIndex, SizeOf(Integer), hash);

    for i := 0 to High(Value.ArgTypes) do
    begin
        if Assigned(Value.ArgTypes[i]) then
            typeHash := Value.ArgTypes[i].GetHashCode
        else
            typeHash := 0;

        hash := THashBobJenkins.GetHashValue(typeHash, SizeOf(Integer), hash);
    end;

    Result := hash;
end;

{ TCompiledFunction }

constructor TCompiledFunction.Create(const AFunc: TDataValue.TFunc; const AStaticType: IStaticType; AIsPure: Boolean);
begin
    Func := AFunc;
    StaticType := AStaticType;
    IsPure := AIsPure;
end;

type
    { TStandardExecutionStrategy }
    TStandardExecutionStrategy = class(TInterfacedObject, IExecutionStrategy)
    public
        function CreateVisitor(const AScope: IExecutionScope): IEvaluatorVisitor;
    end;

    { TDebugExecutionStrategy }
    TDebugExecutionStrategy = class(TInterfacedObject, IExecutionStrategy)
    private
        FLog: TStrings;
        FShowScope: Boolean;
    public
        constructor Create(ALog: TStrings; AShowScope: Boolean);
        function CreateVisitor(const AScope: IExecutionScope): IEvaluatorVisitor;
    end;

    TFunctionDefinitionRegistry = class(TInterfacedObject, IFunctionDefinitionRegistry)
    private
        FMap: TDictionary<TResolvedAddress, IFunctionDefinition>;
    public
        constructor Create;
        destructor Destroy; override;
        procedure Register(const Address: TResolvedAddress; const ADef: IFunctionDefinition);
        function Resolve(const Address: TResolvedAddress): IFunctionDefinition;
    end;

    { TEnvironment }
    TEnvironment = class(TInterfacedObject, IEnvironment)
    private
        FRootScope: IExecutionScope;
        FMacroRegistry: IMacroRegistry;
        FExecutionStrategy: IExecutionStrategy;
        FMonomorphCache: TMonomorphCache;
        FFunctionRegistry: IFunctionDefinitionRegistry;
        function GetRootScope: IExecutionScope;
        function GetMacroRegistry: IMacroRegistry;
        function GetMonomorphCache: TMonomorphCache;
        function GetFunctionRegistry: IFunctionDefinitionRegistry;
    public
        constructor Create(
            const ARootScope: IExecutionScope;
            const AMacroRegistry: IMacroRegistry;
            const AExecutionStrategy: IExecutionStrategy
        );
        destructor Destroy; override;

        function CreateEnvironment: IEnvironment;

        procedure SetExecutionStrategy(const AStrategy: IExecutionStrategy);

        function ExpandMacros(const Node: IAstNode): IAstNode;
        function Bind(const Node: IAstNode; out Layout: IScopeLayout; const ArgTypes: TArray<IStaticType>): IAstNode;
        function Specialize(const Node: IAstNode): IAstNode;

        function Compile(
            const Node: IAstNode;
            const Params: TArray<IIdentifierNode>;
            const ArgTypes: TArray<IStaticType>
        ): TCompiledFunction; overload;
        function Compile(const Node: IFunctionDefinition; const ArgTypes: TArray<IStaticType>): TCompiledFunction; overload;
    end;

{ TAstEnvironment }

constructor TAstEnvironment.Create(const AEnvironment: IEnvironment);
begin
    FEnvironment := AEnvironment;
end;

procedure TAstEnvironment.SetStandardMode;
begin
    FEnvironment.SetExecutionStrategy(TStandardExecutionStrategy.Create);
end;

procedure TAstEnvironment.SetDebugMode(ALog: TStrings; AShowScope: Boolean);
begin
    FEnvironment.SetExecutionStrategy(TDebugExecutionStrategy.Create(ALog, AShowScope));
end;

class operator TAstEnvironment.Implicit(const A: IEnvironment): TAstEnvironment;
begin
    Result.FEnvironment := A;
end;

class operator TAstEnvironment.Implicit(const A: TAstEnvironment): IEnvironment;
begin
    Result := A.FEnvironment;
end;

function TAstEnvironment.Compile(
    const Node: IAstNode;
    const Params: TArray<IIdentifierNode> = [];
    const ArgTypes: TArray<IStaticType> = []
): TCompiledFunction;
begin
    Result := FEnvironment.Compile(Node, Params, ArgTypes);
end;

function TAstEnvironment.Compile(const Node: IFunctionDefinition; const ArgTypes: TArray<IStaticType> = []): TCompiledFunction;
begin
    Result := FEnvironment.Compile(Node, ArgTypes);
end;

class function TAstEnvironment.Construct(const Scope: IExecutionScope): TAstEnvironment;
var
    RootScope: IExecutionScope;
begin
    // Initialize root scope with library registration
    RootScope := TAst.CreateScope(nil, nil, True);

    Result.Create(TEnvironment.Create(RootScope, TMacroRegistry.Create(nil), TStandardExecutionStrategy.Create));
end;

function TAstEnvironment.CreateEnvironment: TAstEnvironment;
begin
    Result := FEnvironment.CreateEnvironment;
end;

procedure TAstEnvironment.Define(const Name: String; const AScript: IAstNode);
begin
    var compiled := Compile(AScript);
    // Note: Define now potentially handles Pure flags internally via the Node metadata,
    // but the runtime value is just the TDataValue.
    RootScope.Define(Name, compiled.Func([]), compiled.StaticType);
end;

function TAstEnvironment.GetRootScope: IExecutionScope;
begin
    Result := FEnvironment.GetRootScope;
end;

function TAstEnvironment.GetMacroRegistry: IMacroRegistry;
begin
    Result := FEnvironment.GetMacroRegistry;
end;

function TAstEnvironment.Run(
    const ANode: IAstNode;
    const Params: TArray<IIdentifierNode> = [];
    const Args: TArray<TDataValue> = []
): TDataValue;
begin
    var compiled := Compile(ANode, Params);
    Result := compiled.Func(Args);
end;

{ TStandardExecutionStrategy }

function TStandardExecutionStrategy.CreateVisitor(const AScope: IExecutionScope): IEvaluatorVisitor;
begin
    Result := TEvaluatorVisitor.Create(AScope);
end;

{ TDebugExecutionStrategy }

constructor TDebugExecutionStrategy.Create(ALog: TStrings; AShowScope: Boolean);
begin
    inherited Create;
    FLog := ALog;
    FShowScope := AShowScope;
end;

function TDebugExecutionStrategy.CreateVisitor(const AScope: IExecutionScope): IEvaluatorVisitor;
begin
    Result := TDebugEvaluatorVisitor.Create(AScope, FLog, FShowScope, 0);
end;

{ TFunctionDefinitionRegistry }

constructor TFunctionDefinitionRegistry.Create;
begin
    inherited Create;
    FMap := TDictionary<TResolvedAddress, IFunctionDefinition>.Create(TResolvedAddressComparer.Create);
end;

destructor TFunctionDefinitionRegistry.Destroy;
begin
    FMap.Free;
    inherited Destroy;
end;

procedure TFunctionDefinitionRegistry.Register(const Address: TResolvedAddress; const ADef: IFunctionDefinition);
begin
    FMap.AddOrSetValue(Address, ADef);
end;

function TFunctionDefinitionRegistry.Resolve(const Address: TResolvedAddress): IFunctionDefinition;
begin
    FMap.TryGetValue(Address, Result);
end;

{ TEnvironment }

constructor TEnvironment.Create(
    const ARootScope: IExecutionScope;
    const AMacroRegistry: IMacroRegistry;
    const AExecutionStrategy: IExecutionStrategy
);
begin
    inherited Create;
    FRootScope := ARootScope;
    FMacroRegistry := AMacroRegistry;
    FExecutionStrategy := AExecutionStrategy;
    FMonomorphCache := TMonomorphCache.Create(TMonoCacheKeyComparer.Create);
    FFunctionRegistry := TFunctionDefinitionRegistry.Create;
end;

destructor TEnvironment.Destroy;
begin
    FMonomorphCache.Free;
    inherited Destroy;
end;

function TEnvironment.GetMonomorphCache: TMonomorphCache;
begin
    Result := FMonomorphCache;
end;

function TEnvironment.Bind(const Node: IAstNode; out Layout: IScopeLayout; const ArgTypes: TArray<IStaticType>): IAstNode;
begin
    // 1. Bind (produces Layout and Bound AST with Addresses)
    var boundAst := TAstBinder.Bind(FRootScope.Descriptor.Layout, Node, Layout, FFunctionRegistry, ArgTypes);

    // 2. Check Types (produces Typed AST with Descriptors baked into nodes)
    Result := TTypeChecker.CheckTypes(boundAst, Layout);
end;

function TEnvironment.GetMacroRegistry: IMacroRegistry;
begin
    Result := FMacroRegistry;
end;

function TEnvironment.GetRootScope: IExecutionScope;
begin
    Result := FRootScope;
end;

procedure TEnvironment.SetExecutionStrategy(const AStrategy: IExecutionStrategy);
begin
    FExecutionStrategy := AStrategy;
end;

function TEnvironment.CreateEnvironment: IEnvironment;
begin
    Result := TEnvironment.Create(TAst.CreateScope(FRootScope), TMacroRegistry.Create(FMacroRegistry), FExecutionStrategy);
end;

function TEnvironment.Compile(
    const Node: IAstNode;
    const Params: TArray<IIdentifierNode>;
    const ArgTypes: TArray<IStaticType>
): TCompiledFunction;
var
    layout: IScopeLayout;
    descriptor: IScopeDescriptor;
    funcType: IStaticType;
    finalFunc: TDataValue.TFunc;

    typedNode: IAstNode;
    specialized: IAstNode;
    tcoOptimized: IAstNode;
    isPure: Boolean;
begin
    // 0. Wrap in Lambda to create an isolatable compilation unit (scope)
    var prg := TAst.LambdaExpr(Params, Node);

    // 1. Expand Macros
    var expanded := ExpandMacros(prg);

    // 2. Bind & Check Types
    typedNode := Bind(expanded, layout, ArgTypes);

    // Retrieve the descriptor from the typed lambda node
    // (The TypeChecker baked it into the new LambdaExpressionNode)
    descriptor := typedNode.AsLambdaExpression.Descriptor;
    Assert(Assigned(descriptor), 'Compiler Logic Error: Descriptor missing after type checking.');

    funcType := typedNode.AsTypedNode.StaticType;

    // 3. Specialize
    specialized := Specialize(typedNode);

    // 4. Optimize (TCO)
    tcoOptimized := TAstTCO.Optimize(specialized);

    // 5. Purity Inference
    // Check the body of the optimized lambda for purity.
    // Note: tcoOptimized is a LambdaExpressionNode. Its body is what we check.
    isPure := TPurityAnalyzer.IsPure(tcoOptimized.AsLambdaExpression.Body);

    // 6. Create Evaluator
    var visitor := FExecutionStrategy.CreateVisitor(descriptor.CreateScope(FRootScope));

    // Execute the AST to get the TFunc (Closure)
    var closure := tcoOptimized.Accept(visitor).AsMethod();

    // Wrap TCO handling for the final result
    finalFunc :=
        function(const Args: TArray<TDataValue>): TDataValue
        begin
            Result := closure(Args);
            TEvaluatorVisitor.HandleTCO(Result);
        end;

    // Pass IsPure flag to the compiled function record
    Result := TCompiledFunction.Create(finalFunc, funcType, isPure);
end;

function TEnvironment.Compile(const Node: IFunctionDefinition; const ArgTypes: TArray<IStaticType>): TCompiledFunction;
var
    layout: IScopeLayout;
    descriptor: IScopeDescriptor;
    funcType: IStaticType;
    finalFunc: TDataValue.TFunc;

    typedNode, specialized, tcoOptimized: IAstNode;
    isPure: Boolean;
begin
    // Same steps as above, but Node is already a definition (Lambda)
    var expanded := ExpandMacros(Node);
    typedNode := Bind(expanded, layout, ArgTypes);

    descriptor := typedNode.AsLambdaExpression.Descriptor;
    Assert(Assigned(descriptor));
    funcType := typedNode.AsTypedNode.StaticType;

    specialized := Specialize(typedNode);
    tcoOptimized := TAstTCO.Optimize(specialized);

    // Purity Inference
    isPure := TPurityAnalyzer.IsPure(tcoOptimized.AsLambdaExpression.Body);

    var visitor := FExecutionStrategy.CreateVisitor(descriptor.CreateScope(FRootScope));
    var closure := tcoOptimized.Accept(visitor);

    finalFunc :=
        function(const Args: TArray<TDataValue>): TDataValue
        begin
            Result := closure.AsMethod()(Args);
            TEvaluatorVisitor.HandleTCO(Result);
        end;

    Result := TCompiledFunction.Create(finalFunc, funcType, isPure);
end;

function TEnvironment.ExpandMacros(const Node: IAstNode): IAstNode;
begin
    var cExecutionStrategy := FExecutionStrategy;
    Result :=
        TMacroExpander.ExpandMacros(
            FMacroRegistry,
            FRootScope,
            Node,
            function(const Scope: IExecutionScope): IEvaluatorVisitor begin Result := cExecutionStrategy.CreateVisitor(Scope); end
        );
end;

function TEnvironment.GetFunctionRegistry: IFunctionDefinitionRegistry;
begin
    Result := FFunctionRegistry;
end;

function TEnvironment.Specialize(const Node: IAstNode): IAstNode;
begin
    Result := TStaticSpecializer.Specialize(Self, Node);
end;

end.