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,
    Myc.Ast.Compiler.Macros,
    Myc.Ast.Compiler.Binder,
    Myc.Ast.Compiler.TypeChecker,
    Myc.Ast.Compiler.Specializer,
    Myc.Ast.Compiler.TCO,
    Myc.Ast.Analysis.Purity;

type
    IEnvironment = interface;
    IExecutionStrategy = interface;

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

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

        procedure SetExecutionStrategy(const AStrategy: IExecutionStrategy);

        function ExpandMacros(const Node: IAstNode): IAstNode;

        function Bind(
            const Node: IAstNode;
            out Layout: IScopeLayout;
            const AArgTypes: TArray<IStaticType>;
            const Log: ICompilerLog
        ): IAstNode;

        function Specialize(const Node: IAstNode): IAstNode;

        function Compile(
            const Node: IFunctionDefinition;
            const ArgTypes: TArray<IStaticType> = [];
            const Log: ICompilerLog = nil
        ): ILambdaExpressionNode;
        function Link(const Node: ILambdaExpressionNode; const Log: ICompilerLog = nil): TCompiledFunction;

        function CreateEnvironment: IEnvironment;

        property RootScope: IExecutionScope read GetRootScope;

        property MacroRegistry: IMacroRegistry read GetMacroRegistry;
        property MonomorphCache: IMonomorphCache 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 ExpandMacros(const Node: IAstNode): IAstNode;
        function Bind(const Node: IAstNode; const AArgTypes: TArray<IStaticType>; const Log: ICompilerLog): IAstNode;
        function Specialize(const Node: IAstNode): IAstNode;

        function Compile(
            const Node: IAstNode;
            const Params: TArray<IIdentifierNode> = [];
            const ArgTypes: TArray<IStaticType> = []
        ): ILambdaExpressionNode; overload;
        function Compile(const Node: IFunctionDefinition; const ArgTypes: TArray<IStaticType> = []): ILambdaExpressionNode; overload;
        function Link(const Node: ILambdaExpressionNode; const Log: ICompilerLog = nil): TCompiledFunction;

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

        // UPDATED: Now supports Documentation
        procedure Define(const Name: String; const AScript: IAstNode; const Doc: string = '');

        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;

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

    TMonomorphCache = class(TInterfacedObject, IMonomorphCache)
    private
        FMonomorphCache: TDictionary<TMonoCacheKey, TSpecializedMethod>;
    public
        constructor Create;
        destructor Destroy; override;
        function TryGetFunction(const Key: TMonoCacheKey; out Func: TSpecializedMethod): Boolean;
        procedure Add(const Key: TMonoCacheKey; const Func: TSpecializedMethod);
    end;

    { 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: IMonomorphCache;
        FFunctionRegistry: IFunctionDefinitionRegistry;
        function GetRootScope: IExecutionScope;
        function GetMacroRegistry: IMacroRegistry;
        function GetMonomorphCache: IMonomorphCache;
        function GetFunctionRegistry: IFunctionDefinitionRegistry;
    public
        constructor Create(
            const ARootScope: IExecutionScope;
            const AMacroRegistry: IMacroRegistry;
            const AExecutionStrategy: IExecutionStrategy
        );
        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>;
            const Log: ICompilerLog
        ): IAstNode;

        function Specialize(const Node: IAstNode): IAstNode;

        function Compile(
            const Node: IFunctionDefinition;
            const ArgTypes: TArray<IStaticType>;
            const Log: ICompilerLog = nil
        ): ILambdaExpressionNode;
        function Link(const Node: ILambdaExpressionNode; const Log: ICompilerLog = nil): TCompiledFunction;
    end;

{ TAstEnvironment }

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

function TAstEnvironment.Bind(const Node: IAstNode; const AArgTypes: TArray<IStaticType>; const Log: ICompilerLog): IAstNode;
var
    layout: IScopeLayout;
begin
    Result := FEnvironment.Bind(Node, layout, AArgTypes, Log);
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> = []
): ILambdaExpressionNode;
begin
    Result := FEnvironment.Compile(TAst.LambdaExpr(Params, Node), ArgTypes);
end;

function TAstEnvironment.Compile(const Node: IFunctionDefinition; const ArgTypes: TArray<IStaticType> = []): ILambdaExpressionNode;
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 := Scope;
    if RootScope = nil then
        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; const Doc: string);
begin
    var compiled := Compile(AScript);
    var linked := Link(compiled);
    // Pass the documentation to the scope
    RootScope.Define(Name, linked.Func([]), compiled.StaticType, Doc);
end;

function TAstEnvironment.ExpandMacros(const Node: IAstNode): IAstNode;
begin
    Result := FEnvironment.ExpandMacros(Node);
end;

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

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

function TAstEnvironment.Link(const Node: ILambdaExpressionNode; const Log: ICompilerLog = nil): TCompiledFunction;
begin
    Result := FEnvironment.Link(Node, Log);
end;

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

function TAstEnvironment.Specialize(const Node: IAstNode): IAstNode;
begin
    Result := FEnvironment.Specialize(Node);
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;
    FFunctionRegistry := TFunctionDefinitionRegistry.Create;
end;

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

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

    // 2. Check Types (produces Typed AST with Descriptors baked into nodes)
    Result := TTypeChecker.CheckTypes(boundAst, Layout, FRootScope, Log);
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: IFunctionDefinition;
    const ArgTypes: TArray<IStaticType>;
    const Log: ICompilerLog = nil
): ILambdaExpressionNode;
var
    layout: IScopeLayout;
    lg: ICompilerLog;
begin
    lg :=
        if Assigned(Log) then Log
        else TCompilerLog.Create as ICompilerLog;

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

    // 2. Bind & TypeCheck (accumulating errors)
    var typedNode := Bind(expanded, layout, ArgTypes, lg);

    // 3. Check for compilation errors
    if lg.HasErrors then
        raise ECompilationFailed.Create(lg.GetEntries);

    var specialized := Specialize(typedNode);

    Result := specialized.AsLambdaExpression;
end;

function TEnvironment.Link(const Node: ILambdaExpressionNode; const Log: ICompilerLog = nil): TCompiledFunction;
var
    descriptor: IScopeDescriptor;
    funcType: IStaticType;
    finalFunc: TDataValue.TFunc;
    isPure: Boolean;
begin
    descriptor := Node.Descriptor;

    // Descriptor might be nil if Bind failed badly, but HasErrors check above covers that.
    Assert(Assigned(descriptor));

    funcType := Node.StaticType;

    var tcoOptimized := TAstTCO.Optimize(Node).AsLambdaExpression;

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

    // 6. Generate Visitor/Closure
    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;

    // The Glue: This callback performs the full compilation pipeline for a macro argument expression.
    var macroEvaluator: TMacroEvaluatorProc :=
        function(const Scope: IExecutionScope; const ANode: IAstNode): TDataValue
        var
            tmpLayout: IScopeLayout;
            evaluator: IEvaluatorVisitor;
            boundSubAst: IAstNode;
            scratchScope: IExecutionScope;
            tempLog: ICompilerLog;
        begin
            // Create temporary log for the macro context
            tempLog := TCompilerLog.Create;

            // 1. Binding
            boundSubAst := TAstBinder.Bind(Scope.Descriptor.Layout, ANode, tmpLayout, tempLog);

            if tempLog.HasErrors then
                raise EMacroException.Create('Macro Argument Error: ' + tempLog.GetEntries[0].Message);

            // 2. Scope Matching
            scratchScope := TScope.CreateScope(Scope, nil, nil);

            // 3. Execution
            evaluator := cExecutionStrategy.CreateVisitor(scratchScope);
            Result := evaluator.Execute(boundSubAst);
        end;

    Result := TMacroExpander.ExpandMacros(FMacroRegistry, FRootScope, Node, macroEvaluator);
end;

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

function TEnvironment.Specialize(const Node: IAstNode): IAstNode;
begin
    Result :=
        TStaticSpecializer.Specialize(
            Node,
            FMonomorphCache,
            FFunctionRegistry,
            function(const Node: IFunctionDefinition; const ArgTypes: TArray<IStaticType>): TCompiledFunction
            begin
                Result := Link(Compile(Node, ArgTypes));
            end
        );
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;

constructor TMonomorphCache.Create;
begin
    inherited Create;
    FMonomorphCache := TDictionary<TMonoCacheKey, TSpecializedMethod>.Create(TMonoCacheKeyComparer.Create);
end;

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

procedure TMonomorphCache.Add(const Key: TMonoCacheKey; const Func: TSpecializedMethod);
begin
    FMonomorphCache.Add(Key, Func);
end;

function TMonomorphCache.TryGetFunction(const Key: TMonoCacheKey; out Func: TSpecializedMethod): Boolean;
begin
    Result := FMonomorphCache.TryGetValue(Key, Func);
end;

end.