defmodule Handoff.ConcurrentExecutionTest do
  # async: false for controlled concurrency testing initially
  # Can be async now
  use ExUnit.Case, async: false

  alias Handoff.DAG
  alias Handoff.DataLocationRegistry
  alias Handoff.DistributedExecutor
  alias Handoff.Function
  alias Handoff.ResultStore

  defp create_simple_dag(dag_id, val_prefix) do
    dag_id
    |> DAG.new()
    |> DAG.add_function(%Function{
      id: :source,
      args: [],
      code: &Elixir.Function.identity/1,
      extra_args: ["#{val_prefix}_source_val"]
    })
    |> DAG.add_function(%Function{
      id: :process,
      args: [:source],
      code: &Handoff.DistributedTestFunctions.g/2,
      extra_args: ["#{val_prefix}_processed"]
    })
  end

  describe "Concurrent Local DAG Execution (Handoff.DistributedExecutor)" do
    test "executes two simple DAGs concurrently with data isolation" do
      dag_id_1 = {self(), 1}
      dag_id_2 = {self(), 2}

      dag1 = create_simple_dag(dag_id_1, "dag1")
      dag2 = create_simple_dag(dag_id_2, "dag2")

      # Execute DAGs. DistributedExecutor.execute is internally async (spawns a Task).
      # The GenServer call to DistributedExecutor itself is synchronous,
      # returning :noreply quickly.
      # The actual result is obtained by awaiting the task implicitly handled
      # by the test if we got {:ok, res}.
      # However, to ensure these start and potentially interleave,
      # we can Task.async them here too.

      task1 = Task.async(fn -> DistributedExecutor.execute(dag1) end)
      task2 = Task.async(fn -> DistributedExecutor.execute(dag2) end)

      {:ok, res1} = Task.await(task1, 15_000)
      {:ok, res2} = Task.await(task2, 15_000)

      # Check DAG 1 results and ID
      assert res1.dag_id == dag_id_1
      assert res1.results[:source] == "dag1_source_val"
      assert res1.results[:process] == ["dag1_source_val", "dag1_processed"]

      # Check DAG 2 results and ID
      assert res2.dag_id == dag_id_2
      assert res2.results[:source] == "dag2_source_val"
      assert res2.results[:process] == ["dag2_source_val", "dag2_processed"]

      # Verify data isolation in ResultStore
      assert {:ok, "dag1_source_val"} = ResultStore.get(dag_id_1, :source)
      assert {:ok, ["dag1_source_val", "dag1_processed"]} = ResultStore.get(dag_id_1, :process)
      assert {:ok, "dag2_source_val"} = ResultStore.get(dag_id_2, :source)

      assert {:ok, ["dag2_source_val", "dag2_processed"]} =
               ResultStore.get(dag_id_2, :process)

      # Verify that an item ID from dag1 cannot be fetched using dag2's ID if
      # that item ID was unique to dag1.
      # Since :source and :process are used in both, their differing values
      # for each dag_id (asserted above) prove isolation.
      # For an explicit non-existent cross-check, imagine dag1 had a unique item:
      # ResultStore.store(dag_id_1, :unique_to_dag1, "unique_val")
      # assert {:error, :not_found} = ResultStore.get(dag_id_2, :unique_to_dag1)
      # This is implicitly covered by the setup clearing both DAG IDs and then only
      # populating as defined.
      # A check for a completely non-existent key in one of the DAGs is fine:
      assert {:error, :not_found} =
               ResultStore.get(dag_id_1, :completely_random_key_not_in_any_dag)
    end
  end

  describe "Concurrent Distributed DAG Execution (Handoff.DistributedExecutor)" do
    defp create_dist_dag(dag_id, val_prefix, node_to_run_on) do
      dag_id
      |> DAG.new()
      |> DAG.add_function(%Function{
        id: :source_op,
        args: [],
        code: &Elixir.Function.identity/1,
        extra_args: ["#{val_prefix}_source_data"],
        # cost: %{cpu: 1}, # Ensure it gets assigned to a node
        # Explicitly assign for testing, though allocator would do it
        node: node_to_run_on
      })
      |> DAG.add_function(%Function{
        id: :process_op,
        args: [:source_op],
        code: &Handoff.DistributedTestFunctions.g/2,
        extra_args: ["#{val_prefix}_processed"],
        # cost: %{cpu: 1},
        node: node_to_run_on
      })
    end

    test "executes two simple DAGs concurrently with data isolation via DistributedExecutor" do
      # Use test PID
      dag_a_id = {self(), 1}
      # Use a spawned PID for the second DAG
      dag_b_id = {self(), 2}

      # For this test, we'll have both run on Node.self() to simplify,
      # focusing on data isolation via DAG ID with DistributedExecutor.
      node_self = Node.self()

      dag_a = create_dist_dag(dag_a_id, "dagA", node_self)
      dag_b = create_dist_dag(dag_b_id, "dagB", node_self)

      task_a = Task.async(fn -> DistributedExecutor.execute(dag_a, []) end)
      task_b = Task.async(fn -> DistributedExecutor.execute(dag_b, []) end)

      res_a = Task.await(task_a, 15_000)
      res_b = Task.await(task_b, 15_000)

      assert {:ok, %{dag_id: ^dag_a_id, results: results_a}} = res_a
      assert {:ok, %{dag_id: ^dag_b_id, results: results_b}} = res_b

      assert results_a[:source_op] == "dagA_source_data"
      assert results_a[:process_op] == ["dagA_source_data", "dagA_processed"]
      assert results_b[:source_op] == "dagB_source_data"
      assert results_b[:process_op] == ["dagB_source_data", "dagB_processed"]

      # Verify ResultStore isolation
      assert {:ok, "dagA_source_data"} = ResultStore.get(dag_a_id, :source_op)
      assert {:ok, "dagB_source_data"} = ResultStore.get(dag_b_id, :source_op)

      # Verify DataLocationRegistry isolation (source_op result is registered)
      assert {:ok, ^node_self} = DataLocationRegistry.lookup(dag_a_id, :source_op)
      assert {:ok, ^node_self} = DataLocationRegistry.lookup(dag_b_id, :source_op)
      # Check that process_op (the final one) location was also registered for its node
      assert {:ok, ^node_self} = DataLocationRegistry.lookup(dag_a_id, :process_op)
      assert {:ok, ^node_self} = DataLocationRegistry.lookup(dag_b_id, :process_op)

      # Ensure dag_a's item isn't findable with dag_b's ID in DLR
      assert {:error, :not_found} =
               DataLocationRegistry.lookup(dag_b_id, :unique_to_a_if_it_existed)
    end
  end

  describe "API Error Handling with DAG IDs" do
    test "Handoff.get_result returns :not_found for non-existent dag_id or item_id" do
      # Use test PID
      dag_id_exists = {self(), 1}
      item_id_exists_for_dag = :item_for_error_test
      # A different PID
      non_existent_dag_id = {self(), 2}

      # Store an item for a known DAG ID to make the DAG ID "valid" in some sense
      ResultStore.store(dag_id_exists, item_id_exists_for_dag, "some_value")
      DataLocationRegistry.register(dag_id_exists, item_id_exists_for_dag, Node.self())

      # Try to get result with a completely non-existent DAG ID
      assert {:error, :timeout} =
               Handoff.get_result(non_existent_dag_id, :any_item, 50)

      # (get_result uses get_with_timeout which tries to fetch,
      # so :timeout is expected if not found after trying)

      # Try to get a non-existent item from an existing DAG ID
      assert {:error, :timeout} =
               Handoff.get_result(dag_id_exists, :non_existent_item_for_error, 50)

      # Try to get an existing item from an existing DAG ID (should succeed)
      assert {:ok, "some_value"} = Handoff.get_result(dag_id_exists, item_id_exists_for_dag, 50)
    end
  end
end
