-module(albedo_openai_bench).
-export([env/1, measure/1, text/1]).

env(Name) ->
    case os:getenv(binary_to_list(Name)) of
        false -> <<>>;
        Value -> unicode:characters_to_binary(Value)
    end.

text(Bytes) when Bytes > 0 ->
    case get(first_text) of undefined -> put(first_text, erlang:monotonic_time(microsecond)); _ -> ok end,
    put(text_bytes, get(text_bytes) + Bytes),
    put(text_chunks, get(text_chunks) + 1),
    nil;
text(_) -> nil.

measure(Run) ->
    {ok, _} = application:ensure_all_started(gun),
    Owner = self(),
    Baseline = erlang:memory(binary),
    {Pid, Ref} = spawn_monitor(fun() ->
        put(text_bytes, 0), put(text_chunks, 0),
        Start = erlang:monotonic_time(microsecond),
        Result = Run(),
        End = erlang:monotonic_time(microsecond),
        First = case get(first_text) of undefined -> none; Time -> {some, (Time - Start) / 1000} end,
        {reductions, Reductions} = process_info(self(), reductions),
        {memory, Memory} = process_info(self(), memory),
        Owner ! {self(), measured, Result, (End - Start) / 1000, First,
            get(text_bytes), get(text_chunks), Memory, Reductions},
        receive release -> ok end
    end),
    sample(Pid, Ref, Baseline, 0, Baseline).

sample(Pid, Ref, Baseline, PeakProcess, PeakBinary) ->
    receive
        {Pid, measured, Result, Elapsed, First, Bytes, Chunks, Memory, Reductions} ->
            Binary = erlang:memory(binary),
            Pid ! release,
            demonitor(Ref, [flush]),
            {Result, {metrics, Elapsed, First, Bytes, Chunks,
                max(PeakProcess, Memory), max(0, max(PeakBinary, Binary) - Baseline), Reductions}};
        {'DOWN', Ref, process, Pid, Reason} -> error({benchmark_worker_crashed, Reason})
    after 10 ->
        Memory = case process_info(Pid, memory) of {memory, M} -> M; undefined -> 0 end,
        sample(Pid, Ref, Baseline, max(PeakProcess, Memory), max(PeakBinary, erlang:memory(binary)))
    end.