defmodule DurableStash do
  @moduledoc """
  Keep Phoenix LiveView state alive across reconnects, crashes, and redeploys.

  A LiveView's assigns live in memory and vanish when the socket drops, the
  process crashes, or you deploy. `DurableStash` saves the assigns you pick to
  object storage and restores them the next time the LiveView mounts. The saved
  copy belongs to the browser session, so every LiveView the user opens shares
  it. Under the hood it plugs into
  [LiveStash](https://hex.pm/packages/live_stash) and stores state through
  [DurableServer](https://hex.pm/packages/durable_server). State dies with the
  browser session.

  ## Usage

      defmodule MyAppWeb.SomeLive do
        use MyAppWeb, :live_view
        use LiveStash, adapter: DurableStash, stored_keys: [:count, :username]

        def mount(_params, _session, socket) do
          socket = assign(socket, count: 0, username: nil)
          {_status, socket} = LiveStash.recover_state(socket)
          {:ok, socket}
        end
      end

  Call `LiveStash.recover_state/1` in `mount/3` *after* assigning defaults
  (recovered values overwrite them), and `LiveStash.stash/1` whenever a
  stored assign changes — or pass `auto_stash: true`.

  Unlike the stock ETS adapter, `DurableStash` recovers on **every** mount —
  fresh navigations included — and never deletes the stash on mount. State is
  keyed by the browser session, not by the socket.

  ## Setup

  1. Register the adapter and configure the backend:

         config :live_stash, adapters: [DurableStash]

         config :durable_stash,
           backend: {DurableServer.Backends.ObjectStore, bucket: "...", ...},
           prefix: "durable_stash/",
           secret: "some-stable-secret"

     With that config, `LiveStash.Application` starts the DurableServer
     supervisor automatically through `child_spec/1`. Alternatively, run your
     own `DurableServer.Supervisor` and point the adapter at it with the
     `:supervisor` option.

  2. Put a session id into the cookie session, in the `:browser` pipeline
     after `plug :fetch_session`:

         plug :ensure_session_id

         defp ensure_session_id(conn, _opts) do
           if get_session(conn, "sid") do
             conn
           else
             sid = 16 |> :crypto.strong_rand_bytes() |> Base.url_encode64(padding: false)
             put_session(conn, "sid", sid)
           end
         end

  ## Scopes

  Not all state wants the same recovery policy. Each stored key declares a
  scope:

      use LiveStash, adapter: DurableStash,
        stored_keys: [
          theme: :session,    # recover on every mount (the default)
          draft: :reconnect   # recover only on reconnects; cleared on fresh mounts
        ]

    * `:session` — recovered on every mount: live navigation, reconnects,
      crashes, redeploys. Right for settings the user expects to stick.
    * `:reconnect` — recovered only when the client *rejoins* an existing
      view (`_mounts > 0`): Wi-Fi drops, LiveView crashes, and redeploys —
      the browser stays on the page through all of these. A fresh navigation
      to the view clears the stored values, so starting a "new thing" starts
      blank. Right for in-progress form drafts.

  ## Options (via `use LiveStash, adapter: DurableStash, ...`)

    * `:stored_keys` (required) — assigns to persist. Bare atoms mean
      `:session` scope; see *Scopes* above for `:reconnect`. `:permanent` is
      reserved and raises for now.
    * `:vsn` (default `1`) — version of this view's stored shape. On
      recovery, a stored slice with a different vsn is discarded to defaults
      unless `:migrate` is given.
    * `:migrate` — 2-arity function `(old_vsn, data) :: data` receiving the
      stored string-keyed data map and returning the migrated one. The
      migrated set is written back under the new vsn.
    * `:supervisor` — DurableServer supervisor name (default
      `config :durable_stash, :supervisor_name`, falling back to
      `DurableStash.Supervisor`).
    * `:secret` — mixed into the storage-key hash (default
      `config :durable_stash, :secret`).
    * `:session_id_key` — cookie-session key holding the session id
      (default `"sid"`).

  ## What's storable

  JSON-safe values only: no structs, tuples, pids, or functions; maps come
  back with string keys. Offending values are skipped with a logged error, or
  raise when `config :durable_stash, on_invalid_value: :raise` is set
  (recommended for dev and test). Values are normalized through a JSON
  round-trip at stash time, so what you recover in dev is byte-for-byte what
  you'd recover after a redeploy in prod.
  """

  @behaviour LiveStash.Adapter

  require Logger

  alias DurableStash.Session
  alias Phoenix.Component
  alias Phoenix.LiveView

  defmodule Context do
    @moduledoc false
    defstruct supervisor: nil,
              storage_key: nil,
              view: nil,
              stored_keys: [],
              vsn: 1,
              migrate: nil,
              fingerprints: %{},
              reconnected?: false
  end

  @private_key :durable_stash
  @default_secret "durable_stash"
  @default_session_id_key "sid"

  ## LiveStash.Adapter callbacks

  @impl true
  def init_stash(socket, session, opts) do
    stored_keys = parse_stored_keys!(Keyword.fetch!(opts, :stored_keys))
    vsn = Keyword.get(opts, :vsn, 1)
    migrate = parse_migrate!(Keyword.get(opts, :migrate))

    supervisor =
      Keyword.get(opts, :supervisor) ||
        Application.get_env(:durable_stash, :supervisor_name, DurableStash.Supervisor)

    secret =
      Keyword.get(opts, :secret) ||
        Application.get_env(:durable_stash, :secret, @default_secret)

    session_id_key = Keyword.get(opts, :session_id_key, @default_session_id_key)

    context = %Context{
      supervisor: supervisor,
      storage_key: derive_storage_key(session, session_id_key, secret),
      view: view_name(socket),
      stored_keys: stored_keys,
      vsn: vsn,
      migrate: migrate,
      reconnected?: reconnected?(socket)
    }

    clear_stale_reconnect_keys(socket, context)

    LiveView.put_private(socket, @private_key, context)
  end

  @impl true
  def stash(socket) do
    case operable_context(socket) do
      {:ok, context} -> do_stash(socket, context)
      :error -> socket
    end
  end

  @impl true
  def recover_state(socket) do
    case operable_context(socket) do
      {:ok, context} -> do_recover(socket, context)
      :error -> {:error, socket}
    end
  end

  @impl true
  def reset_stash(socket) do
    case operable_context(socket) do
      {:ok, context} -> do_reset(socket, context)
      :error -> socket
    end
  end

  @doc """
  Starts the DurableServer supervisor from `:durable_stash` config when a
  `:backend` is configured; otherwise starts an empty, harmless supervisor.
  Invoked automatically by `LiveStash.Application` for registered adapters.
  """
  @impl true
  def child_spec(_args) do
    %{
      id: __MODULE__,
      start: {__MODULE__, :start_configured_supervisor, []},
      type: :supervisor
    }
  end

  @doc false
  def start_configured_supervisor do
    children =
      case Application.get_env(:durable_stash, :backend) do
        nil ->
          []

        backend ->
          supervisor_opts = Application.get_env(:durable_stash, :supervisor_opts, [])

          [
            {DurableServer.Supervisor,
             [
               name:
                 Application.get_env(:durable_stash, :supervisor_name, DurableStash.Supervisor),
               prefix: Application.get_env(:durable_stash, :prefix, "durable_stash/"),
               backend: backend
             ] ++ supervisor_opts}
          ]
      end

    Supervisor.start_link(children, strategy: :one_for_one)
  end

  ## Stash

  defp do_stash(socket, %Context{} = context) do
    {changes, fingerprints} = changed_entries(socket.assigns, context)

    if changes == %{} do
      socket
    else
      case call_session(context, &Session.merge(&1, context.view, changes, context.vsn)) do
        :ok ->
          put_context(socket, %{context | fingerprints: fingerprints})

        {:error, reason} ->
          Logger.error("[DurableStash] stash failed for #{context.view}: #{inspect(reason)}")
          socket
      end
    end
  end

  defp changed_entries(assigns, %Context{} = context) do
    context.stored_keys
    |> Enum.map(fn {key, _scope} -> key end)
    |> Enum.reduce({%{}, context.fingerprints}, fn key, {changes, fingerprints} ->
      with {:ok, value} <- Map.fetch(assigns, key),
           {:ok, encoded} <- encode_value(key, value) do
        name = Atom.to_string(key)
        fingerprint = fingerprint(encoded)

        if fingerprints[name] == fingerprint do
          {changes, fingerprints}
        else
          {Map.put(changes, name, JSON.decode!(encoded)),
           Map.put(fingerprints, name, fingerprint)}
        end
      else
        :error -> {changes, fingerprints}
        {:skip, _key} -> {changes, fingerprints}
      end
    end)
  end

  defp encode_value(key, value) do
    {:ok, JSON.encode!(value)}
  rescue
    error ->
      message =
        "[DurableStash] value for #{inspect(key)} is not JSON-safe " <>
          "(no structs, tuples, pids, or functions): #{Exception.message(error)}"

      case Application.get_env(:durable_stash, :on_invalid_value, :log) do
        :raise ->
          reraise ArgumentError.exception(message), __STACKTRACE__

        _log ->
          Logger.error(message)
          {:skip, key}
      end
  end

  ## Recover

  defp do_recover(socket, %Context{} = context) do
    case call_session(context, &Session.fetch_view(&1, context.view)) do
      :not_found ->
        {:not_found, socket}

      {:ok, %{"vsn" => stored_vsn, "data" => data}} ->
        recover_slice(socket, context, stored_vsn, data)

      {:error, reason} ->
        Logger.error("[DurableStash] recover failed for #{context.view}: #{inspect(reason)}")
        {:error, socket}
    end
  end

  defp recover_slice(socket, %Context{vsn: vsn} = context, vsn, data) do
    apply_recovered(socket, context, data)
  end

  defp recover_slice(socket, %Context{migrate: nil}, _stored_vsn, _data) do
    # Version mismatch without a migration: UI-grade state, discard to
    # defaults. The next stash overwrites the slice under the current vsn.
    {:not_found, socket}
  end

  defp recover_slice(socket, %Context{} = context, stored_vsn, data) do
    migrated = context.migrate.(stored_vsn, data)

    case apply_recovered(socket, context, migrated) do
      {:recovered, recovered_socket} ->
        write_back_migrated(recovered_socket, migrated)

      other ->
        other
    end
  rescue
    error ->
      Logger.error(
        "[DurableStash] migrate from vsn #{stored_vsn} failed for #{context.view}: " <>
          Exception.message(error)
      )

      {:error, socket}
  end

  # Persist the full migrated set under the new vsn so the slice can never
  # end up as a mix of old- and new-shape keys.
  defp write_back_migrated(socket, migrated) do
    context = socket.private[@private_key]
    restored = Map.take(migrated, Enum.map(recoverable_keys(context), &Atom.to_string/1))

    case call_session(context, &Session.merge(&1, context.view, restored, context.vsn)) do
      :ok ->
        {:recovered, socket}

      {:error, reason} ->
        Logger.error(
          "[DurableStash] migrated write-back failed for #{context.view}: #{inspect(reason)}"
        )

        {:recovered, socket}
    end
  end

  defp apply_recovered(socket, %Context{} = context, data) do
    {recovered, fingerprints} =
      context
      |> recoverable_keys()
      |> Enum.reduce({%{}, context.fingerprints}, fn key, {recovered, fingerprints} ->
        # Atom keys come from the declared whitelist — never String.to_atom
        # on stored input.
        case Map.fetch(data, Atom.to_string(key)) do
          {:ok, value} ->
            fingerprint = fingerprint(JSON.encode!(value))

            {Map.put(recovered, key, value),
             Map.put(fingerprints, Atom.to_string(key), fingerprint)}

          :error ->
            {recovered, fingerprints}
        end
      end)

    if recovered == %{} do
      {:not_found, socket}
    else
      socket
      |> Component.assign(recovered)
      |> put_context(%{context | fingerprints: fingerprints})
      |> then(&{:recovered, &1})
    end
  end

  ## Reset

  defp do_reset(socket, %Context{} = context) do
    case call_session(context, &Session.reset_view(&1, context.view)) do
      :ok ->
        put_context(socket, %{context | fingerprints: %{}})

      {:error, reason} ->
        Logger.error("[DurableStash] reset failed for #{context.view}: #{inspect(reason)}")
        socket
    end
  end

  ## Session access

  defp call_session(%Context{} = context, fun) do
    case DurableServer.Supervisor.ensure_started_child(
           context.supervisor,
           {Session, key: context.storage_key, initial_state: %{}}
         ) do
      {:ok, {pid, _meta}} -> fun.(pid)
      {:error, reason} -> {:error, reason}
      :ignore -> {:error, :ignore}
    end
  catch
    kind, reason -> {:error, {kind, reason}}
  end

  ## Context plumbing

  defp operable_context(socket) do
    case socket.private[@private_key] do
      %Context{} = context ->
        if operable?(context), do: {:ok, context}, else: :error

      nil ->
        Logger.error(
          "[DurableStash] no stash context on socket — did you `use LiveStash, adapter: DurableStash`?"
        )

        :error
    end
  end

  defp operable?(%Context{storage_key: storage_key, view: view}) do
    is_binary(storage_key) and is_binary(view)
  end

  defp put_context(socket, %Context{} = context) do
    LiveView.put_private(socket, @private_key, context)
  end

  defp derive_storage_key(session, session_id_key, secret) do
    case session do
      %{^session_id_key => sid} when is_binary(sid) ->
        :sha256
        |> :crypto.hash(sid <> secret)
        |> Base.url_encode64(padding: false)

      _session ->
        Logger.error(
          "[DurableStash] no #{inspect(session_id_key)} in the cookie session — " <>
            "add the ensure_session_id plug to your :browser pipeline (see the DurableStash docs)"
        )

        nil
    end
  end

  defp view_name(%{view: view}) when is_atom(view) and not is_nil(view), do: Atom.to_string(view)

  defp view_name(_socket) do
    Logger.error("[DurableStash] socket has no view — cannot scope the stash")
    nil
  end

  defp keys_for_scope(%Context{stored_keys: stored_keys}, scope) do
    for {key, ^scope} <- stored_keys, do: key
  end

  defp recoverable_keys(%Context{reconnected?: true} = context) do
    keys_for_scope(context, :session) ++ keys_for_scope(context, :reconnect)
  end

  defp recoverable_keys(%Context{} = context) do
    keys_for_scope(context, :session)
  end

  # A reconnect (Wi-Fi drop, LiveView crash, deploy) rejoins the same client
  # view, so `_mounts` is positive. A fresh navigation mounts at zero.
  defp reconnected?(socket) do
    match?(%{"_mounts" => mounts} when is_integer(mounts) and mounts > 0, connect_params(socket))
  end

  defp connect_params(socket) do
    if LiveView.connected?(socket) do
      LiveView.get_connect_params(socket) || %{}
    else
      %{}
    end
  rescue
    # get_connect_params is only available while mounting; treat anything
    # else as a fresh mount.
    _error -> %{}
  end

  # Stock LiveStash semantics for :reconnect keys: a fresh mount starts the
  # task over, so the stored values must go — otherwise a crash right after
  # navigating here would resurrect a stale draft.
  defp clear_stale_reconnect_keys(socket, %Context{} = context) do
    reconnect_keys = keys_for_scope(context, :reconnect)

    if reconnect_keys != [] and LiveView.connected?(socket) and not context.reconnected? and
         operable?(context) do
      keys = Enum.map(reconnect_keys, &Atom.to_string/1)

      case call_session(context, &Session.drop(&1, context.view, keys)) do
        :ok ->
          :ok

        {:error, reason} ->
          Logger.error(
            "[DurableStash] clearing reconnect keys failed for #{context.view}: #{inspect(reason)}"
          )
      end
    end

    :ok
  end

  defp fingerprint(encoded) when is_binary(encoded) do
    :crypto.hash(:sha256, encoded)
  end

  ## Option parsing — config errors should be loud, so these raise.

  defp parse_stored_keys!(stored_keys) when is_list(stored_keys) do
    Enum.map(stored_keys, fn
      key when is_atom(key) ->
        {key, :session}

      {key, scope} when is_atom(key) and scope in [:session, :reconnect] ->
        {key, scope}

      {key, :permanent} when is_atom(key) ->
        raise ArgumentError,
              "[DurableStash] the :permanent scope is not yet supported (key #{inspect(key)})"

      other ->
        raise ArgumentError,
              "[DurableStash] invalid stored_keys entry: #{inspect(other)} — " <>
                "expected an atom, `{atom, :session}`, or `{atom, :reconnect}`"
    end)
  end

  defp parse_stored_keys!(other) do
    raise ArgumentError, "[DurableStash] stored_keys must be a list, got: #{inspect(other)}"
  end

  defp parse_migrate!(nil), do: nil
  defp parse_migrate!(migrate) when is_function(migrate, 2), do: migrate

  defp parse_migrate!(other) do
    raise ArgumentError,
          "[DurableStash] :migrate must be a 2-arity function (old_vsn, data), got: #{inspect(other)}"
  end
end