iroh_ex

defmodule IrohEx.NodeManager do
  use GenServer
  alias IrohEx.Native
  require Logger

  @default_log_collector IrohEx.LogCollector
  # @delay_after_batch 500

  def start_link(log_collector \\ @default_log_collector) do

    IO.inspect(log_collector, label: "Log collector impl")

    case GenServer.start_link(__MODULE__, log_collector) do
      # case Sensocto.RegistryUtils.dynamic_start_child(Sensocto.SensorsDynamicSupervisor, __MODULE__, child_spec) do
      {:ok, pid} when is_pid(pid) ->
        Logger.debug("#{__MODULE__} started")
        {:ok, pid}

      {:error, {:already_started, _pid}} ->
        Logger.debug("#{__MODULE__} already started")
        {:ok, :already_started}

      {:error, reason} ->
        Logger.debug("Error starting #{__MODULE__}, reason: #{reason}")
        {:error, reason}
    end
  end

  @impl true
  def init(log_pid) do
    children = [
      # log_collector,
      # log_collector,
      {IrohEx.NodeSupervisor, []}
    ]

    IO.inspect(children, label: "Childspec")
    # IO.inspect(ticket, label: "init Ticket")

    opts = [strategy: :one_for_one, name: Iroh.Supervisor]

    { :ok, sup } = case Supervisor.start_link(children, opts) do
      # case Sensocto.RegistryUtils.dynamic_start_child(Sensocto.SensorsDynamicSupervisor, __MODULE__, child_spec) do
      {:ok, pid} when is_pid(pid) ->
        Logger.debug("#{__MODULE__} Supervisor started")
        {:ok, pid}

      {:error, {:already_started, _pid}} ->
        Logger.debug("#{__MODULE__} Supervisor already started")
        {:ok, :already_started}

      {:error, reason} ->
        Logger.debug("Error starting #{__MODULE__} Supervisor, reason: #{reason}")
        {:error, reason}
    end

    {:ok, %{ sup: sup, log_pid: log_pid}}
  end

  @impl true
  def handle_cast({:send_message, pid, payload}, state) do
    #GenServer.cast(pid, {:send_message, payload })
    Process.send_after(pid, {:send_message, payload }, 0)
    {:noreply, state}
  end

  @impl true
  def handle_call(:get_random_node, _from, state) do
    node = IrohEx.NodeSupervisor.get_children()
    |> Enum.random()
    {:reply, {:ok, node}, state}
  end

  @impl true
  def handle_call({:create_nodes, whale_node_prob, node_count, log_pid, delay_after_batch}, _from, state) do

    # kill all existing nodes
    IrohEx.NodeSupervisor.reset()

    mothership_node_ref = Native.create_node(self())
    ticket = Native.create_ticket(mothership_node_ref)

    # IO.inspect(state, label: "State")
    # 1..node_count
    # |> Enum.map(fn _ ->
    #   Task.async(fn ->
    #     IO.puts("Start node")
    #     #Process.send_after(pid, :create_node, 0)
    #     Task.async(fn -> IrohEx.NodeSupervisor.start_node(ticket, log_pid) end)
    #   end)
    # end)

    Process.send_after(self(), {:create_nodes_batch, whale_node_prob, delay_after_batch}, 0)

    #|> Enum.map(&Task.await/1)
    {:reply, IrohEx.NodeSupervisor.get_children(), state
      |> Map.put(:nodes_to_create, node_count)
      |> Map.put(:mothership_node_ref, mothership_node_ref)
      |> Map.put(:ticket, ticket)
      |> Map.put(:log_pid, log_pid)
    }
  end

  @impl true
  def handle_call(:reset, _from, state) do
    # kill all existing nodes
    IrohEx.NodeSupervisor.reset()

    {:reply, IrohEx.NodeSupervisor.get_children(), state
      |> Map.put(:mothership_node_ref, nil)
      |> Map.put(:ticket, nil)
    }
  end

  @impl true
  def handle_info({:create_nodes_batch, whale_node_prob, delay_ms}, %{nodes_to_create: nodes_to_create, ticket: ticket, log_pid: log_pid} = state) do
    max_concurrency = System.schedulers_online()

    IO.puts("#{__MODULE__} handle_info :create_nodes_batch #{whale_node_prob} #{delay_ms} #{nodes_to_create}")
    batch_size = min(max_concurrency, nodes_to_create)

    stream =
      1..batch_size
      |> Stream.map(fn x ->
        fn ->
          # N% chance
          is_whale_node = case whale_node_prob > 0 do
            true -> :rand.uniform() <= max(1, whale_node_prob) / 100
            false -> false
          end

          IO.puts("#{__MODULE__} Node: #{x} is whale: #{is_whale_node}")

          IrohEx.NodeSupervisor.start_node(is_whale_node, ticket, log_pid)
          IO.puts("Node #{x} started")
          :ok
        end
      end)

    results =
      stream
      |> Task.async_stream(fn action -> action.() end, max_concurrency: max_concurrency)
      |> Enum.to_list()

    started_nodes = Enum.count(results, fn
      {:ok, :ok} -> true
      _ -> false
    end)

    IO.puts("#{__MODULE__} Started #{started_nodes} nodes in batch")

    remaining = nodes_to_create - started_nodes

    if remaining > 0 do
      Process.send_after(self(), {:create_nodes_batch, whale_node_prob, delay_ms}, delay_ms)
    end

    {:noreply, state
      |> Map.put(:nodes_to_create, remaining)
    }
  end

  @impl true
  def handle_info(_msg, state) do
    # IO.inspect(msg, label: "#{__MODULE__} handle_info catchall #{inspect(msg)}")
    {:noreply, state}
  end

end

defmodule IrohEx.NodeSupervisor do
  use DynamicSupervisor
  require Logger

  def start_link(args) do
    DynamicSupervisor.start_link(__MODULE__, args, name: __MODULE__)
  end

  @impl true
  def init(_args) do
    DynamicSupervisor.init(strategy: :one_for_one)
  end

  def start_node(is_whale_node, ticket, log_pid) do
    node_id = "Node:#{:rand.uniform(1_000_000)}"
    child_spec = %{
      id: node_id,
      start: {IrohEx.Node, :start_link, [%{is_whale_node: is_whale_node, ticket: ticket, log_pid: log_pid}]},
      shutdown: 5_000,
      restart: :permanent,
      type: :worker
    }

    # IO.puts("Test")

    case DynamicSupervisor.start_child(__MODULE__, child_spec) do
      # case Sensocto.RegistryUtils.dynamic_start_child(Sensocto.SensorsDynamicSupervisor, __MODULE__, child_spec) do
      {:ok, pid} when is_pid(pid) ->
        Logger.debug("Added node #{node_id}")
        {:ok, pid}

      {:error, {:already_started, _pid}} ->
        Logger.debug("Node already started #{node_id}")
        {:ok, :already_started}

      {:error, reason} ->
        Logger.debug("error adding node #{node_id}: #{inspect(reason)}")
        {:error, reason}
    end
  end

  # def remove_node(sensor_id) do
  #   case Registry.lookup(Sensocto.SensorPairRegistry, sensor_id) do
  #     [{pid, _}] ->
  #       DynamicSupervisor.terminate_child(__MODULE__, pid)
  #       Logger.debug("Stopped sensor #{sensor_id}")
  #       :ok

  #     [] ->
  #       Logger.debug("Error removing sensor: #{sensor_id}")
  #       :error
  #   end
  # end

  def get_children() do
    # children = DynamicSupervisor.which_children(__MODULE__)
    # |> Enum.each(fn {_, pid, _, _} -> IO.inspect(pid) end)
    # IO.inspect(children)
    # children
    children = DynamicSupervisor.which_children(__MODULE__)
            |> Enum.map(fn {_, pid, _, _} -> pid end)

    # IO.inspect(children, label: "Children PIDs")
    children
  end

  def reset() do
    DynamicSupervisor.which_children(__MODULE__)
    |> Enum.each(fn {_, pid, _, _} -> DynamicSupervisor.terminate_child(__MODULE__, pid) end)
  end
end


defmodule IrohEx.Node do
  use GenServer
  require Logger
  alias IrohEx.Native

  def start_link(state) do
    GenServer.start_link(__MODULE__, state)
  end

  @impl true
  def init(%{ticket: ticket, log_pid: log_pid, is_whale_node: is_whale_node} = _state) do
    # Process.flag(:trap_exit, true)
    IO.puts("#{__MODULE__} init")
    Process.send_after(self(), :setup_node, 0)
    {:ok, %{ticket: ticket, is_whale_node: is_whale_node, log_pid: log_pid, peers: [], messages_sent: [], messages_received: []}}
  end

  @impl true
  def handle_info(:setup_node, %{is_whale_node: is_whale_node, ticket: _ticket, log_pid: log_pid } = state) do
    # IO.puts("setup_node ticket: #{inspect(ticket)}, log_pid: #{inspect(log_pid)}")
    #IO.inspect(state, label: "Node init state")
    pid = self()

    IO.puts("#{__MODULE__} before create_node")
    node_ref = Native.create_node(pid, is_whale_node)
    Process.sleep(500)
    node_addr = Native.gen_node_addr(node_ref)
    IO.puts("#{__MODULE__} after create_node #{node_addr}")

    # send immediately, otherwise all :continue callbacks will need to finish before
    Process.send_after(log_pid, {:iroh_node_setup, node_addr}, 0)
    Process.send_after(self(), :connect_node, 0)

    #IO.inspect(node_addr, "Node init")
    {:noreply, state
      |> Map.put(:node_ref, node_ref)
      |> Map.put(:node_addr, node_addr),
      # {:continue, :connect_node}
    }
  end

  def handle_info(:connect_node, %{node_ref: node_ref, node_addr: node_addr, ticket: ticket, log_pid: _log_pid } = state) do
    # IO.puts("connect_node Node addr: #{inspect(node_addr)}")
    IO.puts("#{__MODULE__} before connect #{node_addr}")
    Native.connect_node(node_ref, ticket)
    IO.puts("#{__MODULE__} after connect #{node_addr}")
    Process.sleep(500)
    #IO.inspect(node_addr, "Node init")
    {:noreply, state}
  end

  def handle_info({:send_message, payload }, %{ node_ref: node_ref, node_addr: _node_addr, messages_sent: messages_sent } = state) do
    # IO.puts("#{__MODULE__} handle_info send_message #{node_addr} #{payload}")
    # IO.inspect(state, "handle_info send_message state")
    timestamp = DateTime.utc_now() |> DateTime.to_unix(:millisecond)
    Native.send_message(node_ref, payload)
    {:noreply, state
      |> Map.put(:messages_sent, messages_sent ++ [%{timestamp: timestamp, payload: payload}])
    }
  end

  # IrohEx.Node.handle_cast({:send_message, "Test manager"},
  # def handle_info({:send_message, payload }, %{node_ref: node_ref, log_pid: log_pid, messages_sent: messages_sent} = state) do
  #   timestamp = DateTime.utc_now() |> DateTime.to_unix(:millisecond)
  #   Native.send_message(node_ref, payload)
  #   {:noreply, %{state | messages_sent: messages_sent ++ [%{timestamp: timestamp, payload: payload}]}}
  # end

  @impl true
  def handle_info({:iroh_gossip_message_received, source, msg}, %{ log_pid: log_pid, messages_received: messages_received } = state) do
    # IO.puts("#{__MODULE__} handle_info iroh_gossip_message_received PID alive?: #{Process.alive?(log_pid)} source: #{source} msg: #{msg} #{inspect(Enum.count(messages_received))}")
    timestamp = DateTime.utc_now() |> DateTime.to_unix(:millisecond)

    Process.send_after(log_pid, {:iroh_gossip_message_received, source, msg}, 0)

    {:noreply, state
    |> Map.put(:messages_received, messages_received ++ [%{timestamp: timestamp, payload: msg}])
  }
  end

  def handle_info(msg, %{log_pid: log_pid} = state) do
    # IO.inspect(msg, label: "#{__MODULE__} handle_info catchall #{inspect(msg)}")
    Process.send_after(log_pid, msg, 0)
    {:noreply, state}
  end

  def handle_info(msg, %{ node_ref: _node_ref, node_addr: node_addr } = state) do
    IO.puts("#{__MODULE__} Catchall #{node_addr} #{inspect(msg)}")
    {:noreply, state}
  end

  @impl true
  def handle_call(:get_node_addr, _from, %{node_ref: node_ref } = state) do
    node_addr = Native.gen_node_addr(node_ref)
    {:reply, node_addr, state}
  end

  @impl true
  def handle_call(:report, _from, %{node_addr: _node_addr, peers: peers, messages_sent: messages_sent, messages_received: messages_received} = state) do

    # IO.puts("Hello there")

    # IO.inspect(state, label: "#{__MODULE__} #{node_addr}")

    {:reply, %{
      peers: peers,
      messages_sent_cnt: Enum.count(messages_sent),
      messages_received_cnt: Enum.count(messages_received),
      messages_sent: messages_sent,
      messages_received: messages_received
    }, state}
  end

  @impl true
  def terminate(reason, %{ node_ref: node_ref, node_addr: node_addr } = state) do
    Native.cleanup(node_ref)
    IO.puts("#{__MODULE__} #{node_addr} #{inspect(node_ref)} Going Down: #{inspect(reason)} #{inspect(state)}")
    :normal
  end


end

defmodule IrohEx.LogCollector do
  use GenServer

  def init(init_arg) do
    {:ok, init_arg}
  end

  def start_link(_args) do
    GenServer.start_link(__MODULE__, %{}, name: __MODULE__)
  end

  def handle_info(_msg, state) do
    # IO.inspect(msg, label: "#{__MODULE__} handle_info catchall #{inspect(msg)}")
    {:noreply, state}
  end
end