-module(glean@testing@simulation).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/glean/testing/simulation.gleam").
-export([new/2, with_system_prompt/2, with_max_steps/2, then_respond_text/2, then_respond_text_and_assert/3, then_call_tools/2, then_call_tools_and_assert/3, then_provider_error/2, assert_finally/2, assert_passed/1, run/1]).
-export_type([simulation_step/0, scripted_tool_execution/0, tool_execution_result/0, assert_fn/0, assertion_error/0, simulation_state/0, simulation/0, simulation_result/0]).

-if(?OTP_RELEASE >= 27).
-define(MODULEDOC(Str), -moduledoc(Str)).
-define(DOC(Str), -doc(Str)).
-else.
-define(MODULEDOC(Str), -compile([])).
-define(DOC(Str), -compile([])).
-endif.

-type simulation_step() :: {expect_text, binary(), assert_fn()} |
    {expect_tool_calls, list(scripted_tool_execution()), assert_fn()} |
    {expect_error, glean@error:glean_error(), assert_fn()}.

-type scripted_tool_execution() :: {scripted_tool_execution,
        glean@testing@provider:scripted_tool_call(),
        tool_execution_result()}.

-type tool_execution_result() :: {tool_success, binary()} |
    {tool_failure, binary()}.

-type assert_fn() :: no_assert |
    {assert,
        fun((simulation_state()) -> {ok, nil} | {error, assertion_error()})}.

-type assertion_error() :: {assertion_error, integer(), binary(), binary()}.

-type simulation_state() :: {simulation_state,
        list(glean@message:message()),
        integer(),
        glean@stream:usage(),
        list(glean@message:tool_result()),
        list(glean@error:glean_error())}.

-opaque simulation() :: {simulation,
        binary(),
        binary(),
        binary(),
        list(simulation_step()),
        integer(),
        list(fun((simulation_state()) -> {ok, nil} | {error, assertion_error()}))}.

-type simulation_result() :: {simulation_passed, binary(), simulation_state()} |
    {simulation_failed, binary(), integer(), assertion_error()} |
    {simulation_exceeded_max_steps, binary(), integer(), simulation_state()}.

-file("src/glean/testing/simulation.gleam", 89).
?DOC(" Create a new simulation.\n").
-spec new(binary(), binary()) -> simulation().
new(Name, Prompt) ->
    {simulation, Name, Prompt, <<""/utf8>>, [], 10, []}.

-file("src/glean/testing/simulation.gleam", 101).
?DOC(" Set the system prompt.\n").
-spec with_system_prompt(simulation(), binary()) -> simulation().
with_system_prompt(Sim, Prompt) ->
    {simulation,
        erlang:element(2, Sim),
        erlang:element(3, Sim),
        Prompt,
        erlang:element(5, Sim),
        erlang:element(6, Sim),
        erlang:element(7, Sim)}.

-file("src/glean/testing/simulation.gleam", 106).
?DOC(" Set max steps before the simulation is considered stuck.\n").
-spec with_max_steps(simulation(), integer()) -> simulation().
with_max_steps(Sim, Max) ->
    case Max > 0 of
        true -> nil;
        _assert_fail ->
            erlang:error(#{gleam_error => let_assert,
                        message => <<"Pattern match failed, no pattern matched the value."/utf8>>,
                        file => <<?FILEPATH/utf8>>,
                        module => <<"glean/testing/simulation"/utf8>>,
                        function => <<"with_max_steps"/utf8>>,
                        line => 107,
                        value => _assert_fail,
                        start => 2836,
                        'end' => 2861,
                        pattern_start => 2847,
                        pattern_end => 2851})
    end,
    {simulation,
        erlang:element(2, Sim),
        erlang:element(3, Sim),
        erlang:element(4, Sim),
        erlang:element(5, Sim),
        Max,
        erlang:element(7, Sim)}.

-file("src/glean/testing/simulation.gleam", 112).
?DOC(" Add a step where the provider returns text.\n").
-spec then_respond_text(simulation(), binary()) -> simulation().
then_respond_text(Sim, Text) ->
    Step = {expect_text, Text, no_assert},
    {simulation,
        erlang:element(2, Sim),
        erlang:element(3, Sim),
        erlang:element(4, Sim),
        lists:append(erlang:element(5, Sim), [Step]),
        erlang:element(6, Sim),
        erlang:element(7, Sim)}.

-file("src/glean/testing/simulation.gleam", 118).
?DOC(" Add a text step with assertion.\n").
-spec then_respond_text_and_assert(
    simulation(),
    binary(),
    fun((simulation_state()) -> {ok, nil} | {error, assertion_error()})
) -> simulation().
then_respond_text_and_assert(Sim, Text, Assertion) ->
    Step = {expect_text, Text, {assert, Assertion}},
    {simulation,
        erlang:element(2, Sim),
        erlang:element(3, Sim),
        erlang:element(4, Sim),
        lists:append(erlang:element(5, Sim), [Step]),
        erlang:element(6, Sim),
        erlang:element(7, Sim)}.

-file("src/glean/testing/simulation.gleam", 128).
?DOC(" Add a step where the provider calls tools.\n").
-spec then_call_tools(simulation(), list(scripted_tool_execution())) -> simulation().
then_call_tools(Sim, Executions) ->
    Step = {expect_tool_calls, Executions, no_assert},
    {simulation,
        erlang:element(2, Sim),
        erlang:element(3, Sim),
        erlang:element(4, Sim),
        lists:append(erlang:element(5, Sim), [Step]),
        erlang:element(6, Sim),
        erlang:element(7, Sim)}.

-file("src/glean/testing/simulation.gleam", 138).
?DOC(" Add a tool call step with assertion.\n").
-spec then_call_tools_and_assert(
    simulation(),
    list(scripted_tool_execution()),
    fun((simulation_state()) -> {ok, nil} | {error, assertion_error()})
) -> simulation().
then_call_tools_and_assert(Sim, Executions, Assertion) ->
    Step = {expect_tool_calls, Executions, {assert, Assertion}},
    {simulation,
        erlang:element(2, Sim),
        erlang:element(3, Sim),
        erlang:element(4, Sim),
        lists:append(erlang:element(5, Sim), [Step]),
        erlang:element(6, Sim),
        erlang:element(7, Sim)}.

-file("src/glean/testing/simulation.gleam", 149).
?DOC(" Add a step where the provider returns an error.\n").
-spec then_provider_error(simulation(), glean@error:glean_error()) -> simulation().
then_provider_error(Sim, Error) ->
    Step = {expect_error, Error, no_assert},
    {simulation,
        erlang:element(2, Sim),
        erlang:element(3, Sim),
        erlang:element(4, Sim),
        lists:append(erlang:element(5, Sim), [Step]),
        erlang:element(6, Sim),
        erlang:element(7, Sim)}.

-file("src/glean/testing/simulation.gleam", 155).
?DOC(" Add a final assertion that runs after all steps complete.\n").
-spec assert_finally(
    simulation(),
    fun((simulation_state()) -> {ok, nil} | {error, assertion_error()})
) -> simulation().
assert_finally(Sim, Assertion) ->
    {simulation,
        erlang:element(2, Sim),
        erlang:element(3, Sim),
        erlang:element(4, Sim),
        erlang:element(5, Sim),
        erlang:element(6, Sim),
        lists:append(erlang:element(7, Sim), [Assertion])}.

-file("src/glean/testing/simulation.gleam", 187).
?DOC(" Assert that a simulation passed. Panics with descriptive message if not.\n").
-spec assert_passed(simulation_result()) -> simulation_state().
assert_passed(Result) ->
    case Result of
        {simulation_passed, _, State} ->
            State;

        {simulation_failed, Name, Step, Err} ->
            erlang:error(#{gleam_error => panic,
                    message => erlang:list_to_binary(
                        [<<"Simulation '"/utf8>>,
                            Name,
                            <<"' failed at step "/utf8>>,
                            erlang:integer_to_binary(Step),
                            <<": expected "/utf8>>,
                            erlang:element(3, Err),
                            <<" but got "/utf8>>,
                            erlang:element(4, Err)]
                    ),
                    file => <<?FILEPATH/utf8>>,
                    module => <<"glean/testing/simulation"/utf8>>,
                    function => <<"assert_passed"/utf8>>,
                    line => 191});

        {simulation_exceeded_max_steps, Name@1, Max, _} ->
            erlang:error(#{gleam_error => panic,
                    message => erlang:list_to_binary(
                        [<<"Simulation '"/utf8>>,
                            Name@1,
                            <<"' exceeded max steps ("/utf8>>,
                            erlang:integer_to_binary(Max),
                            <<")"/utf8>>]
                    ),
                    file => <<?FILEPATH/utf8>>,
                    module => <<"glean/testing/simulation"/utf8>>,
                    function => <<"assert_passed"/utf8>>,
                    line => 202})
    end.

-file("src/glean/testing/simulation.gleam", 311).
-spec run_assert(simulation_state(), assert_fn()) -> {ok, simulation_state()} |
    {error, assertion_error()}.
run_assert(State, Assert_fn) ->
    case Assert_fn of
        no_assert ->
            {ok, State};

        {assert, F} ->
            case F(State) of
                {ok, nil} ->
                    {ok, State};

                {error, Err} ->
                    {error, Err}
            end
    end.

-file("src/glean/testing/simulation.gleam", 243).
-spec execute_step(simulation_state(), simulation_step()) -> {ok,
        simulation_state()} |
    {error, assertion_error()}.
execute_step(State, Step) ->
    case Step of
        {expect_text, Text, Then_assert} ->
            New_state = {simulation_state,
                lists:append(
                    erlang:element(2, State),
                    [glean@message:assistant(Text)]
                ),
                erlang:element(3, State) + 1,
                erlang:element(4, State),
                erlang:element(5, State),
                erlang:element(6, State)},
            run_assert(New_state, Then_assert);

        {expect_tool_calls, Executions, Then_assert@1} ->
            Tool_calls = gleam@list:map(
                Executions,
                fun(Exec) ->
                    {assistant_tool_call,
                        erlang:element(2, erlang:element(2, Exec)),
                        erlang:element(3, erlang:element(2, Exec)),
                        erlang:element(4, erlang:element(2, Exec))}
                end
            ),
            Tool_results = gleam@list:map(
                Executions,
                fun(Exec@1) -> case erlang:element(3, Exec@1) of
                        {tool_success, Content} ->
                            {tool_result_ok,
                                erlang:element(2, erlang:element(2, Exec@1)),
                                erlang:element(3, erlang:element(2, Exec@1)),
                                Content};

                        {tool_failure, Err} ->
                            {tool_result_error,
                                erlang:element(2, erlang:element(2, Exec@1)),
                                erlang:element(3, erlang:element(2, Exec@1)),
                                Err}
                    end end
            ),
            New_state@1 = {simulation_state,
                lists:append(
                    erlang:element(2, State),
                    [{assistant_message, Tool_calls},
                        {tool_result_message, Tool_results}]
                ),
                erlang:element(3, State) + 1,
                erlang:element(4, State),
                lists:append(erlang:element(5, State), Tool_results),
                erlang:element(6, State)},
            run_assert(New_state@1, Then_assert@1);

        {expect_error, Error, Then_assert@2} ->
            New_state@2 = {simulation_state,
                erlang:element(2, State),
                erlang:element(3, State) + 1,
                erlang:element(4, State),
                erlang:element(5, State),
                lists:append(erlang:element(6, State), [Error])},
            run_assert(New_state@2, Then_assert@2)
    end.

-file("src/glean/testing/simulation.gleam", 325).
-spec run_final_assertions(
    simulation(),
    simulation_state(),
    list(fun((simulation_state()) -> {ok, nil} | {error, assertion_error()}))
) -> simulation_result().
run_final_assertions(Sim, State, Assertions) ->
    case Assertions of
        [] ->
            {simulation_passed, erlang:element(2, Sim), State};

        [Assertion | Rest] ->
            case Assertion(State) of
                {ok, nil} ->
                    run_final_assertions(Sim, State, Rest);

                {error, Err} ->
                    {simulation_failed,
                        erlang:element(2, Sim),
                        erlang:element(3, State),
                        Err}
            end
    end.

-file("src/glean/testing/simulation.gleam", 214).
-spec run_steps(simulation(), list(simulation_step()), simulation_state()) -> simulation_result().
run_steps(Sim, Remaining, State) ->
    case erlang:element(3, State) >= erlang:element(6, Sim) of
        true ->
            {simulation_exceeded_max_steps,
                erlang:element(2, Sim),
                erlang:element(6, Sim),
                State};

        false ->
            case Remaining of
                [] ->
                    run_final_assertions(Sim, State, erlang:element(7, Sim));

                [Step | Rest] ->
                    case execute_step(State, Step) of
                        {error, Err} ->
                            {simulation_failed,
                                erlang:element(2, Sim),
                                erlang:element(3, State),
                                Err};

                        {ok, New_state} ->
                            run_steps(Sim, Rest, New_state)
                    end
            end
    end.

-file("src/glean/testing/simulation.gleam", 168).
?DOC(" Run the simulation. Pure function — all responses come from the script.\n").
-spec run(simulation()) -> simulation_result().
run(Sim) ->
    Initial_messages = case erlang:element(4, Sim) of
        <<""/utf8>> ->
            [glean@message:user(erlang:element(3, Sim))];

        Prompt ->
            [glean@message:system(Prompt),
                glean@message:user(erlang:element(3, Sim))]
    end,
    Initial_state = {simulation_state,
        Initial_messages,
        0,
        {usage, 0, 0},
        [],
        []},
    run_steps(Sim, erlang:element(5, Sim), Initial_state).