wampex_router

defmodule Wampex.Router.Session do
  @moduledoc """
  A state machine based process for managing a WAMP Session. Utilizes StatesLanguage to implement the gen_statem process. See priv/session.json for the JSON representation
  """
  use StatesLanguage, data: "priv/router.json"

  alias StatesLanguage, as: SL
  alias Wampex.Roles.{Broker, Dealer, Peer}
  alias Wampex.Roles.Broker.{Event, Published, Subscribed, Unsubscribed}
  alias Wampex.Roles.Callee.{Register, Unregister, Yield}
  alias Wampex.Roles.Caller.Call
  alias Wampex.Roles.Dealer.{Invocation, Registered, Result, Unregistered}
  alias Wampex.Roles.Peer.{Abort, Authenticate, Challenge, Error, Hello, Welcome}
  alias Wampex.Roles.Publisher.Publish
  alias Wampex.Roles.Subscriber.{Subscribe, Unsubscribe}
  alias Wampex.Router
  alias Wampex.Router.Realms
  alias Wampex.Router.Realms.Session, as: RealmSession
  alias __MODULE__, as: Sess

  @enforce_keys [:transport, :transport_pid]
  defstruct [
    :id,
    :db,
    :proxy,
    :transport,
    :transport_pid,
    :message,
    :hello,
    :authenticate,
    :register,
    :unregister,
    :call,
    :yield,
    :subscribe,
    :unsubscribe,
    :publish,
    :realm,
    :name,
    :goodbye,
    :peer,
    :challenge,
    :authentication_module,
    :authorization_module,
    error: "wamp.error.protocol_violation",
    roles: [Peer, Broker, Dealer],
    registrations: [],
    subscriptions: [],
    invocations: [],
    message_queue: [],
    request_id: 0,
    requests: [],
    hello_received: false
  ]

  @type t :: %__MODULE__{
          id: integer() | nil,
          db: module() | nil,
          proxy: module() | nil,
          transport_pid: pid() | module() | nil,
          message: Wampex.message() | nil,
          hello: Hello.t() | nil,
          authenticate: Authenticate.t() | nil,
          register: Register.t() | nil,
          unregister: Unregister.t() | nil,
          call: Call.t() | nil,
          yield: Yield.t() | nil,
          subscribe: Subscribe.t() | nil,
          unsubscribe: Unsubscribe.t() | nil,
          publish: Publish.t() | nil,
          goodbye: binary() | nil,
          realm: String.t() | nil,
          name: module() | nil,
          challenge: Challenge.t() | nil,
          authentication_module: module() | map(),
          authorization_module: module() | map(),
          roles: [module()],
          registrations: [],
          subscriptions: [],
          invocations: [],
          message_queue: [],
          request_id: integer(),
          requests: [],
          hello_received: boolean()
        }

  ## States
  @init "Init"
  @established "Established"
  @hello "Hello"
  @authenticate "Authenticate"
  @call "Call"
  @yield "Yield"
  @register "Register"
  @unregister "Unregister"
  @subscribe "Subscribe"
  @unsubscribe "Unsubscribe"
  @publish "Publish"
  @invocation_error "InvocationError"
  @abort "Abort"
  @goodbye "Goodbye"

  ## Resources
  @handle_init "HandleInit"
  @handle_established "HandleEstablished"
  @handle_message "HandleMessage"
  @handle_hello "HandleHello"
  @handle_authenticate "HandleAuthenticate"
  @handle_call "HandleCall"
  @handle_yield "HandleYield"
  @handle_register "HandleRegister"
  @handle_unregister "HandleUnregister"
  @handle_subscribe "HandleSubscribe"
  @handle_unsubscribe "HandleUnsubscribe"
  @handle_publish "HandlePublish"
  @handle_invocation_error "HandleInvocationError"
  # @handle_interrupt "HandleInterrupt"
  @handle_abort "HandleAbort"
  @handle_goodbye "HandleGoodbye"
  # @handle_cancel "HandleCancel"

  @impl true
  def handle_resource(
        @handle_init,
        _,
        @init,
        data
      ) do
    Logger.debug("Init")

    {:ok, %SL{data | data: %Sess{data.data | id: RealmSession.get_id()}}, [{:next_event, :internal, :transition}]}
  end

  @impl true
  def handle_resource(
        @handle_established,
        _,
        @established,
        data
      ) do
    Logger.debug("Established")
    {:ok, data, []}
  end

  @impl true
  def handle_resource(
        @handle_message,
        _,
        _,
        %SL{data: %Sess{message: msg, roles: roles}} = data
      ) do
    Logger.debug("Handling Message #{inspect(msg)}")

    {data, actions} =
      case handle_message(msg, roles) do
        {actions, _id, response} ->
          {data, response} = maybe_update_response(data, response)
          Logger.debug("Response: #{inspect(response)}")
          {data, actions}

        nil ->
          Logger.error("Message not supported by roles: #{inspect(roles)}")
          {data, [{:next_event, :internal, :abort}]}
      end

    {:ok, %SL{data | data: %Sess{data.data | message: nil}}, actions}
  rescue
    er ->
      Logger.error(inspect(er))
      {:ok, data, [{:next_event, :internal, :transition}]}
  end

  @impl true
  def handle_resource(
        @handle_hello,
        _,
        @hello,
        %SL{
          data:
            %Sess{
              id: session_id,
              transport: tt,
              transport_pid: t,
              hello: %Hello{realm: realm, options: options},
              hello_received: false,
              authentication_module: auth
            } = data
        } = sl
      ) do
    {actions, challenge, mod} = maybe_challenge(auth, realm, options, session_id, tt, t)

    {:ok,
     %SL{
       sl
       | data: %Sess{
           data
           | challenge: challenge,
             hello_received: true,
             realm: realm,
             authentication_module: mod
         }
     }, [{:next_event, :internal, :transition}] ++ actions}
  end

  @impl true
  def handle_resource(
        @handle_authenticate,
        _,
        @authenticate,
        %SL{
          data: %Sess{
            id: session_id,
            transport: tt,
            transport_pid: t,
            authentication_module: auth,
            challenge: %Challenge{extra: challenge},
            realm: realm,
            name: name,
            authenticate: %Authenticate{signature: sig}
          }
        } = sl
      ) do
    info = auth.parse_challenge(challenge)

    {actions, proxy, peer} = maybe_welcome({auth, session_id, sig, realm, name, challenge, info, tt, t})

    {:ok, %SL{sl | data: %Sess{sl.data | peer: peer, proxy: proxy}}, actions}
  end

  @impl true
  def handle_resource(
        @handle_subscribe,
        _,
        @subscribe,
        %SL{
          data:
            %Sess{
              transport: tt,
              transport_pid: t,
              subscriptions: subs,
              challenge: %Challenge{auth_method: method},
              realm: realm,
              authorization_module: am,
              peer: peer,
              subscribe: %Subscribe{request_id: rid, options: opts, topic: topic},
              db: db
            } = data
        } = sl
      ) do
    {data, actions} =
      case is_authorized?(am, peer, :subscribe, topic, method) do
        true ->
          id = RealmSession.get_id()
          wc = RealmSession.subscribe(db, realm, topic, {id, {self(), Node.self()}}, opts)
          send_to_peer(Broker.subscribed(%Subscribed{request_id: rid, subscription_id: id}), tt, t)

          Logger.info("Subscription #{id}: #{realm} - #{topic}")

          {%SL{sl | data: %Sess{data | subscriptions: [{id, topic, wc} | subs]}},
           [{:next_event, :internal, :transition}]}

        false ->
          send_to_peer(Broker.subscribe_error(%Error{request_id: rid, error: "wamp.error.not_authorized"}), tt, t)
          {sl, [{:next_event, :internal, :transition}]}
      end

    {:ok, data, actions}
  end

  @impl true
  def handle_resource(
        @handle_unsubscribe,
        _,
        @unsubscribe,
        %SL{
          data: %Sess{
            realm: realm,
            transport: tt,
            transport_pid: t,
            db: db,
            subscriptions: subs,
            unsubscribe: %Unsubscribe{request_id: rid, subscription_id: subscription_id}
          }
        } = sl
      ) do
    subs = RealmSession.unsubscribe(db, realm, {subscription_id, {self(), Node.self()}}, subs)

    send_to_peer(Broker.unsubscribed(%Unsubscribed{request_id: rid}), tt, t)
    Logger.info("Unsubscribe: #{realm} - #{subscription_id}")
    {:ok, %SL{sl | data: %Sess{sl.data | subscriptions: subs}}, [{:next_event, :internal, :transition}]}
  end

  @impl true
  def handle_resource(
        @handle_publish,
        _,
        @publish,
        %SL{
          data: %Sess{
            proxy: proxy,
            realm: realm,
            transport: tt,
            transport_pid: t,
            db: db,
            challenge: %Challenge{auth_method: method},
            authorization_module: am,
            peer: peer,
            id: session_id,
            publish: %Publish{request_id: rid, options: opts, topic: topic, arg_list: arg_l, arg_kw: arg_kw}
          }
        } = sl
      ) do
    {data, actions} =
      case is_authorized?(am, peer, :publish, topic, method) do
        true ->
          pub_id = RealmSession.get_id()

          subs = RealmSession.subscriptions(db, realm, topic)
          details = Map.put_new(opts, "topic", topic)
          details = Map.put_new(details, "session_id", session_id)

          groups =
            Enum.reduce(subs, %{}, fn {id, {pid, node}}, acc ->
              event = {{:event, id, pub_id, arg_l, arg_kw, details}, pid}
              Map.update(acc, node, [event], &[event | &1])
            end)

          Enum.each(groups, fn {node, messages} ->
            send({proxy, node}, messages)
          end)

          case opts do
            %{"acknowledge" => true} ->
              send_to_peer(Broker.published(%Published{request_id: rid, publication_id: pub_id}), tt, t)

            %{} ->
              :noop
          end

          Logger.info("Published #{pub_id}: #{realm} - #{topic}")
          {sl, [{:next_event, :internal, :transition}]}

        false ->
          send_to_peer(Broker.publish_error(%Error{request_id: rid, error: "wamp.error.not_authorized"}), tt, t)
          {sl, [{:next_event, :internal, :transition}]}
      end

    {:ok, %SL{data | data: %Sess{data.data | publish: nil}}, actions}
  end

  @impl true
  def handle_resource(
        @handle_unregister,
        _,
        @unregister,
        %SL{
          data:
            %Sess{
              db: db,
              transport: tt,
              transport_pid: t,
              registrations: regs,
              realm: realm,
              unregister: %Unregister{request_id: request_id, registration_id: registration_id}
            } = data
        } = sl
      ) do
    regs = RealmSession.unregister(db, realm, {registration_id, {self(), Node.self()}}, regs)

    send_to_peer(Dealer.unregistered(%Unregistered{request_id: request_id}), tt, t)
    Logger.info("Unregister: #{realm} - #{registration_id}")
    {:ok, %SL{sl | data: %Sess{data | registrations: regs}}, [{:next_event, :internal, :transition}]}
  end

  @impl true
  def handle_resource(
        @handle_register,
        _,
        @register,
        %SL{
          data:
            %Sess{
              transport: tt,
              transport_pid: t,
              registrations: regs,
              db: db,
              realm: realm,
              challenge: %Challenge{auth_method: method},
              authorization_module: am,
              peer: peer,
              register: %Register{request_id: rid, procedure: procedure}
            } = data
        } = sl
      ) do
    {data, actions} =
      case is_authorized?(am, peer, :register, procedure, method) do
        true ->
          id = RealmSession.get_id()
          RealmSession.register(db, realm, procedure, {id, {self(), Node.self()}})
          regd = %Registered{request_id: rid, registration_id: id}
          send_to_peer(Dealer.registered(regd), tt, t)

          Logger.info("Registered #{id}: #{realm} - #{procedure}")

          {%SL{sl | data: %Sess{data | registrations: [{id, procedure} | regs]}},
           [{:next_event, :internal, :transition}]}

        false ->
          send_to_peer(Dealer.register_error(%Error{request_id: rid, error: "wamp.error.not_authorized"}), tt, t)
          {sl, [{:next_event, :internal, :transition}]}
      end

    {:ok, data, actions}
  end

  @impl true
  def handle_resource(
        @handle_call,
        _,
        @call,
        %SL{
          data:
            %Sess{
              db: db,
              proxy: proxy,
              transport: tt,
              transport_pid: t,
              realm: realm,
              request_id: ri,
              challenge: %Challenge{auth_method: method},
              authorization_module: am,
              id: session_id,
              peer: peer,
              call: %Call{request_id: call_id, options: opts, procedure: proc, arg_list: al, arg_kw: akw}
            } = data
        } = sl
      ) do
    {data, actions} =
      case is_authorized?(am, peer, :call, proc, method) do
        true ->
          with {callees, _index} <- RealmSession.callees(db, realm, proc),
               {id, {pid, node}} <- get_live_callee(proxy, callees, 1000) do
            req_id = RealmSession.get_request_id(ri)
            opts = Map.put(opts, "procedure", proc)
            opts = Map.put(opts, "session_id", session_id)

            send(
              {proxy, node},
              {
                {
                  call_id,
                  %Invocation{
                    request_id: req_id,
                    registration_id: id,
                    details: opts,
                    arg_list: al,
                    arg_kw: akw
                  },
                  {self(), Node.self()}
                },
                pid
              }
            )

            Logger.info("Call #{req_id}: #{realm} - #{proc}")
            {%SL{sl | data: %Sess{data | request_id: req_id}}, [{:next_event, :internal, :transition}]}
          else
            {:error, :no_live_callees} ->
              send_to_peer(
                Dealer.call_error(%Error{
                  request_id: call_id,
                  error: "wamp.error.no_callees",
                  details: %{procedure: proc}
                }),
                tt,
                t
              )

              Logger.info("No Callees: #{realm} - #{proc}")
              {sl, [{:next_event, :internal, :transition}]}

            :not_found ->
              send_to_peer(
                Dealer.call_error(%Error{
                  request_id: call_id,
                  error: "wamp.error.no_registration",
                  details: %{procedure: proc}
                }),
                tt,
                t
              )

              {sl, [{:next_event, :internal, :transition}]}
          end

        false ->
          send_to_peer(Dealer.call_error(%Error{request_id: call_id, error: "wamp.error.not_authorized"}), tt, t)
          {sl, [{:next_event, :internal, :transition}]}
      end

    {:ok, data, actions}
  end

  @impl true
  def handle_resource(
        @handle_yield,
        _,
        @yield,
        %SL{
          data: %Sess{
            proxy: proxy,
            invocations: inv,
            realm: realm,
            yield: %Yield{request_id: id, options: opts, arg_list: arg_l, arg_kw: arg_kw}
          }
        } = data
      ) do
    case Enum.find(inv, fn
           {_, %Invocation{request_id: ^id}, _} -> true
           _ -> false
         end) do
      {call_id, _, {pid, node}} ->
        send(
          {proxy, node},
          {%Result{request_id: call_id, arg_list: arg_l, arg_kw: arg_kw, details: opts}, pid}
        )

        Logger.info("Yield: #{realm} - #{id}")

      _ ->
        Logger.error("Unable to yield: #{id} not present in invocations: #{inspect(inv)}")
        :ok
    end

    inv =
      Enum.filter(inv, fn
        {_, %Invocation{request_id: ^id}, _} -> false
        _ -> true
      end)

    {:ok, %SL{data | data: %Sess{data.data | invocations: inv, yield: nil}}, [{:next_event, :internal, :transition}]}
  end

  @impl true
  def handle_resource(
        @handle_invocation_error,
        _,
        @invocation_error,
        %SL{
          data: %Sess{
            proxy: proxy,
            invocations: inv,
            realm: realm,
            error: %Error{request_id: id, details: dets, arg_list: arg_l, error: err, arg_kw: arg_kw}
          }
        } = data
      ) do
    {call_id, _, {pid, node}} =
      Enum.find(inv, fn
        {_, %Invocation{request_id: ^id}, _} -> true
        _ -> false
      end)

    send(
      {proxy, node},
      {%Error{request_id: call_id, type: 48, error: err, arg_list: arg_l, arg_kw: arg_kw, details: dets}, pid}
    )

    inv =
      Enum.filter(inv, fn
        {_, %Invocation{request_id: ^id}, _} -> false
        _ -> true
      end)

    Logger.info("Call Error #{id}: #{realm} - #{get_in(dets, ["procedure"])}")
    {:ok, %SL{data | data: %Sess{data.data | invocations: inv}}, [{:next_event, :internal, :transition}]}
  end

  @impl true
  def handle_resource(
        @handle_abort,
        _,
        @abort,
        %SL{data: %Sess{transport: tt, transport_pid: t, realm: realm, error: er}} = data
      ) do
    send_to_peer(Peer.abort(%Abort{reason: er}), tt, t)

    Logger.info("Abort: #{realm} #{inspect(er)}")
    {:ok, data, []}
  end

  @impl true
  def handle_resource(
        @handle_goodbye,
        _,
        @goodbye,
        %SL{data: %Sess{transport: tt, transport_pid: t, goodbye: goodbye, realm: realm}} = data
      ) do
    send_to_peer(Peer.goodbye(goodbye), tt, t)

    Logger.info("Goodbye: #{realm}")
    {:ok, data, []}
  end

  @impl true
  def handle_resource(resource, _, state, data) do
    Logger.error("No specific resource handler for #{resource} in state #{state}")
    {:ok, data, []}
  end

  @impl true
  def handle_info(
        {call_id, %Invocation{} = e, from},
        _,
        %SL{data: %Sess{transport: tt, transport_pid: t, invocations: inv}} = data
      ) do
    send_to_peer(Dealer.invocation(e), tt, t)
    {:ok, %SL{data | data: %Sess{data.data | invocations: [{call_id, e, from} | inv]}}, []}
  end

  @impl true
  def handle_info(
        {:event, id, pub_id, arg_l, arg_kw, details},
        _,
        %SL{data: %Sess{transport: tt, transport_pid: t}} = data
      ) do
    send_to_peer(
      Broker.event(%Event{
        subscription_id: id,
        publication_id: pub_id,
        arg_list: arg_l,
        arg_kw: arg_kw,
        details: details
      }),
      tt,
      t
    )

    {:ok, data, []}
  end

  @impl true
  def handle_info(%Result{} = e, _, %SL{data: %Sess{transport: tt, transport_pid: t}} = data) do
    send_to_peer(Dealer.result(e), tt, t)
    {:ok, data, []}
  end

  @impl true
  def handle_info(%Error{} = e, _, %SL{data: %Sess{transport: tt, transport_pid: t}} = data) do
    send_to_peer(Dealer.call_error(e), tt, t)
    {:ok, data, []}
  end

  @impl true
  def handle_info({:set_message, message}, _, %SL{data: sess} = data) do
    {:ok, %SL{data | data: %Sess{sess | message: message}}, [{:next_event, :internal, :message_received}]}
  end

  @impl true
  def handle_info({:EXIT, t, reason}, _, %SL{data: %Sess{transport: tt, transport_pid: t}} = data) do
    Logger.debug("Transport #{tt} exited because #{inspect(reason)}: Shutting down")
    {:ok, data, [{:next_event, :internal, :abort}]}
  end

  @impl true
  def handle_info(event, state, data) do
    Logger.warn("Got event #{inspect(event)} in state #{state} with data #{inspect(data)}")
    {:ok, data, []}
  end

  @impl true
  def handle_termination(_reason, _, %SL{
        data: %Sess{
          id: session_id,
          db: db,
          registrations: regs,
          realm: realm,
          subscriptions: subs,
          invocations: inv,
          proxy: proxy
        }
      }) do
    Enum.each(inv, fn {call_id, _, {pid, node}} ->
      send(
        {proxy, node},
        {%Error{request_id: call_id, type: @call, error: "wamp.error.callee_went_away"}, pid}
      )
    end)

    pub_id = RealmSession.get_id()
    topic = "router.peer.disconnect"
    d_subs = RealmSession.subscriptions(db, realm, topic)
    details = %{topic: topic, session_id: session_id}

    groups =
      Enum.reduce(d_subs, %{}, fn {id, {pid, node}}, acc ->
        event = {{:event, id, pub_id, [], %{}, details}, pid}
        Map.update(acc, node, [event], &[event | &1])
      end)

    Enum.each(groups, fn {node, messages} ->
      send({proxy, node}, messages)
    end)

    Enum.each(regs, fn {id, _proc} ->
      RealmSession.unregister(db, realm, {id, {self(), Node.self()}}, regs)
    end)

    Enum.each(subs, fn {id, _proc, _} ->
      RealmSession.unsubscribe(db, realm, {id, {self(), Node.self()}}, subs)
    end)
  end

  defp maybe_welcome(
         {auth, session_id, sig, realm, name, challenge, {authid, authrole, authmethod, authprovider}, tt, t}
       ) do
    case auth.authenticate(sig, realm, authid, challenge) do
      {true, peer} ->
        proxy = Realms.start_realm(Router.realms_name(name), realm)

        send_to_peer(
          Peer.welcome(%Welcome{
            session_id: session_id,
            options: %{
              agent: "WAMPex Router",
              authid: authid,
              authrole: authrole,
              authmethod: authmethod,
              authprovider: authprovider,
              roles: %{broker: %{}, dealer: %{}}
            }
          }),
          tt,
          t
        )

        Logger.info("Welcome #{authid}: #{realm}")

        {[{:next_event, :internal, :transition}], proxy, peer}

      {false, peer} ->
        Logger.error("Invalid Authentication: #{inspect(authid)}")
        {[{:next_event, :internal, :transition}, {:next_event, :internal, :abort}], nil, peer}
    end
  end

  defp maybe_challenge(auth, realm, details, session_id, tt, t) do
    case details do
      %{"authid" => ai, "authmethods" => am} ->
        auth = get_auth_module(auth, am)

        case auth.authenticate?(am) do
          true ->
            ch = auth.challenge(realm, ai, session_id)

            chal = %Challenge{
              auth_method: auth.method(),
              extra: ch
            }

            send_to_peer(
              Peer.challenge(chal),
              tt,
              t
            )

            {[], chal, auth}

          false ->
            Logger.error("Not a supported authentication method: #{inspect(am)}")
            {[{:next_event, :internal, :abort}], nil}
        end

      %{} ->
        Logger.error("No authentication details given")
        {[{:next_event, :internal, :abort}], nil}
    end
  end

  defp get_auth_module(auth, methods) when is_map(auth) do
    case Enum.find(auth, fn {k, _v} -> k in methods end) do
      {_, mod} -> mod
      nil -> Logger.error("Unsupported authentication methods #{inspect(methods)}: Supported #{inspect(auth)}")
    end
  end

  defp get_auth_module(auth, _methods), do: auth

  defp get_live_callee(_proxy, [], 0) do
    Logger.error("No live callees, tried all replicas")
    {:error, :no_live_callees}
  end

  defp get_live_callee(_proxy, [], _) do
    Logger.error("No live callees, empty result from lookup")
    {:error, :no_live_callees}
  end

  defp get_live_callee(proxy, callees, tries) when is_list(callees) do
    {_id, {pid, node}} = c = Enum.random(callees)

    nodes = [Node.self() | Node.list()]

    with true <- node in nodes,
         true <- GenServer.call({proxy, node}, {:is_up, pid}) do
      c
    else
      false ->
        get_live_callee(proxy, callees, tries - 1)
    end
  end

  defp get_live_callee(proxy, result, tries) do
    Logger.error("No live callees, Proxy: #{inspect(proxy)} Result: #{inspect(result)} Tries: #{tries}")

    {:error, :no_live_callees}
  end

  defp send_to_peer(msg, transport, pid) do
    transport.send_request(pid, remove_nil_values(msg))
  end

  defp maybe_update_response(data, {:update, key, resp}) do
    {%SL{data | data: Map.put(data.data, key, resp)}, resp}
  end

  defp maybe_update_response(data, resp), do: {data, resp}

  def is_authorized?(am, peer, type, uri, method) do
    auth = get_auth_module(am, [method])
    auth.authorized?(type, uri, peer)
  end

  defp remove_nil_values(message) do
    message =
      case Enum.reverse(message) do
        [map | t] when map == %{} -> t
        m -> m
      end

    message =
      case message do
        [list | t] when list == [] ->
          t

        m ->
          m
      end

    Enum.reverse(message)
  end

  defp handle_message(msg, roles) do
    Enum.reduce_while(roles, nil, fn r, _ ->
      try do
        res = r.handle(msg)
        {:halt, res}
      rescue
        FunctionClauseError -> {:cont, nil}
      end
    end)
  end
end