iroh_ex

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

  @default_log_collector IrohEx.LogCollector

  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

  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

    # log_pid =
    #   Supervisor.which_children(sup)
    #   |> Enum.find(fn {id, _, _, _} -> id == IrohEx.LogCollector end)
    #   |> case do
    #     {_, pid, _, _} -> pid
    #     nil -> raise "LogCollector not found in supervisor"
    #   end

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

  # def handle_info(:create_node, %{ticket: ticket, sup: sup, log_pid: log_pid} = state) do
  #   IrohEx.NodeSupervisor.start_node(ticket, log_pid)
  #   {:noreply, state}
  # end

  def handle_call(:get_random_node, _from, %{ticket: ticket, sup: sup, log_pid: log_pid} = state) do
    node = IrohEx.NodeSupervisor.get_children()
    |> Enum.random()
    {:reply, {:ok, node}, state}
  end

  def handle_cast({:send_message, pid, payload}, %{ticket: ticket, sup: sup, log_pid: log_pid} = state) do
    #GenServer.cast(pid, {:send_message, payload })
    Process.send_after(pid, {:send_message, payload }, 0)
    {:noreply, state}
  end

  def handle_call({:create_nodes, node_count, log_pid}, _from, %{ sup: sup} = state) do

    pid = self()

    # 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)
        IrohEx.NodeSupervisor.start_node(ticket, log_pid)
      end)
    end)
    #|> Enum.map(&Task.await/1)
    {:reply, IrohEx.NodeSupervisor.get_children(), state
      |> Map.put(:mothership_node_ref, mothership_node_ref)
      |> Map.put(:ticket, ticket)
      |> Map.put(:log_pid, log_pid)
    }
  end

  def handle_call(:reset, _from, %{ sup: sup} = 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

  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

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

  def start_node(ticket, log_pid) do
    node_id = "Node:#{:rand.uniform(1_000_000)}"
    child_spec = %{
      id: node_id,
      start: {IrohEx.Node, :start_link, [%{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(%{ticket: ticket, log_pid: log_pid} = state) do
    GenServer.start_link(__MODULE__, state)
  end

  def init(%{ticket: ticket, log_pid: log_pid} = state) do
    Process.flag(:trap_exit, true)
    # IO.puts("#{__MODULE__} init #{inspect(opts)}")
    {:ok, %{ticket: ticket, log_pid: log_pid, peers: [], messages_sent: [], messages_received: []}, {:continue, :setup_node}}
  end

  def handle_continue(:setup_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()

    node_ref = Native.create_node(pid)
    node_addr = Native.gen_node_addr(node_ref)
    Process.send_after(log_pid, {:iroh_node_setup, node_addr}, 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_continue(: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)}")
    Native.connect_node(node_ref, ticket)
    #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_call(:get_node_addr, _from, %{node_ref: node_ref } = state) do
    node_addr = Native.gen_node_addr(node_ref)
    {:reply, node_addr, state}
  end

  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

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

end

defmodule IrohEx.LogCollector do
  use GenServer

  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