% SPDX-License-Identifier: LicenseRef-Sovereign-Leviathan-SL-AGPL3-001
% Licensing: see LICENSE in the repository root.

%%%-------------------------------------------------------------------
%%% @doc `/agents' API domain.
%%%
%%%   POST   /api/agents                  create an agent
%%%   GET    /api/agents                  list agents
%%%   GET    /api/agents/:id              describe one agent
%%%   POST   /api/agents/:id/objective    assign its objective
%%%   POST   /api/agents/:id/actions      run one loop iteration (a tool call)
%%%   POST   /api/agents/:id/complete     mark completed
%%%   POST   /api/agents/:id/terminate    terminate
%%%
%%% Every POST here goes through `hk_api_util:run_operation/4', so
%%% every mutating call returns an `operation_id' alongside its
%%% result -- including a failed one, per the "every mutating
%%% operation should return an operation identifier" requirement
%%% (an operation identifier is returned whether the operation
%%% succeeded or not; only its `status' differs).
%%% @end
%%%-------------------------------------------------------------------
-module(hk_api_agents_h).

-export([init/2]).

init(Req0 = #{method := <<"POST">>}, collection) ->
    RequestId = hk_api_util:request_id(Req0),
    {ok, Body, Req1} = hk_api_util:read_json_body(Req0),
    Capabilities = parse_capabilities(maps:get(<<"capabilities">>, Body, [])),
    Limits = maps:get(<<"limits">>, Body, #{}),
    Req = hk_api_util:run_operation(agent_create, RequestId, fun() ->
        AgentId = hk_id:new(<<"agent">>),
        case hk_agent_fsm_sup:start_agent(#{agent_id => AgentId, capabilities => Capabilities,
                                              limits => atomize_limit_keys(Limits)}) of
            {ok, _Pid} -> {ok, {agent, AgentId}, #{agent_id => AgentId}};
            {error, Reason} -> {error, Reason}
        end
    end, Req1),
    {ok, Req, collection};

init(Req0 = #{method := <<"GET">>}, collection) ->
    Described = [describe_or_skip(Id) || {Id, _Pid} <- hk_agent_registry:list_agents()],
    Agents = [M || {true, M} <- Described],
    Req = hk_api_util:reply_json(200, #{agents => Agents}, Req0),
    {ok, Req, collection};

init(Req0 = #{method := <<"GET">>}, item) ->
    AgentId = cowboy_req:binding(id, Req0),
    Req = case lookup_agent(AgentId) of
              {ok, Pid} ->
                  {ok, Agent} = hk_agent_fsm:describe(Pid),
                  hk_api_util:reply_json(200, hk_schema:to_map(Agent), Req0);
              {error, not_found} ->
                  hk_api_util:reply_error(404, agent_not_found, Req0)
          end,
    {ok, Req, item};

init(Req0 = #{method := <<"POST">>}, objective) ->
    AgentId = cowboy_req:binding(id, Req0),
    RequestId = hk_api_util:request_id(Req0),
    {ok, Body, Req1} = hk_api_util:read_json_body(Req0),
    Description = maps:get(<<"description">>, Body, <<>>),
    Req = hk_api_util:run_operation(agent_assign_objective, RequestId, fun() ->
        case lookup_agent(AgentId) of
            {ok, Pid} ->
                case hk_agent_fsm:assign_objective(Pid, Description) of
                    ok -> {ok, {agent, AgentId}, #{agent_id => AgentId}};
                    {error, Reason} -> {error, Reason}
                end;
            {error, Reason} -> {error, Reason}
        end
    end, Req1),
    {ok, Req, objective};

init(Req0 = #{method := <<"POST">>}, actions) ->
    AgentId = cowboy_req:binding(id, Req0),
    RequestId = hk_api_util:request_id(Req0),
    {ok, Body, Req1} = hk_api_util:read_json_body(Req0),
    Tool = binary_to_atom(maps:get(<<"tool">>, Body, <<"browser">>), utf8),
    Op = binary_to_atom(maps:get(<<"op">>, Body, <<"read_page">>), utf8),
    Args = atomize_arg_keys(maps:get(<<"args">>, Body, #{})),
    Req = hk_api_util:run_operation(agent_action, RequestId, fun() ->
        case lookup_agent(AgentId) of
            {ok, Pid} ->
                case hk_agent_fsm:next_action(Pid, Tool, Op, Args) of
                    {ok, Result} -> {ok, {agent, AgentId}, #{observation => Result}};
                    {error, Reason} -> {error, Reason}
                end;
            {error, Reason} -> {error, Reason}
        end
    end, Req1),
    {ok, Req, actions};

init(Req0 = #{method := <<"POST">>}, complete) ->
    AgentId = cowboy_req:binding(id, Req0),
    RequestId = hk_api_util:request_id(Req0),
    {ok, Body, Req1} = hk_api_util:read_json_body(Req0),
    Result = maps:get(<<"result">>, Body, null),
    Req = hk_api_util:run_operation(agent_complete, RequestId, fun() ->
        case lookup_agent(AgentId) of
            {ok, Pid} -> ok = hk_agent_fsm:complete(Pid, Result), {ok, {agent, AgentId}, #{}};
            {error, Reason} -> {error, Reason}
        end
    end, Req1),
    {ok, Req, complete};

init(Req0 = #{method := <<"POST">>}, terminate) ->
    AgentId = cowboy_req:binding(id, Req0),
    RequestId = hk_api_util:request_id(Req0),
    {ok, Body, Req1} = hk_api_util:read_json_body(Req0),
    Reason0 = maps:get(<<"reason">>, Body, <<"operator_requested">>),
    Req = hk_api_util:run_operation(agent_terminate, RequestId, fun() ->
        case lookup_agent(AgentId) of
            {ok, Pid} -> ok = hk_agent_fsm:terminate_agent(Pid, Reason0), {ok, {agent, AgentId}, #{}};
            {error, Reason} -> {error, Reason}
        end
    end, Req1),
    {ok, Req, terminate};

init(Req0, State) ->
    Req = hk_api_util:reply_error(405, method_not_allowed, Req0),
    {ok, Req, State}.

%%% internal

lookup_agent(undefined) -> {error, missing_agent_id};
lookup_agent(AgentId) -> hk_agent_registry:lookup(AgentId).

describe_or_skip(AgentId) ->
    case lookup_agent(AgentId) of
        {ok, Pid} ->
            case catch hk_agent_fsm:describe(Pid) of
                {ok, Agent} -> {true, hk_schema:to_map(Agent)};
                _ -> false
            end;
        _ -> false
    end.

parse_capabilities(List) when is_list(List) ->
    [binary_to_atom(C, utf8) || C <- List, is_binary(C)];
parse_capabilities(_) ->
    [].

atomize_limit_keys(Map) when is_map(Map) ->
    maps:fold(fun(K, V, Acc) -> Acc#{binary_to_atom(K, utf8) => V} end, #{}, Map);
atomize_limit_keys(_) -> #{}.

atomize_arg_keys(Map) when is_map(Map) ->
    maps:fold(fun(K, V, Acc) -> Acc#{binary_to_atom(K, utf8) => V} end, #{}, Map);
atomize_arg_keys(_) -> #{}.