codex_sdk

defmodule Codex.Files.Registry do
  @moduledoc """
  GenServer-backed manifest that tracks staged file attachments, deduplicates by checksum,
  and prunes expired entries on a schedule. This powers the public `Codex.Files` helpers.
  """

  use GenServer

  alias Codex.Config.Defaults
  alias Codex.Files.Attachment

  @registry __MODULE__
  @manifest_table :codex_files_manifest

  @type stage_opts :: %{
          required(:checksum) => String.t(),
          required(:name) => String.t(),
          required(:persist) => boolean(),
          required(:ttl_ms) => :infinity | pos_integer(),
          required(:size) => non_neg_integer(),
          required(:source_path) => Path.t(),
          required(:destination_path) => Path.t()
        }

  @type stage_request :: {GenServer.from(), stage_opts(), DateTime.t()}

  @type work_item ::
          {:stage, String.t(), stage_opts()}
          | {:force_cleanup, GenServer.from(), DateTime.t()}
          | {:cleanup_tick, DateTime.t()}
          | {:reset, GenServer.from(), Path.t()}
          | :cleanup_orphaned_staging

  @spec start_link(keyword()) :: GenServer.on_start()
  def start_link(opts \\ []) do
    name = Keyword.get(opts, :name, @registry)
    init_opts = Keyword.delete(opts, :name)
    GenServer.start_link(__MODULE__, init_opts, name: name)
  end

  @doc """
  Ensures the registry is running under application supervision.
  """
  @spec ensure_started() :: {:ok, pid()} | {:error, term()}
  def ensure_started do
    with :ok <- ensure_application_started(),
         {:ok, pid} <- fetch_or_restart_registry() do
      {:ok, pid}
    else
      {:error, _} = error -> error
    end
  catch
    :exit, reason -> {:error, reason}
  end

  @doc """
  Inserts or refreshes a staged attachment using the supplied options, returning the
  canonical `Attachment` struct stored in ETS.
  """
  @spec stage(stage_opts()) :: {:ok, Attachment.t()} | {:error, term()}
  def stage(opts) when is_map(opts) do
    GenServer.call(@registry, {:stage, opts})
  end

  @doc """
  Lists all staged attachments currently tracked in the manifest.
  """
  @spec list() :: [Attachment.t()]
  def list do
    GenServer.call(@registry, :list)
  end

  @doc """
  Aggregates counts, sizes, and TTL information for staged attachments.
  """
  @spec metrics() :: map()
  def metrics do
    GenServer.call(@registry, :metrics)
  end

  @doc """
  Triggers an immediate cleanup pass to remove expired attachments.
  """
  @spec force_cleanup() :: :ok | {:error, term()}
  def force_cleanup do
    GenServer.call(@registry, :force_cleanup)
  end

  @doc """
  Clears the manifest and deletes staged files within the provided staging directory.
  """
  @spec reset(Path.t()) :: :ok | {:error, term()}
  def reset(staging_dir) do
    GenServer.call(@registry, {:reset, staging_dir})
  end

  defp ensure_application_started do
    case Application.ensure_all_started(:codex_sdk) do
      {:ok, _} -> :ok
      {:error, {:already_started, _}} -> :ok
      {:error, reason} -> {:error, reason}
    end
  end

  defp fetch_or_restart_registry do
    case wait_for_registry_pid(200) do
      {:ok, pid} -> {:ok, pid}
      :error -> restart_registry_child()
    end
  end

  defp restart_registry_child do
    with true <- is_pid(Process.whereis(Codex.Supervisor)) || {:error, :not_started},
         {:ok, pid} <- restart_registry_child_under_supervisor() do
      {:ok, pid}
    else
      {:error, _} = error -> error
    end
  end

  defp restart_registry_child_under_supervisor do
    case Supervisor.restart_child(Codex.Supervisor, @registry) do
      {:ok, pid} when is_pid(pid) -> {:ok, pid}
      {:error, :running} -> wait_for_registry_or_not_started()
      {:error, :restarting} -> wait_for_registry_or_not_started()
      {:error, :not_found} -> {:error, :not_started}
      {:error, reason} -> {:error, reason}
    end
  end

  defp wait_for_registry_or_not_started do
    case wait_for_registry_pid(200) do
      {:ok, pid} -> {:ok, pid}
      :error -> {:error, :not_started}
    end
  end

  defp wait_for_registry_pid(timeout_ms) when is_integer(timeout_ms) and timeout_ms >= 0 do
    start = System.monotonic_time(:millisecond)
    do_wait_for_registry_pid(start, timeout_ms)
  end

  defp do_wait_for_registry_pid(start, timeout_ms) do
    case Process.whereis(@registry) do
      pid when is_pid(pid) ->
        {:ok, pid}

      nil ->
        if System.monotonic_time(:millisecond) - start >= timeout_ms do
          :error
        else
          Process.sleep(10)
          do_wait_for_registry_pid(start, timeout_ms)
        end
    end
  end

  # GenServer callbacks

  @impl true
  def init(opts) do
    manifest_table = Keyword.get(opts, :manifest_table, @manifest_table)
    file_module = Keyword.get(opts, :file_module, File)

    table =
      :ets.new(manifest_table, [
        :named_table,
        :set,
        :protected,
        read_concurrency: true,
        write_concurrency: true
      ])

    interval = cleanup_interval()

    state = %{
      table: table,
      table_name: manifest_table,
      file_module: file_module,
      cleanup_interval_ms: interval,
      cleanup_timer: schedule_cleanup(interval),
      work_queue: :queue.new(),
      in_flight: nil,
      pending_stage_requests: %{}
    }

    {:ok, state, {:continue, :cleanup_orphaned_staging}}
  end

  @impl true
  def handle_continue(:cleanup_orphaned_staging, state) do
    state = state |> enqueue_work(:cleanup_orphaned_staging) |> maybe_start_work()
    {:noreply, state}
  end

  @impl true
  def handle_call({:stage, opts}, from, state) do
    now = DateTime.utc_now() |> DateTime.truncate(:millisecond)

    case :ets.lookup(state.table, opts.checksum) do
      [{_checksum, %Attachment{} = existing}] ->
        updated = merge_attachment(existing, opts, now)
        :ets.insert(state.table, {updated.checksum, updated})
        emit_staged(updated, cached?: true)
        {:reply, {:ok, updated}, state}

      [] ->
        request = {from, opts, now}

        case Map.get(state.pending_stage_requests, opts.checksum) do
          nil ->
            pending_stage_requests =
              Map.put(state.pending_stage_requests, opts.checksum, [request])

            state =
              %{state | pending_stage_requests: pending_stage_requests}
              |> enqueue_work({:stage, opts.checksum, opts})
              |> maybe_start_work()

            {:noreply, state}

          requests ->
            pending_stage_requests =
              Map.put(state.pending_stage_requests, opts.checksum, [request | requests])

            {:noreply, %{state | pending_stage_requests: pending_stage_requests}}
        end
    end
  end

  def handle_call(:list, from, state) do
    reply_async(from, fn -> list_attachments(state.table) end)
    {:noreply, state}
  end

  def handle_call(:metrics, from, state) do
    reply_async(from, fn ->
      :ets.foldl(&accumulate_metrics/2, initial_metrics(), state.table)
    end)

    {:noreply, state}
  end

  def handle_call(:force_cleanup, from, state) do
    now = DateTime.utc_now() |> DateTime.truncate(:millisecond)
    state = state |> enqueue_work({:force_cleanup, from, now}) |> maybe_start_work()
    {:noreply, state}
  end

  def handle_call({:reset, staging_dir}, from, state) do
    state = state |> enqueue_work({:reset, from, staging_dir}) |> maybe_start_work()
    {:noreply, state}
  end

  @impl true
  def handle_info(:cleanup_tick, state) do
    now = DateTime.utc_now() |> DateTime.truncate(:millisecond)
    state = state |> enqueue_work({:cleanup_tick, now}) |> maybe_start_work()
    {:noreply, state}
  end

  def handle_info(
        {:work_result, pid, result},
        %{in_flight: %{pid: pid, ref: ref, work: work}} = state
      ) do
    Process.demonitor(ref, [:flush])
    state = %{state | in_flight: nil}
    state = handle_completed_work(state, work, result)
    {:noreply, maybe_start_work(state)}
  end

  def handle_info({:work_result, _pid, _result}, state) do
    {:noreply, state}
  end

  def handle_info(
        {:DOWN, ref, :process, _pid, reason},
        %{in_flight: %{ref: ref, work: work}} = state
      ) do
    state = %{state | in_flight: nil}
    state = handle_failed_work(state, work, {:worker_down, reason})
    {:noreply, maybe_start_work(state)}
  end

  def handle_info({:DOWN, _ref, :process, _pid, _reason}, state) do
    {:noreply, state}
  end

  def handle_info(_msg, state) do
    {:noreply, state}
  end

  @impl true
  def terminate(_reason, state) do
    _ = Process.cancel_timer(state.cleanup_timer)

    reply_pending_stage_requests(state.pending_stage_requests, :closed)

    if state.in_flight do
      Process.demonitor(state.in_flight.ref, [:flush])

      if Process.alive?(state.in_flight.pid) do
        Process.exit(state.in_flight.pid, :shutdown)
      end

      reply_work_error(state.in_flight.work, :closed)
    end

    state.work_queue
    |> :queue.to_list()
    |> Enum.each(&reply_work_error(&1, :closed))

    :ok
  end

  defp enqueue_work(state, work) do
    %{state | work_queue: :queue.in(work, state.work_queue)}
  end

  defp maybe_start_work(%{in_flight: nil, work_queue: work_queue} = state) do
    case :queue.out(work_queue) do
      {:empty, _} ->
        state

      {{:value, work}, rest} ->
        start_work(%{state | work_queue: rest}, work)
    end
  end

  defp maybe_start_work(state), do: state

  defp start_work(state, work) do
    parent = self()

    runner = fn ->
      result = perform_work(state, work)
      send(parent, {:work_result, self(), result})
    end

    {:ok, pid} = start_task(runner)
    ref = Process.monitor(pid)
    %{state | in_flight: %{ref: ref, pid: pid, work: work}}
  end

  @spec start_task((-> any())) :: {:ok, pid()}
  defp start_task(fun) do
    case Task.Supervisor.start_child(Codex.TaskSupervisor, fun) do
      {:ok, pid} -> {:ok, pid}
      {:error, {:already_started, pid}} -> {:ok, pid}
      {:error, _} -> Task.start(fun)
    end
  catch
    :exit, _ -> Task.start(fun)
  end

  defp perform_work(%{file_module: file_module}, {:stage, _checksum, opts}) do
    file_module.mkdir_p!(Path.dirname(opts.destination_path))
    file_module.cp!(opts.source_path, opts.destination_path)
    :ok
  end

  defp perform_work(%{file_module: file_module, table: table}, {:force_cleanup, _from, now}) do
    entries = expired_entries(table, now)

    Enum.each(entries, fn {_checksum, attachment} ->
      _ = file_module.rm_rf(attachment.path)
    end)

    {:ok, entries}
  end

  defp perform_work(%{file_module: file_module, table: table}, {:cleanup_tick, now}) do
    entries = expired_entries(table, now)

    Enum.each(entries, fn {_checksum, attachment} ->
      _ = file_module.rm_rf(attachment.path)
    end)

    {:ok, entries}
  end

  defp perform_work(%{file_module: file_module, table: table}, {:reset, _from, staging_dir}) do
    attachments = list_attachments(table)

    Enum.each(attachments, fn attachment ->
      _ = file_module.rm_rf(attachment.path)
    end)

    _ = file_module.rm_rf(staging_dir)
    {:ok, attachments}
  end

  defp perform_work(%{file_module: file_module}, :cleanup_orphaned_staging) do
    _ = file_module.rm_rf(Codex.Files.staging_dir())
    :ok
  end

  defp handle_completed_work(state, work, {:error, reason}) do
    handle_failed_work(state, work, reason)
  end

  defp handle_completed_work(state, work, result) do
    handle_succeeded_work(state, work, result)
  end

  defp handle_succeeded_work(state, {:stage, checksum, _opts}, :ok) do
    {requests, pending_stage_requests} =
      Map.pop(state.pending_stage_requests, checksum, [])

    requests =
      requests
      |> Enum.reverse()

    case requests do
      [] ->
        %{state | pending_stage_requests: pending_stage_requests}

      [{from, opts, now} | rest] ->
        first_attachment = build_attachment(opts, now)
        :ets.insert(state.table, {first_attachment.checksum, first_attachment})
        emit_staged(first_attachment, cached?: false)
        GenServer.reply(from, {:ok, first_attachment})

        _final_attachment =
          Enum.reduce(rest, first_attachment, fn {reply_from, req_opts, req_now}, current ->
            updated = merge_attachment(current, req_opts, req_now)
            :ets.insert(state.table, {updated.checksum, updated})
            emit_staged(updated, cached?: true)
            GenServer.reply(reply_from, {:ok, updated})
            updated
          end)

        %{state | pending_stage_requests: pending_stage_requests}
    end
  end

  defp handle_succeeded_work(state, {:force_cleanup, from, _now}, {:ok, entries}) do
    apply_cleanup_entries(state.table, entries)
    GenServer.reply(from, :ok)
    reschedule_cleanup(state)
  end

  defp handle_succeeded_work(state, {:cleanup_tick, _now}, {:ok, entries}) do
    apply_cleanup_entries(state.table, entries)
    reschedule_cleanup(state)
  end

  defp handle_succeeded_work(state, {:reset, from, _staging_dir}, {:ok, _attachments}) do
    :ets.delete_all_objects(state.table)
    GenServer.reply(from, :ok)
    reschedule_cleanup(state)
  end

  defp handle_succeeded_work(state, :cleanup_orphaned_staging, :ok), do: state

  defp handle_succeeded_work(state, work, other) do
    handle_failed_work(state, work, {:unexpected_work_result, other})
  end

  defp handle_failed_work(state, {:stage, checksum, _opts}, reason) do
    {requests, pending_stage_requests} =
      Map.pop(state.pending_stage_requests, checksum, [])

    Enum.each(requests, fn {from, _opts, _now} ->
      GenServer.reply(from, {:error, reason})
    end)

    %{state | pending_stage_requests: pending_stage_requests}
  end

  defp handle_failed_work(state, {:force_cleanup, from, _now}, reason) do
    GenServer.reply(from, {:error, reason})
    reschedule_cleanup(state)
  end

  defp handle_failed_work(state, {:cleanup_tick, _now}, _reason) do
    reschedule_cleanup(state)
  end

  defp handle_failed_work(state, {:reset, from, _staging_dir}, reason) do
    GenServer.reply(from, {:error, reason})
    reschedule_cleanup(state)
  end

  defp handle_failed_work(state, :cleanup_orphaned_staging, _reason), do: state

  defp reply_work_error({:stage, _checksum, _opts}, _reason), do: :ok

  defp reply_work_error({:force_cleanup, from, _now}, reason) do
    GenServer.reply(from, {:error, reason})
  end

  defp reply_work_error({:reset, from, _staging_dir}, reason) do
    GenServer.reply(from, {:error, reason})
  end

  defp reply_work_error({:cleanup_tick, _now}, _reason), do: :ok
  defp reply_work_error(:cleanup_orphaned_staging, _reason), do: :ok

  defp reply_pending_stage_requests(pending_stage_requests, reason) do
    Enum.each(pending_stage_requests, fn {_checksum, requests} ->
      Enum.each(requests, fn {from, _opts, _now} ->
        GenServer.reply(from, {:error, reason})
      end)
    end)
  end

  defp reply_async(from, fun) when is_function(fun, 0) do
    runner = fn ->
      reply =
        try do
          fun.()
        rescue
          _ -> []
        catch
          :exit, _ -> []
        end

      GenServer.reply(from, reply)
    end

    {:ok, _pid} = start_task(runner)
    :ok
  end

  defp merge_attachment(%Attachment{} = existing, opts, now) do
    ttl_ms = normalize_ttl(existing.ttl_ms, opts.ttl_ms, opts.persist)
    persist = existing.persist || opts.persist

    %Attachment{
      existing
      | persist: persist,
        ttl_ms: ttl_ms,
        inserted_at: now
    }
  end

  defp build_attachment(opts, now) do
    %Attachment{
      id: opts.checksum,
      name: opts.name,
      path: opts.destination_path,
      checksum: opts.checksum,
      size: opts.size,
      persist: opts.persist,
      inserted_at: now,
      ttl_ms: opts.ttl_ms
    }
  end

  defp emit_staged(%Attachment{} = attachment, metadata) do
    :telemetry.execute(
      [:codex, :attachment, :staged],
      %{size_bytes: attachment.size},
      Map.merge(
        %{
          checksum: attachment.checksum,
          name: attachment.name,
          persist?: attachment.persist,
          ttl_ms: attachment.ttl_ms
        },
        Map.new(metadata)
      )
    )
  end

  defp expired_entries(table, now) do
    table
    |> :ets.tab2list()
    |> Enum.filter(fn {_checksum, attachment} ->
      expirable?(attachment) and expired?(attachment, now)
    end)
  end

  defp apply_cleanup_entries(table, entries) do
    Enum.each(entries, fn {checksum, attachment} ->
      :ets.delete(table, checksum)

      :telemetry.execute(
        [:codex, :attachment, :cleaned],
        %{count: 1, bytes: attachment.size},
        %{
          checksum: attachment.checksum,
          name: attachment.name,
          ttl_ms: attachment.ttl_ms
        }
      )
    end)
  end

  defp list_attachments(table) do
    table
    |> :ets.tab2list()
    |> Enum.map(fn {_checksum, attachment} -> attachment end)
  end

  defp expirable?(%Attachment{persist: true}), do: false
  defp expirable?(%Attachment{ttl_ms: :infinity}), do: false
  defp expirable?(%Attachment{}), do: true

  defp expired?(%Attachment{inserted_at: inserted_at, ttl_ms: ttl_ms}, now)
       when is_integer(ttl_ms) do
    DateTime.diff(now, inserted_at, :millisecond) >= ttl_ms
  end

  defp normalize_ttl(current_ttl, new_ttl, persist?) do
    cond do
      persist? -> :infinity
      current_ttl == :infinity -> :infinity
      new_ttl == :infinity -> :infinity
      is_integer(new_ttl) -> new_ttl
      true -> current_ttl
    end
  end

  defp initial_metrics do
    %{
      total_count: 0,
      total_bytes: 0,
      persistent_count: 0,
      persistent_bytes: 0,
      expirable_count: 0
    }
  end

  defp accumulate_metrics({_checksum, attachment}, acc) do
    acc
    |> Map.update!(:total_count, &(&1 + 1))
    |> Map.update!(:total_bytes, &(&1 + attachment.size))
    |> bump_persist_counter(attachment.persist, attachment.size)
  end

  defp bump_persist_counter(acc, true, size) do
    acc
    |> Map.update!(:persistent_count, &(&1 + 1))
    |> Map.update!(:persistent_bytes, &(&1 + size))
  end

  defp bump_persist_counter(acc, false, _size), do: Map.update!(acc, :expirable_count, &(&1 + 1))

  defp cleanup_interval, do: Defaults.attachment_cleanup_interval_ms()

  defp schedule_cleanup(interval_ms) do
    Process.send_after(self(), :cleanup_tick, interval_ms)
  end

  defp reschedule_cleanup(state) do
    _ = Process.cancel_timer(state.cleanup_timer)
    %{state | cleanup_timer: schedule_cleanup(state.cleanup_interval_ms)}
  end
end