require 'json'
require 'fileutils'

# A durable, bounded export spool. Web and supervised workers share the same POSIX volume.
class CountsExportJobs
  class Cancelled < StandardError; end
  class WorkerUnavailable < StandardError; end
  ACTIVE = ['queued', 'running', 'cancelling'].freeze
  HEARTBEAT_SECONDS = 15
  QUEUE_SECONDS = CountsExportDownload::MAX_EXPORT_SECONDS + 30

  def self.worker_id
    @worker_id ||= ".worker-#{ Process.pid }-#{ SecureRandom.hex(8) }"
  end

  def self.pulse
    FileUtils.mkdir_p(directory, mode: 0o700)
    write({ 'id' => worker_id, 'heartbeat_at' => Time.current.to_f })
  end

  def self.worker_available?
    Dir.glob(directory.join('.worker-*.json')).any? do |path|
      JSON.parse(File.read(path)).fetch('heartbeat_at') > Time.current.to_f - HEARTBEAT_SECONDS
    rescue Errno::ENOENT
      false
    end
  end

  def self.expire_stalled(state)
    now = Time.current.to_f
    queued = state['status'] == 'queued'
    stale = if queued
              now - state.fetch('created_at', now) > QUEUE_SECONDS ||
                  (!worker_available? && now - state.fetch('created_at', now) > HEARTBEAT_SECONDS)
            else
              ['running', 'cancelling'].include?(state['status']) &&
                  now - state.fetch('heartbeat_at', now) > HEARTBEAT_SECONDS
            end
    return unless stale

    state.merge!('status' => 'failed', 'error' => queued ? 'worker_unavailable' : 'worker_stalled',
        'updated_at' => now)
    remove_download(state)
    state.delete('token')
    write(state)
  end

  def self.directory = CountsExportDownload.directory.join('jobs')
  def self.verifier = Rails.application.message_verifier('counts_export_job')

  def self.synchronise
    FileUtils.mkdir_p(directory, mode: 0o700)
    File.open(directory.join('.registry.lock'), File::RDWR | File::CREAT, 0o600) do |lock|
      lock.flock(File::LOCK_EX)
      yield
    end
  end

  def self.write(state)
    path = directory.join("#{ state.fetch('id') }.json")
    temporary = "#{ path }.#{ SecureRandom.hex(8) }.tmp"
    File.open(temporary, 'w', 0o600) do |file|
      file.write(JSON.generate(state))
      file.flush
      file.fsync
    end
    File.rename(temporary, path)
  ensure
    File.unlink(temporary) if temporary && File.exist?(temporary)
  end

  def self.read(id)
    JSON.parse(File.read(directory.join("#{ id }.json")))
  rescue Errno::ENOENT
    nil
  end

  def self.remove_download(state)
    name = state['download_file']
    return unless name&.match?(/\Acounts-[a-zA-Z0-9_-]+\.export\z/)

    File.unlink(CountsExportDownload.directory.join(name))
  rescue Errno::ENOENT
    nil
  end

  def self.remove_partial(id)
    Dir.glob(CountsExportDownload.directory.join("counts-#{ id }-*")).each do |path|
      File.unlink(path)
    rescue Errno::ENOENT
      nil
    end
  end

  def self.clean
    Dir.glob(directory.join('*.tmp')).each do |path|
      File.unlink(path) if File.mtime(path) < Time.current - CountsExportDownload::EXPIRY
    end
    Dir.glob(directory.join('*.run')).each do |path|
      id = File.basename(path, '.run')
      next if directory.join("#{ id }.json").exist?

      File.open(path, File::RDWR) do |lock|
        File.unlink(path) if lock.flock(File::LOCK_EX | File::LOCK_NB)
      end
    end
    Dir.glob(directory.join('*.json')).each do |path|
      state = JSON.parse(File.read(path))
      next if state.fetch('expires_at') > Time.current.to_f

      run = directory.join("#{ state.fetch('id') }.run")
      File.open(run, File::RDWR | File::CREAT, 0o600) do |lock|
        unless lock.flock(File::LOCK_EX | File::LOCK_NB)
          state['status'] = 'cancelling'
          write(state)
          next
        end
        remove_download(state)
        remove_partial(state.fetch('id'))
        File.unlink(path)
        File.unlink(run)
      end
    end
  end

  def self.start(payload, owner, language)
    payload = payload.dup
    time_key = payload.key?('definition') ? 'execution_at' : 'at'
    if payload[time_key].nil? || payload[time_key] == ''
      payload[time_key] = Time.current.utc.iso8601(6)
    end
    exporter = CountsExporter.new(payload, language_code: language)
    authority = HistoryAccess.authority
    synchronise do
      clean
      raise WorkerUnavailable unless worker_available?
      active = Dir.glob(directory.join('*.json')).count do |path|
        existing = JSON.parse(File.read(path))
        expire_stalled(existing)
        ACTIVE.include?(existing['status'])
      end
      raise CountsExportDownload::Busy if active >= CountsExportDownload::MAX_CONCURRENT_EXPORTS

      id = SecureRandom.hex(24)
      state = { 'id' => id, 'owner' => owner, 'payload' => payload,
          'language' => language, 'resource' => exporter.resource,
          'format' => exporter.format, 'status' => 'queued',
          'phase' => 'queued', 'created_at' => Time.current.to_f,
          'updated_at' => Time.current.to_f, 'heartbeat_at' => Time.current.to_f,
          'processed' => 0, 'total' => nil,
          **authority,
          'expires_at' => (Time.current + CountsExportDownload::EXPIRY).to_f }
      write(state)
      verifier.generate({ 'id' => id, 'owner' => owner },
          expires_in: CountsExportDownload::EXPIRY, purpose: 'counts_export_job')
    end
  end

  def self.authorised_id(token, owner)
    value = verifier.verified(token, purpose: 'counts_export_job')
    return unless value.is_a?(Hash) && value['owner'] == owner
    return unless value['id']&.match?(/\A[a-f0-9]{48}\z/)

    value['id']
  end

  def self.status(token, owner, cancel: false)
    id = authorised_id(token, owner)
    return unless id
    snapshot = read(id)
    return unless snapshot && snapshot['owner'] == owner
    return unless HistoryAccess.authorised_artifact?(snapshot)

    synchronise do
      state = read(id)
      return unless state && state['owner'] == owner && state['expires_at'] > Time.current.to_f

      expire_stalled(state)

      if cancel && ['queued', 'running', 'ready'].include?(state['status'])
        state['status'] = state['status'] == 'running' ? 'cancelling' : 'cancelled'
        remove_download(state)
        state.delete('token')
        write(state)
      end
      state.slice('status', 'phase', 'processed', 'total', 'format', 'token', 'error',
          'updated_at', 'heartbeat_at')
    end
  end

  def self.checkpoint(id, processed = nil, total = nil, phase: nil)
    synchronise do
      state = read(id)
      unless state && state['status'] == 'running' && state['expires_at'] > Time.current.to_f
        raise Cancelled
      end
      if processed || phase
        state.merge!('processed' => processed, 'total' => total) unless processed.nil?
        state['phase'] = phase if phase
        state['updated_at'] = Time.current.to_f
        write(state)
      end
    end
  end

  def self.work_one
    pulse
    FileUtils.mkdir_p(directory, mode: 0o700)
    Dir.glob(directory.join('*.json')).sort.each do |path|
      id = File.basename(path, '.json')
      File.open(directory.join("#{ id }.run"), File::RDWR | File::CREAT, 0o600) do |lock|
        next unless lock.flock(File::LOCK_EX | File::LOCK_NB)

        state = synchronise do
          candidate = read(id)
          next unless candidate && ACTIVE.include?(candidate['status'])

          if candidate['status'] != 'queued'
            candidate['status'] = candidate['status'] == 'cancelling' ? 'cancelled' : 'failed'
            remove_partial(id)
            candidate['error'] = 'interrupted'
            write(candidate)
            next
          end
          candidate['status'] = 'running'
          candidate['phase'] = 'aggregating'
          candidate['heartbeat_at'] = Time.current.to_f
          candidate['updated_at'] = Time.current.to_f
          write(candidate)
          candidate
        end
        next unless state

        supervise(state)
        return true
      end
    end
    synchronise { clean }
    false
  end

  def self.supervise(state)
    id = state.fetch('id')
    ApplicationRecord.connection_handler.clear_all_connections!
    pid = Process.fork do
      Signal.trap('TERM') { raise Cancelled }
      Signal.trap('INT') { raise Cancelled }
      begin
        perform(state)
      ensure
        ApplicationRecord.connection_handler.clear_all_connections!
        Process.exit!(0)
      end
    end
    started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
    termination = nil
    reaped = false
    begin
      loop do
        if Process.waitpid(pid, Process::WNOHANG)
          reaped = true
          break
        end
        pulse
        now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
        terminal = synchronise do
          current = read(id)
          if current && current['status'] == 'running'
            if now - started > CountsExportDownload::MAX_EXPORT_SECONDS
              current.merge!('status' => 'failed', 'error' => 'generation_timeout')
            end
            current['heartbeat_at'] = Time.current.to_f
            write(current)
          end
          !current || !['running', 'ready'].include?(current['status'])
        end
        begin
          if terminal && termination.nil?
            Process.kill('TERM', pid)
            termination = now
          elsif termination && now - termination > 5
            Process.kill('KILL', pid)
          end
        rescue Errno::ESRCH
          # The child may exit between the nonblocking wait and signal delivery.
          nil
        end
        sleep(1)
      end
    ensure
      unless reaped
        begin
          Process.kill('KILL', pid)
        rescue Errno::ESRCH
          nil
        end
        Process.waitpid(pid)
      end
      synchronise do
        current = read(id)
        if current && current['status'] != 'ready'
          if ['running', 'cancelling'].include?(current['status'])
            current['status'] = current['status'] == 'cancelling' ? 'cancelled' : 'failed'
            current['error'] = 'interrupted' if current['status'] == 'failed'
          end
          remove_download(current)
          remove_partial(id)
          current.delete('token')
          current['updated_at'] = Time.current.to_f
          write(current)
        end
      end
    end
  end

  def self.perform(state)
    actor = state['actor_id'] && User.find_by(id: state['actor_id'])
    actor = nil unless state['admin'] && actor&.admin? && !actor.banned?
    HistoryAccess.set(user: actor) do
      unless HistoryAccess.authorised_artifact?(state)
        synchronise do
          current = read(state.fetch('id'))
          current['status'] = 'failed'
          write(current)
        end
        return
      end
      perform_authorised(state)
    end
  end

  def self.perform_authorised(state)
    id = state.fetch('id')
    token = nil
    begin
      checkpoint(id)
      exporter = CountsExporter.new(state.fetch('payload'), language_code: state.fetch('language'))
      last = Process.clock_gettime(Process::CLOCK_MONOTONIC)
      progress = lambda do |processed, total|
        now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
        if processed.zero? || now - last >= 0.25 || processed == total
          checkpoint(id, processed, total)
          last = now
        end
      end
      check = lambda do
        now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
        if now - last >= 0.25
          checkpoint(id)
          last = now
        end
      end
      phase = ->(value) { checkpoint(id, phase: value) }
      token = CountsExportDownload.prepare(exporter, state.fetch('owner'),
          progress:, phase:, check:, identifier: id)
      file = CountsExportDownload.find(token, state.fetch('owner'))
      raise Cancelled unless file

      synchronise do
        current = read(id)
        if (!current || current['status'] != 'running' ||
            current['expires_at'] <= Time.current.to_f)
          File.unlink(file.fetch('path')) if File.exist?(file.fetch('path'))
          raise Cancelled
        end
        current.merge!('status' => 'ready', 'token' => token,
            'updated_at' => Time.current.to_f,
            'download_file' => File.basename(file.fetch('path')))
        current['total'] = current['processed']
        write(current)
      end
    rescue StandardError => error
      remove_partial(id)
      synchronise do
        current = read(id)
        if current
          unless current['status'] == 'failed'
            current['status'] = error.is_a?(Cancelled) ? 'cancelled' : 'failed'
            current['error'] = error.class.name
          end
          current['updated_at'] = Time.current.to_f
          write(current)
        end
      end
      unless error.is_a?(Cancelled)
        Rails.logger.error("counts_export_job #{ id }: #{ error.class }")
      end
    end
  end
end