-module(extreme_stress_test).
-export([run_extreme_tests/0]).

run_extreme_tests() ->
    io:format("=== EXTREME MLX STRESS TESTS ===~n~n"),
    io:format("⚠️  WARNING: These tests push MLX to its absolute limits ⚠️~n~n"),
    
    % Setup
    code:add_path("mlx/_build/default/lib/mlx/ebin"),
    
    ExtremeSuites = [
        {"Massive Arrays", fun test_massive_arrays/0},
        {"Ultra-High Frequency", fun test_ultra_high_frequency/0},
        {"Memory Saturation", fun test_memory_saturation/0},
        {"Computational Intensity", fun test_computational_intensity/0},
        {"Dimensional Extremes", fun test_dimensional_extremes/0},
        {"Endurance Testing", fun test_endurance/0},
        {"Pathological Cases", fun test_pathological_cases/0},
        {"Resource Exhaustion", fun test_resource_exhaustion/0}
    ],
    
    Results = run_extreme_suites(ExtremeSuites),
    
    % Analysis
    TotalSuites = length(ExtremeSuites),
    PassedSuites = length([Result || {_, pass} = Result <- Results]),
    
    io:format("~n🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥~n"),
    io:format("                    EXTREME STRESS RESULTS                      ~n"),
    io:format("🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥~n"),
    io:format("Total Extreme Tests: ~p~n", [TotalSuites]),
    io:format("Survived: ~p~n", [PassedSuites]),
    io:format("Survival Rate: ~.1f%~n", [PassedSuites * 100.0 / TotalSuites]),
    
    % Show results
    lists:foreach(fun({Suite, Status}) ->
        case Status of
            pass -> io:format("💪 ~s: SURVIVED~n", [Suite]);
            {fail, _} -> io:format("💥 ~s: FAILED (as expected under extreme stress)~n", [Suite]);
            {partial, Details} -> io:format("⚡ ~s: PARTIAL (~p)~n", [Suite, Details])
        end
    end, Results),
    
    case PassedSuites of
        TotalSuites ->
            io:format("~n🏆 UNBREAKABLE! MLX SURVIVED ALL EXTREME TESTS! 🏆~n");
        _ when PassedSuites >= TotalSuites * 0.7 ->
            io:format("~n💪 ROBUST! MLX handled most extreme conditions! 💪~n");
        _ ->
            io:format("~n⚡ TESTED TO LIMITS! Some extreme conditions exceeded capacity ⚡~n")
    end,
    
    io:format("🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥~n~n").

run_extreme_suites(ExtremeSuites) ->
    lists:map(fun({SuiteName, TestFun}) ->
        io:format("💥 EXTREME TEST: ~s...~n", [SuiteName]),
        StartTime = erlang:monotonic_time(millisecond),
        
        try
            TestFun(),
            EndTime = erlang:monotonic_time(millisecond),
            Duration = EndTime - StartTime,
            io:format("  💪 ~s: SURVIVED (~pms)~n", [SuiteName, Duration]),
            {SuiteName, pass}
        catch
            ExtremeClass:ExtremeReason:_Stack ->
                ExtremeEndTime = erlang:monotonic_time(millisecond),
                ExtremeDuration = ExtremeEndTime - StartTime,
                io:format("  💥 ~s: FAILED after ~pms - ~p:~p~n", [SuiteName, ExtremeDuration, ExtremeClass, ExtremeReason]),
                {SuiteName, {fail, {ExtremeClass, ExtremeReason}}}
        end
    end, ExtremeSuites).

test_massive_arrays() ->
    io:format("    Creating arrays of massive proportions...~n"),
    
    % Test progressively larger arrays
    MassiveSizes = [2000, 3000, 4000, 5000],
    
    lists:foreach(fun(Size) ->
        io:format("      Attempting ~px~p array...~n", [Size, Size]),
        StartMemory = erlang:memory(total),
        
        {ok, Massive} = mlx_nif:zeros([Size, Size], float32),
        ok = mlx_nif:eval(Massive),
        
        EndMemory = erlang:memory(total),
        MemoryUsed = (EndMemory - StartMemory) / (1024 * 1024), % MB
        
        io:format("      ✓ ~px~p array: ~.1f MB used~n", [Size, Size, MemoryUsed])
    end, MassiveSizes),
    
    % Test extremely large 1D array
    io:format("      Testing massive 1D array...~n"),
    {ok, MegaArray} = mlx_nif:zeros([10000000], float32), % 10 million elements
    ok = mlx_nif:eval(MegaArray),
    io:format("      ✓ 10M element 1D array successful~n"),
    
    ok.

test_ultra_high_frequency() ->
    io:format("    Testing ultra-high frequency operations...~n"),
    
    % Extremely rapid operations
    NumOperations = 10000,
    io:format("      Performing ~p rapid operations...~n", [NumOperations]),
    
    StartTime = erlang:monotonic_time(microsecond),
    
    lists:foreach(fun(_) ->
        {ok, A} = mlx_nif:ones([5, 5], float32),
        {ok, B} = mlx_nif:ones([5, 5], float32),
        {ok, C} = mlx_nif:add(A, B),
        ok = mlx_nif:eval(C)
    end, lists:seq(1, NumOperations)),
    
    EndTime = erlang:monotonic_time(microsecond),
    TotalTime = (EndTime - StartTime) / 1000,
    OpsPerSecond = NumOperations / (TotalTime / 1000),
    
    io:format("      ✓ ~p operations in ~.1fms (~.0f ops/sec)~n", 
              [NumOperations, TotalTime, OpsPerSecond]),
    
    % Test with no delays
    io:format("      Testing burst operations...~n"),
    BurstSize = 1000,
    StartBurst = erlang:monotonic_time(microsecond),
    
    % Create all arrays first
    Arrays = lists:map(fun(_) ->
        {ok, A} = mlx_nif:ones([10, 10], float32),
        A
    end, lists:seq(1, BurstSize)),
    
    % Operate on all at once
    lists:foreach(fun(A) ->
        ok = mlx_nif:eval(A)
    end, Arrays),
    
    EndBurst = erlang:monotonic_time(microsecond),
    BurstTime = (EndBurst - StartBurst) / 1000,
    
    io:format("      ✓ Burst: ~p arrays in ~.1fms~n", [BurstSize, BurstTime]),
    
    ok.

test_memory_saturation() ->
    io:format("    Testing memory saturation scenarios...~n"),
    
    % Create arrays until near memory limit
    ArraySize = 500,
    Arrays = create_arrays_until_limit(ArraySize, []),
    NumArrays = length(Arrays),
    
    TotalElements = NumArrays * ArraySize * ArraySize,
    MemoryUsed = TotalElements * 4 / (1024 * 1024), % 4 bytes per float32, convert to MB
    
    io:format("      ✓ Created ~p arrays (~px~p each)~n", [NumArrays, ArraySize, ArraySize]),
    io:format("      ✓ Total memory: ~.1f MB~n", [MemoryUsed]),
    
    % Test operations on all arrays
    io:format("      Testing operations under memory pressure...~n"),
    lists:foreach(fun(A) ->
        ok = mlx_nif:eval(A)
    end, lists:sublist(Arrays, 100)), % Test first 100 to avoid timeout
    
    io:format("      ✓ Operations successful under memory pressure~n"),
    
    ok.

create_arrays_until_limit(Size, Acc) ->
    try
        {ok, Array} = mlx_nif:ones([Size, Size], float32),
        ok = mlx_nif:eval(Array),
        create_arrays_until_limit(Size, [Array | Acc])
    catch
        _:_ ->
            % Hit limit, return what we have
            Acc
    end.

test_computational_intensity() ->
    io:format("    Testing computational intensity...~n"),
    
    % Chain many matrix operations
    io:format("      Testing operation chains...~n"),
    {ok, Start} = mlx_nif:ones([100, 100], float32),
    
    Result = lists:foldl(fun(_, Acc) ->
        {ok, B} = mlx_nif:ones([100, 100], float32),
        {ok, Sum} = mlx_nif:add(Acc, B),
        {ok, Prod} = mlx_nif:multiply(Sum, B),
        ok = mlx_nif:eval(Prod),
        Prod
    end, Start, lists:seq(1, 100)),
    
    ok = mlx_nif:eval(Result),
    io:format("      ✓ 100-step operation chain completed~n"),
    
    % Large matrix multiplication chains
    io:format("      Testing matrix multiplication chains...~n"),
    MatrixSize = 200,
    {ok, Matrix1} = mlx_nif:ones([MatrixSize, MatrixSize], float32),
    
    FinalMatrix = lists:foldl(fun(_, Acc) ->
        {ok, Next} = mlx_nif:ones([MatrixSize, MatrixSize], float32),
        {ok, Product} = mlx_nif:matmul(Acc, Next),
        ok = mlx_nif:eval(Product),
        Product
    end, Matrix1, lists:seq(1, 10)),
    
    ok = mlx_nif:eval(FinalMatrix),
    io:format("      ✓ 10-step ~px~p matrix multiplication chain~n", [MatrixSize, MatrixSize]),
    
    ok.

test_dimensional_extremes() ->
    io:format("    Testing dimensional extremes...~n"),
    
    % Test very high dimensional arrays
    HighDimShapes = [
        [2, 2, 2, 2, 2, 2, 2],           % 7D
        [1, 1, 1, 1, 1, 1, 1, 1],       % 8D
        [2, 1, 2, 1, 2, 1, 2, 1, 2],    % 9D
        [1, 2, 1, 2, 1, 2, 1, 2, 1, 2] % 10D
    ],
    
    lists:foreach(fun(Shape) ->
        Dimensions = length(Shape),
        try
            {ok, Array} = mlx_nif:zeros(Shape, float32),
            {ok, Shape} = mlx_nif:shape(Array),
            ok = mlx_nif:eval(Array),
            io:format("      ✓ ~pD array: ~p~n", [Dimensions, Shape])
        catch
            _:_ ->
            io:format("      ⚠ ~pD array: Not supported~n", [Dimensions])
        end
    end, HighDimShapes),
    
    % Test extremely long 1D arrays
    LongSizes = [1000000, 5000000, 10000000],
    lists:foreach(fun(Size) ->
        try
            {ok, LongArray} = mlx_nif:zeros([Size], float32),
            ok = mlx_nif:eval(LongArray),
            io:format("      ✓ 1D array: ~p elements~n", [Size])
        catch
            _:_ ->
            io:format("      ⚠ 1D array: ~p elements failed~n", [Size])
        end
    end, LongSizes),
    
    ok.

test_endurance() ->
    io:format("    Testing endurance over extended periods...~n"),
    
    % Long-running test
    Duration = 30, % seconds
    StartTime = erlang:monotonic_time(second),
    EndTime = StartTime + Duration,
    
    io:format("      Running continuous operations for ~p seconds...~n", [Duration]),
    
    OperationCount = endurance_loop(EndTime, 0),
    
    ActualDuration = erlang:monotonic_time(second) - StartTime,
    OpsPerSecond = OperationCount / ActualDuration,
    
    io:format("      ✓ Endurance: ~p operations in ~p seconds (~.1f ops/sec)~n", 
              [OperationCount, ActualDuration, OpsPerSecond]),
    
    % Memory stability test
    io:format("      Testing memory stability...~n"),
    InitialMemory = erlang:memory(total),
    
    lists:foreach(fun(_) ->
        _Arrays = lists:map(fun(_) ->
            {ok, A} = mlx_nif:ones([50, 50], float32),
            ok = mlx_nif:eval(A),
            A
        end, lists:seq(1, 100)),
        erlang:garbage_collect()
    end, lists:seq(1, 50)),
    
    FinalMemory = erlang:memory(total),
    MemoryIncrease = (FinalMemory - InitialMemory) / (1024 * 1024),
    
    io:format("      ✓ Memory stability: ~.1f MB increase after 5000 arrays~n", [MemoryIncrease]),
    
    ok.

endurance_loop(EndTime, Count) ->
    case erlang:monotonic_time(second) < EndTime of
        true ->
            % Perform operation
            {ok, A} = mlx_nif:ones([20, 20], float32),
            {ok, B} = mlx_nif:ones([20, 20], float32),
            {ok, C} = mlx_nif:add(A, B),
            ok = mlx_nif:eval(C),
            endurance_loop(EndTime, Count + 1);
        false ->
            Count
    end.

test_pathological_cases() ->
    io:format("    Testing pathological edge cases...~n"),
    
    % Test with very unbalanced matrices
    UnbalancedShapes = [
        {[1, 1000], [1000, 1]},
        {[10, 1000], [1000, 10]},
        {[1, 10000], [10000, 1]}
    ],
    
    lists:foreach(fun({ShapeA, ShapeB}) ->
        try
            {ok, A} = mlx_nif:ones(ShapeA, float32),
            {ok, B} = mlx_nif:ones(ShapeB, float32),
            {ok, C} = mlx_nif:matmul(A, B),
            ok = mlx_nif:eval(C),
            io:format("      ✓ Unbalanced matmul: ~p × ~p~n", [ShapeA, ShapeB])
        catch
            _:_ ->
            io:format("      ⚠ Unbalanced matmul failed: ~p × ~p~n", [ShapeA, ShapeB])
        end
    end, UnbalancedShapes),
    
    % Test rapid device switching
    io:format("      Testing rapid device switching...~n"),
    lists:foreach(fun(I) ->
        Device = case I rem 2 of
            0 -> cpu;
            1 -> gpu
        end,
        case mlx_nif:set_default_device(Device) of
            ok ->
                {ok, A} = mlx_nif:ones([10, 10], float32),
                ok = mlx_nif:eval(A);
            {error, _} ->
                ok % GPU might not be available
        end
    end, lists:seq(1, 100)),
    
    io:format("      ✓ Rapid device switching: 100 switches~n"),
    
    ok.

test_resource_exhaustion() ->
    io:format("    Testing resource exhaustion scenarios...~n"),
    
    % Try to exhaust array handles
    io:format("      Testing array handle exhaustion...~n"),
    
    try
        LargeArrayList = create_many_arrays(1000, []),
        NumCreated = length(LargeArrayList),
        io:format("      ✓ Created ~p arrays before exhaustion~n", [NumCreated]),
        
        % Try to use some of them
        lists:foreach(fun(A) ->
            ok = mlx_nif:eval(A)
        end, lists:sublist(LargeArrayList, 100)),
        
        io:format("      ✓ Successfully operated on arrays near exhaustion~n")
    catch
        _:_ ->
        io:format("      ⚠ Resource exhaustion handled gracefully~n")
    end,
    
    % Force cleanup
    erlang:garbage_collect(),
    io:format("      ✓ Resource cleanup after exhaustion test~n"),
    
    ok.

create_many_arrays(0, Acc) ->
    Acc;
create_many_arrays(N, Acc) ->
    try
        {ok, Array} = mlx_nif:zeros([10, 10], float32),
        create_many_arrays(N - 1, [Array | Acc])
    catch
        _:_ ->
            Acc
    end.