scout

defmodule Scout.Store.ETSHardened do
  @moduledoc """
  Hardened ETS storage adapter with proper race condition protection.
  
  FIXES:
  - Protected ETS tables (not public)
  - Study-scoped operations prevent cross-study data corruption  
  - GenServer serialization prevents race conditions
  - Proper error handling with structured returns
  
  Tables:
  - :scout_core_studies -> {study_id, study_map}
  - :scout_core_trials -> {trial_id, study_id, trial_index, status, payload}
  - :scout_core_observations -> {trial_id, bracket, rung, score}
  """

  @behaviour Scout.Store.Adapter
  
  use GenServer
  require Logger

  alias Scout.Util.SafeAtoms

  @tbl_studies :scout_core_studies
  @tbl_trials :scout_core_trials  
  @tbl_observations :scout_core_observations

  ## Client API

  def start_link(_opts) do
    GenServer.start_link(__MODULE__, [], name: __MODULE__)
  end

  def child_spec(opts) do
    %{
      id: __MODULE__,
      start: {__MODULE__, :start_link, [opts]},
      type: :worker,
      restart: :permanent,
      shutdown: 500
    }
  end

  @impl Scout.Store.Adapter
  def put_study(study) when is_map(study) do
    with {:ok, study_id} <- extract_study_id(study),
         :ok <- validate_study(study) do
      GenServer.call(__MODULE__, {:put_study, study_id, study})
    end
  end

  @impl Scout.Store.Adapter
  def set_study_status(study_id, status) when is_binary(study_id) do
    with {:ok, status_atom} <- safe_status_atom(status) do
      GenServer.call(__MODULE__, {:set_study_status, study_id, status_atom})
    end
  end

  @impl Scout.Store.Adapter
  def get_study(study_id) when is_binary(study_id) do
    case :ets.lookup(@tbl_studies, study_id) do
      [{^study_id, study}] -> {:ok, study}
      [] -> :error
    end
  end

  @impl Scout.Store.Adapter
  def add_trial(study_id, trial) when is_binary(study_id) and is_map(trial) do
    with {:ok, trial_id} <- extract_or_generate_trial_id(trial),
         {:ok, index} <- extract_trial_index(trial),
         :ok <- validate_study_exists(study_id) do
      GenServer.call(__MODULE__, {:add_trial, trial_id, study_id, index, trial})
    end
  end

  @impl Scout.Store.Adapter
  def update_trial(study_id, trial_id, updates) 
      when is_binary(study_id) and is_binary(trial_id) and is_map(updates) do
    GenServer.call(__MODULE__, {:update_trial, study_id, trial_id, updates})
  end

  @impl Scout.Store.Adapter
  def record_observation(study_id, trial_id, bracket, rung, score) 
      when is_binary(study_id) and is_binary(trial_id) and is_integer(bracket) and is_integer(rung) and is_number(score) do
    GenServer.call(__MODULE__, {:record_observation, study_id, trial_id, bracket, rung, score})
  end

  @impl Scout.Store.Adapter
  def observations_at_rung(study_id, bracket, rung) 
      when is_binary(study_id) and is_integer(bracket) and is_integer(rung) do
    # Get trials for this study at bracket
    trial_ids = get_trial_ids_for_study_bracket(study_id, bracket)
    
    for trial_id <- trial_ids,
        [{^trial_id, ^bracket, ^rung, score}] <- [:ets.lookup(@tbl_observations, {trial_id, bracket, rung})] do
      {trial_id, score}
    end
  end

  @impl Scout.Store.Adapter
  def list_trials(study_id, filters) when is_binary(study_id) do
    status_filter = Keyword.get(filters, :status)
    limit = Keyword.get(filters, :limit)
    
    trials = :ets.foldl(fn
      {trial_id, ^study_id, index, status, payload}, acc ->
        trial = %{
          id: trial_id,
          study_id: study_id,
          index: index,
          status: status,
          payload: payload
        }
        
        if status_filter == nil or status == status_filter do
          [trial | acc]
        else
          acc
        end
      _, acc -> acc
    end, [], @tbl_trials)
    
    trials = Enum.sort_by(trials, & &1.index)
    
    if limit, do: Enum.take(trials, limit), else: trials
  end

  @impl Scout.Store.Adapter
  def fetch_trial(study_id, trial_id) when is_binary(study_id) and is_binary(trial_id) do
    case :ets.lookup(@tbl_trials, trial_id) do
      [{^trial_id, study_id, index, status, payload}] ->
        {:ok, %{
          id: trial_id,
          study_id: study_id, 
          index: index,
          status: status,
          payload: payload
        }}
      [] -> :error
    end
  end

  @impl Scout.Store.Adapter
  def delete_study(study_id) when is_binary(study_id) do
    GenServer.call(__MODULE__, {:delete_study, study_id})
  end

  @impl Scout.Store.Adapter
  def list_studies() do
    :ets.foldl(fn {_id, study}, acc -> [study | acc] end, [], @tbl_studies)
    |> Enum.reverse()
  end

  @impl Scout.Store.Adapter
  def health_check() do
    try do
      case Process.alive?(Process.whereis(__MODULE__)) do
        true -> :ok
        false -> {:error, :process_dead}
      end
    rescue
      _ -> {:error, :health_check_failed}
    end
  end

  ## GenServer Callbacks

  @impl GenServer
  def init([]) do
    # Create PROTECTED tables (not public - prevents external tampering)
    :ets.new(@tbl_studies, [:set, :protected, :named_table, {:read_concurrency, true}])
    :ets.new(@tbl_trials, [:set, :protected, :named_table, {:read_concurrency, true}])
    :ets.new(@tbl_observations, [:set, :protected, :named_table, {:read_concurrency, true}, {:write_concurrency, true}])
    
    Logger.info("ETS storage initialized with protected tables")
    {:ok, %{}}
  end

  @impl GenServer
  def handle_call({:put_study, study_id, study}, _from, state) do
    try do
      :ets.insert(@tbl_studies, {study_id, study})
      {:reply, :ok, state}
    rescue
      error ->
        Logger.error("Failed to put study #{study_id}: #{inspect(error)}")
        {:reply, {:error, error}, state}
    end
  end

  def handle_call({:set_study_status, study_id, status}, _from, state) do
    case :ets.lookup(@tbl_studies, study_id) do
      [{^study_id, study}] ->
        updated_study = Map.put(study, :status, status)
        :ets.insert(@tbl_studies, {study_id, updated_study})
        {:reply, :ok, state}
      [] ->
        {:reply, {:error, :study_not_found}, state}
    end
  end

  def handle_call({:add_trial, trial_id, study_id, index, trial}, _from, state) do
    # Check for duplicate trial index within study
    existing = :ets.foldl(fn
      {_tid, ^study_id, ^index, _status, _payload}, _acc -> :found
      _, acc -> acc
    end, :not_found, @tbl_trials)
    
    case existing do
      :found ->
        {:reply, {:error, :trial_index_exists}, state}
      :not_found ->
        status = Map.get(trial, :status, :pending)
        payload = Map.drop(trial, [:id, :study_id, :index, :status])
        :ets.insert(@tbl_trials, {trial_id, study_id, index, status, payload})
        {:reply, {:ok, trial_id}, state}
    end
  end

  def handle_call({:update_trial, trial_id, updates}, _from, state) do
    case :ets.lookup(@tbl_trials, trial_id) do
      [{^trial_id, study_id, index, current_status, current_payload}] ->
        new_status = Map.get(updates, :status, current_status)
        new_payload = Map.merge(current_payload, Map.drop(updates, [:status, :id, :study_id, :index]))
        
        :ets.insert(@tbl_trials, {trial_id, study_id, index, new_status, new_payload})
        {:reply, :ok, state}
      [] ->
        {:reply, {:error, :trial_not_found}, state}
    end
  end

  def handle_call({:record_observation, trial_id, bracket, rung, score}, _from, state) do
    # Verify trial exists
    case :ets.lookup(@tbl_trials, trial_id) do
      [{^trial_id, _study_id, _index, _status, _payload}] ->
        observation_key = {trial_id, bracket, rung}
        :ets.insert(@tbl_observations, {observation_key, score})
        {:reply, :ok, state}
      [] ->
        {:reply, {:error, :trial_not_found}, state}
    end
  end

  def handle_call({:delete_study, study_id}, _from, state) do
    try do
      # Get all trials for this study FIRST
      trial_ids = get_trial_ids_for_study(study_id)
      
      # Delete observations for these trials
      for trial_id <- trial_ids do
        :ets.match_delete(@tbl_observations, {{trial_id, :_, :_}, :_})
      end
      
      # Delete trials for this study
      :ets.foldl(fn
        {trial_id, ^study_id, _index, _status, _payload}, _acc ->
          :ets.delete(@tbl_trials, trial_id)
          :ok
        _, acc -> acc
      end, :ok, @tbl_trials)
      
      # Delete the study itself
      :ets.delete(@tbl_studies, study_id)
      
      Logger.info("Deleted study #{study_id} and #{length(trial_ids)} trials")
      {:reply, :ok, state}
    rescue
      error ->
        Logger.error("Failed to delete study #{study_id}: #{inspect(error)}")
        {:reply, {:error, error}, state}
    end
  end

  ## Private Helpers

  defp extract_study_id(%{id: id}) when is_binary(id), do: {:ok, id}
  defp extract_study_id(%{"id" => id}) when is_binary(id), do: {:ok, id}
  defp extract_study_id(_), do: {:error, :missing_study_id}

  defp extract_or_generate_trial_id(%{id: id}) when is_binary(id), do: {:ok, id}
  defp extract_or_generate_trial_id(%{"id" => id}) when is_binary(id), do: {:ok, id}
  defp extract_or_generate_trial_id(_), do: {:ok, UUID.uuid4()}

  defp extract_trial_index(%{index: index}) when is_integer(index) and index >= 0, do: {:ok, index}
  defp extract_trial_index(%{"index" => index}) when is_integer(index) and index >= 0, do: {:ok, index}
  defp extract_trial_index(_), do: {:error, :missing_trial_index}

  defp validate_study(study) when is_map(study) do
    required_fields = [:goal, :search_space]
    missing = Enum.filter(required_fields, &(not Map.has_key?(study, &1)))
    
    case missing do
      [] -> :ok
      fields -> {:error, {:missing_fields, fields}}
    end
  end

  defp validate_study_exists(study_id) do
    case :ets.lookup(@tbl_studies, study_id) do
      [{^study_id, _study}] -> :ok
      [] -> {:error, :study_not_found}
    end
  end

  defp safe_status_atom(status) when is_binary(status) do
    try do
      {:ok, SafeAtoms.status_from_string!(status)}
    rescue
      ArgumentError -> {:error, :invalid_status}
    end
  end
  defp safe_status_atom(status) when is_atom(status) do
    case status in SafeAtoms.valid_statuses() do
      true -> {:ok, status}
      false -> {:error, :invalid_status}
    end
  end

  defp get_trial_ids_for_study(study_id) do
    :ets.foldl(fn
      {trial_id, ^study_id, _index, _status, _payload}, acc -> [trial_id | acc]
      _, acc -> acc
    end, [], @tbl_trials)
  end

  defp get_trial_ids_for_study_bracket(study_id, bracket) do
    # For now, return all trials for study - bracket filtering can be added to trial payload
    get_trial_ids_for_study(study_id)
  end
end