!> @file protocol_codec.f90
!> @brief JSON encode/decode for the Cortex IPC opcodes (Phase G).
!!
!! Uses json-fortran. Mirrors the Python cortex/protocol.py opcode set.

module protocol_codec
  use, intrinsic :: iso_fortran_env, only: int8, real64
  use json_module, only: json_core, json_file, json_value
  use config_schema, only: cfg_kind_int, cfg_kind_real, cfg_kind_logical, cfg_kind_string, cfg_kind_choice, cfg_kind_real3, &
                           config_schema_entry_t, find_config_schema_entry, get_config_schema_entry
  use config_schema_2d, only: schema_lookup_2d
  implicit none
  private

  public :: decode_frame_header
  public :: decode_load_namelist
  public :: decode_advance
  public :: decode_copy_solution
  public :: decode_write_result
  public :: decode_write_checkpoint
  public :: encode_reply_ok
  public :: encode_reply_error
  public :: encode_push
  public :: encode_hello_push
  public :: encode_initialize_reply
  public :: encode_advance_reply
  public :: encode_run_to_end_reply
  public :: encode_copy_solution_reply
  public :: encode_copy_solution_reply_2d
  public :: encode_write_result_reply
  public :: encode_write_checkpoint_reply
  public :: schema_lookup
  public :: encode_progress_tick
  public :: encode_error_push
  public :: decode_keyed_request
  public :: encode_get_integer_reply
  public :: encode_get_real_reply
  public :: encode_get_logical_reply
  public :: encode_get_string_reply
  public :: encode_get_real3_reply
  public :: encode_get_progress_reply
  public :: encode_get_point_count_reply
  public :: decode_set_value
  public :: decode_load_config_inline
  public :: decode_run_to_end

  integer, parameter, public :: protocol_version = 1

  !> Upper bound on the number of primitives a single COPY_SOLUTION
  !! request may list. Any legitimate request names at most the four
  !! 1D Euler primitives (x, rho, u, p); the cap is set generously
  !! above that to reject an untrusted request that lists a huge (or
  !! duplicated) primitive set purely to force a large allocation.
  integer, parameter, public :: MAX_PRIMITIVES = 64

  !> Upper bound on the number of config key-value pairs in a single
  !! LOAD_CONFIG_INLINE request. The schema has 62 valid keys; any
  !! larger count cannot be all-valid, so we reject it before allocating.
  integer, parameter, public :: MAX_CONFIG_ENTRIES = 128

  !> A typed wire value, populated by =decode_set_value= or
  !! =decode_load_config_inline=. Exactly one of the typed fields is
  !! meaningful per entry, identified by =kind= (see =cfg_kind_*= in
  !! =config_schema=). Strings up to 256 chars; longer strings are
  !! rejected by the decoder with =code="type_mismatch"= so the
  !! dispatch arm has a fixed-size scratch buffer to hand to
  !! =solver_session_set_string=.
  type, public :: wire_value_t
    integer :: kind = -1
    integer :: int_value = 0
    real(real64) :: real_value = 0.0_real64
    logical :: log_value = .false.
    character(len=256) :: str_value = ''
    real(real64) :: vec_value(3) = 0.0_real64
  end type wire_value_t

  !> A (key, value) pair produced by =decode_load_config_inline=. The
  !! dispatch arm walks the array and calls =solver_session_set_*= per
  !! entry — transactional only at the codec level (validation runs
  !! before any =session % cfg= mutation; if it fails, the array is
  !! never returned).
  type, public :: wire_entry_t
    character(len=128) :: key = ''
    type(wire_value_t) :: value
  end type wire_entry_t

contains

  ! ----- decoders -------------------------------------------------------

  subroutine decode_frame_header(bytes, op, id)
    integer(int8), intent(in) :: bytes(:)
    character(len=*), intent(out) :: op, id
    type(json_file) :: f
    character(len=:), allocatable :: s, tmp
    logical :: found

    call bytes_to_string(bytes, s)
    call f % initialize()
    call f % deserialize(s)
    call f % get("op", tmp, found)
    if (.not. found) error stop 'protocol_codec: missing "op" in header'
    op = tmp
    call f % get("id", tmp, found)
    if (found) then
      id = tmp
    else
      id = ""
    end if
    call f % destroy()
  end subroutine decode_frame_header

  subroutine decode_load_namelist(bytes, path, case_dir)
    integer(int8), intent(in) :: bytes(:)
    character(len=*), intent(out) :: path
    !> Optional case directory to anchor relative file paths (grid, …) against.
    !! Absent in the request -> returned empty (solver falls back to the
    !! namelist's own directory). Backward-compatible with older clients.
    character(len=*), intent(out), optional :: case_dir
    type(json_file) :: f
    type(json_value), pointer :: args_node
    character(len=:), allocatable :: s, tmp
    logical :: found

    call bytes_to_string(bytes, s)
    call f % initialize()
    call f % deserialize(s)
    ! Try dotted-path syntax first; fall back to two-step get
    call f % get("args.path", tmp, found)
    if (.not. found) then
      call f % get("args", args_node, found)
      if (found) then
        block
          type(json_core) :: c
          call c % initialize()
          call c % get(args_node, "path", tmp, found)
          call c % destroy()
        end block
      end if
    end if
    if (.not. found) error stop 'protocol_codec: LOAD_NAMELIST missing args.path'
    path = tmp

    ! Optional args.case_dir — absent for older clients / CLI-style loads.
    if (present(case_dir)) then
      case_dir = ''
      call f % get("args.case_dir", tmp, found)
      if (found) case_dir = tmp
    end if
    call f % destroy()
  end subroutine decode_load_namelist

  !> Decode the {"args":{"key":<string>}} shape shared by GET_<type>
  !! and SET_<type>. SET adds a `value` field; SET decoders call this
  !! first to extract the key, then decode value separately.
  subroutine decode_keyed_request(bytes, key)
    integer(int8), intent(in) :: bytes(:)
    character(len=*), intent(out) :: key
    type(json_file) :: f
    type(json_value), pointer :: args_node
    character(len=:), allocatable :: s, tmp
    logical :: found

    call bytes_to_string(bytes, s)
    call f % initialize()
    call f % deserialize(s)
    call f % get("args.key", tmp, found)
    if (.not. found) then
      call f % get("args", args_node, found)
      if (found) then
        block
          type(json_core) :: c
          call c % initialize()
          call c % get(args_node, "key", tmp, found)
          call c % destroy()
        end block
      end if
    end if
    if (.not. found) error stop 'protocol_codec: keyed request missing args.key'
    key = tmp
    call f % destroy()
  end subroutine decode_keyed_request

  !> Decode the {"args":{"key":"...","value":<v>}} request shape used
  !! by SET_<type>. Caller passes =expected_kind= obtained from
  !! =schema_lookup= so this routine knows which JSON field to pull.
  !! Returns =ok=.false.= with =code="type_mismatch"= when the value
  !! is missing or the wrong JSON type.
  subroutine decode_set_value(bytes, expected_kind, value, ok, code, message)
    integer(int8), intent(in) :: bytes(:)
    integer, intent(in) :: expected_kind
    type(wire_value_t), intent(out) :: value
    logical, intent(out) :: ok
    character(len=*), intent(out) :: code, message
    type(json_file) :: f
    character(len=:), allocatable :: s, tmp
    real(real64) :: r
    integer :: ival, n
    logical :: lval, found
    type(json_value), pointer :: arr_node

    ok = .false.
    code = "type_mismatch"
    message = ""
    value % kind = expected_kind

    call bytes_to_string(bytes, s)
    call f % initialize()
    call f % deserialize(s)

    select case (expected_kind)
    case (cfg_kind_int)
      call f % get("args.value", ival, found)
      if (.not. found) then
        message = "args.value missing or not integer"
        call f % destroy()
        return
      end if
      value % int_value = ival
    case (cfg_kind_real)
      call f % get("args.value", r, found)
      if (.not. found) then
        message = "args.value missing or not real"
        call f % destroy()
        return
      end if
      value % real_value = r
    case (cfg_kind_logical)
      call f % get("args.value", lval, found)
      if (.not. found) then
        message = "args.value missing or not boolean"
        call f % destroy()
        return
      end if
      value % log_value = lval
    case (cfg_kind_string, cfg_kind_choice)
      call f % get("args.value", tmp, found)
      if (.not. found) then
        message = "args.value missing or not string"
        call f % destroy()
        return
      end if
      if (len_trim(tmp) > len(value % str_value)) then
        message = "args.value string exceeds 256-char wire limit"
        call f % destroy()
        return
      end if
      value % str_value = tmp
    case (cfg_kind_real3)
      call f % get("args.value", arr_node, found)
      if (.not. found) then
        message = "args.value missing"
        call f % destroy()
        return
      end if
      block
        type(json_core) :: c
        call c % initialize()
        call c % info(arr_node, n_children=n)
        if (n /= 3) then
          message = "args.value must be a 3-element array"
          call c % destroy()
          call f % destroy()
          return
        end if
        call c % destroy()
      end block
      call f % get("args.value(1)", value % vec_value(1), found)
      if (found) call f % get("args.value(2)", value % vec_value(2), found)
      if (found) call f % get("args.value(3)", value % vec_value(3), found)
      if (.not. found) then
        message = "args.value real3 elements must all be numeric"
        call f % destroy()
        return
      end if
    case default
      message = "unhandled schema kind"
      call f % destroy()
      return
    end select

    ok = .true.
    code = ""
    call f % destroy()
  end subroutine decode_set_value

  !> Decode {"args":{"config":{key1:v1,key2:v2,...}}}. Walks every key,
  !! resolves it against =config_schema= via =schema_lookup=, and
  !! validates the corresponding value's JSON type. On the first
  !! unknown / mistyped entry returns =ok=.false.= with the offending
  !! key carried in =message=. Bounds checking is delegated to the
  !! caller (the SET dispatch path runs =solver_session_set_*= which
  !! re-validates against schema bounds); this decoder only validates
  !! type-level shape.
  subroutine decode_load_config_inline(bytes, entries, n_entries, ok, code, message, dim, case_dir)
    integer(int8), intent(in) :: bytes(:)
    type(wire_entry_t), allocatable, intent(out) :: entries(:)
    integer, intent(out) :: n_entries
    logical, intent(out) :: ok
    character(len=*), intent(out) :: code, message
    !> Problem dimension the inline config targets. Absent -> 1 (older clients,
    !! 1D-only). Controls which schema (config_schema vs config_schema_2d) the
    !! keys are validated against, so a dim=2 config with 2D-only keys decodes.
    integer, intent(out), optional :: dim
    !> Optional case directory anchoring relative file paths (grid_file, …).
    !! Absent -> returned empty; the caller then leaves paths untouched.
    character(len=*), intent(out), optional :: case_dir
    type(json_file) :: f
    type(json_value), pointer :: config_node, child
    type(json_core) :: c
    character(len=:), allocatable :: s, child_name, str_tmp
    character(len=:), allocatable :: cfgpath, dir_tmp
    integer :: i, n_children, idx, expected_kind, ival, alloc_stat, local_dim
    real(real64) :: r
    logical :: lval, lookup_ok, found

    ok = .false.
    code = "unknown_key"
    message = ""
    n_entries = 0
    if (present(dim)) dim = 1
    if (present(case_dir)) case_dir = ''

    call bytes_to_string(bytes, s)
    call f % initialize()
    call f % deserialize(s)

    ! Optional args.dim / args.case_dir — absent for older 1D-only clients.
    local_dim = 1
    call f % get("args.dim", local_dim, found)
    if (.not. found) local_dim = 1
    ! Only 1 and 2 are valid problem dimensions; reject anything else up front
    ! rather than letting the handler store a bogus session%dim that later
    ! dim-dispatch (initialize, resolve, copy) silently mis-handles.
    if (local_dim /= 1 .and. local_dim /= 2) then
      ok = .false.
      code = "value_out_of_range"
      message = "protocol_codec: LOAD_CONFIG_INLINE args.dim must be 1 or 2"
      call f % destroy()
      return
    end if
    if (present(dim)) dim = local_dim
    if (present(case_dir)) then
      call f % get("args.case_dir", dir_tmp, found)
      if (found) case_dir = dir_tmp
    end if

    call f % get("args.config", config_node, found)
    if (.not. found) then
      message = "args.config missing"
      call f % destroy()
      return
    end if
    call c % initialize()
    call c % info(config_node, n_children=n_children)
    ! n_children is a json member count (>= 0 for a parsed object), so only the
    ! upper bound is reachable; cap it before the allocate to block the
    ! memory-amplification DoS (mirrors decode_copy_solution's MAX_PRIMITIVES).
    if (n_children > MAX_CONFIG_ENTRIES) then
      ok = .false.
      code = 'too_many_entries'
      message = 'protocol_codec: LOAD_CONFIG_INLINE member count exceeds MAX_CONFIG_ENTRIES'
      call c % destroy()
      call f % destroy()
      return
    end if
    allocate (entries(n_children), stat=alloc_stat)
    if (alloc_stat /= 0) then
      ! Status-return (not error stop): a rank-local allocation failure must not
      ! diverge from peers and reintroduce the np>1 hang class this decoder avoids.
      ok = .false.
      code = 'internal_error'
      message = 'protocol_codec: decode_load_config_inline: entries allocation failed'
      call c % destroy()
      call f % destroy()
      return
    end if

    do i = 1, n_children
      call c % get_child(config_node, i, child, found)
      if (.not. found) cycle
      call c % info(child, name=child_name)
      entries(i) % key = child_name
      call schema_lookup_dim(local_dim, trim(child_name), idx, expected_kind, lookup_ok, message)
      if (.not. lookup_ok) then
        code = "unknown_key"
        call c % destroy()
        call f % destroy()
        return
      end if
      entries(i) % value % kind = expected_kind
      cfgpath = "args.config."//trim(child_name)
      select case (expected_kind)
      case (cfg_kind_int)
        call f % get(cfgpath, ival, found)
        if (.not. found) then
          code = "type_mismatch"
          message = "value for "//trim(child_name)//" is not integer"
          call c % destroy(); call f % destroy(); return
        end if
        entries(i) % value % int_value = ival
      case (cfg_kind_real)
        call f % get(cfgpath, r, found)
        if (.not. found) then
          code = "type_mismatch"
          message = "value for "//trim(child_name)//" is not real"
          call c % destroy(); call f % destroy(); return
        end if
        entries(i) % value % real_value = r
      case (cfg_kind_logical)
        call f % get(cfgpath, lval, found)
        if (.not. found) then
          code = "type_mismatch"
          message = "value for "//trim(child_name)//" is not boolean"
          call c % destroy(); call f % destroy(); return
        end if
        entries(i) % value % log_value = lval
      case (cfg_kind_string, cfg_kind_choice)
        call f % get(cfgpath, str_tmp, found)
        if (.not. found) then
          code = "type_mismatch"
          message = "value for "//trim(child_name)//" is not string"
          call c % destroy(); call f % destroy(); return
        end if
        if (len_trim(str_tmp) > len(entries(i) % value % str_value)) then
          code = "type_mismatch"
          message = "value for "//trim(child_name)//" exceeds 256-char wire limit"
          call c % destroy(); call f % destroy(); return
        end if
        entries(i) % value % str_value = str_tmp
      case (cfg_kind_real3)
        call f % get(cfgpath//"(1)", entries(i) % value % vec_value(1), found)
        if (found) call f % get(cfgpath//"(2)", entries(i) % value % vec_value(2), found)
        if (found) call f % get(cfgpath//"(3)", entries(i) % value % vec_value(3), found)
        if (.not. found) then
          code = "type_mismatch"
          message = "value for "//trim(child_name)//" must be a 3-element real array"
          call c % destroy(); call f % destroy(); return
        end if
      case default
        code = "type_mismatch"
        message = "unhandled schema kind for "//trim(child_name)
        call c % destroy(); call f % destroy(); return
      end select
    end do

    n_entries = n_children
    ok = .true.
    code = ""
    call c % destroy()
    call f % destroy()
  end subroutine decode_load_config_inline

  subroutine decode_advance(bytes, max_steps)
    integer(int8), intent(in) :: bytes(:)
    integer, intent(out) :: max_steps
    type(json_file) :: f
    type(json_value), pointer :: args_node
    character(len=:), allocatable :: s
    logical :: found

    call bytes_to_string(bytes, s)
    call f % initialize()
    call f % deserialize(s)
    call f % get("args.max_steps", max_steps, found)
    if (.not. found) then
      call f % get("args", args_node, found)
      if (found) then
        block
          type(json_core) :: c
          call c % initialize()
          call c % get(args_node, "max_steps", max_steps, found)
          call c % destroy()
        end block
      end if
    end if
    if (.not. found) error stop 'protocol_codec: ADVANCE missing args.max_steps'
    call f % destroy()
  end subroutine decode_advance

  subroutine decode_run_to_end(bytes, every_steps, every_seconds)
    integer(int8), intent(in) :: bytes(:)
    integer, intent(out) :: every_steps
    real(real64), intent(out) :: every_seconds
    type(json_file) :: f
    character(len=:), allocatable :: s
    logical :: found
    integer :: tmp_int
    real(real64) :: tmp_real

    every_steps = 100
    every_seconds = 1.0_real64
    call bytes_to_string(bytes, s)
    call f % initialize()
    call f % deserialize(s)
    call f % get("args.progress_interval_steps", tmp_int, found)
    if (found) every_steps = tmp_int
    call f % get("args.progress_interval_seconds", tmp_real, found)
    if (found) every_seconds = tmp_real
    call f % destroy()
  end subroutine decode_run_to_end

  subroutine decode_copy_solution(bytes, n_prims, prim_names)
    integer(int8), intent(in) :: bytes(:)
    integer, intent(out) :: n_prims
    character(len=16), allocatable, intent(out) :: prim_names(:)
    type(json_file) :: f
    character(len=:), allocatable :: s, item
    logical :: found
    integer :: i, alloc_stat

    call bytes_to_string(bytes, s)
    call f % initialize()
    call f % deserialize(s)
    call f % info("args.primitives", n_children=n_prims, found=found)
    if (.not. found) then
      n_prims = 4
      allocate (prim_names(4), stat=alloc_stat)
      if (alloc_stat /= 0) error stop 'protocol_codec: decode_copy_solution: prim_names allocation failed'
      prim_names = [character(len=16) :: "x", "rho", "u", "p"]
    else
      ! Cap the untrusted primitive count before allocating: a request
      ! listing a huge/duplicated primitive set must not force a large
      ! allocation (DoS), and a negative child count is nonsensical.
      if (n_prims < 0 .or. n_prims > MAX_PRIMITIVES) then
        error stop 'protocol_codec: COPY_SOLUTION n_prims out of range (0..MAX_PRIMITIVES)'
      end if
      allocate (prim_names(n_prims), stat=alloc_stat)
      if (alloc_stat /= 0) error stop 'protocol_codec: decode_copy_solution: prim_names allocation failed'
      do i = 1, n_prims
        call f % get("args.primitives("//itoa(i)//")", item, found)
        if (found) prim_names(i) = item
      end do
    end if
    call f % destroy()
  end subroutine decode_copy_solution

  subroutine decode_write_result(bytes, path)
    integer(int8), intent(in) :: bytes(:)
    character(len=*), intent(out) :: path
    type(json_file) :: f
    character(len=:), allocatable :: s, tmp
    logical :: found

    call bytes_to_string(bytes, s)
    call f % initialize()
    call f % deserialize(s)
    call f % get("args.path", tmp, found)
    if (found) then
      path = tmp
    else
      path = ""
    end if
    call f % destroy()
  end subroutine decode_write_result

  subroutine decode_write_checkpoint(bytes, base)
    integer(int8), intent(in) :: bytes(:)
    character(len=*), intent(out) :: base
    type(json_file) :: f
    character(len=:), allocatable :: s, tmp
    logical :: found

    call bytes_to_string(bytes, s)
    call f % initialize()
    call f % deserialize(s)
    call f % get("args.base", tmp, found)
    if (found) then
      base = tmp
    else
      base = ""
    end if
    call f % destroy()
  end subroutine decode_write_checkpoint

  ! ----- encoders -------------------------------------------------------

  subroutine encode_reply_ok(req_id, request_op, header_bytes)
    character(len=*), intent(in) :: req_id, request_op
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', trim(request_op)//'.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_reply_ok

  subroutine encode_reply_error(req_id, request_op, code, message, header_bytes)
    character(len=*), intent(in) :: req_id, request_op, code, message
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, err
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', trim(request_op)//'.reply')
    call c % add(root, 'ok', .false.)
    call c % create_object(err, 'error')
    call c % add(err, 'code', code)
    call c % add(err, 'message', message)
    call c % add(root, err)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_reply_error

  subroutine encode_push(op, header_bytes)
    character(len=*), intent(in) :: op
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, args
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', '')                  ! json-fortran has no null literal helper; "" is fine for stub
    call c % add(root, 'op', op)
    call c % create_object(args, 'args')
    call c % add(root, args)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_push

  subroutine encode_hello_push(build, n_ranks, header_bytes)
    !> HELLO push carries handshake metadata inside ``args`` per IPC
    !  protocol spec §12.7. The plain ``encode_push`` emits an empty
    !  args object, which Cortex rejects with KeyError 'v' when it
    !  reads ``header.args.v`` during the handshake. Worker-side HELLO
    !  must go through this dedicated encoder.
    character(len=*), intent(in) :: build
    integer, intent(in) :: n_ranks
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, args
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', '')
    call c % add(root, 'op', 'HELLO')
    call c % create_object(args, 'args')
    call c % add(args, 'v', protocol_version)
    call c % add(args, 'build', build)
    call c % add(args, 'n_ranks', n_ranks)
    call c % add(root, args)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_hello_push

  !> Encode a PROGRESS_TICK push frame (spec §3.2). Rank 0 only.
  subroutine encode_progress_tick(iteration, sim_time, dt, residual, wallclock_s, header_bytes)
    integer, intent(in) :: iteration
    real(real64), intent(in) :: sim_time, dt, residual, wallclock_s
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, args
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', '')
    call c % add(root, 'op', 'PROGRESS_TICK')
    call c % create_object(args, 'args')
    call c % add(args, 'iteration', iteration)
    call c % add(args, 'sim_time', sim_time)
    call c % add(args, 'dt', dt)
    call c % add(args, 'residual', residual)
    call c % add(args, 'wallclock_s', wallclock_s)
    call c % add(root, args)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_progress_tick

  !> Encode a server-pushed ERROR frame (spec §3.2). Rank 0 only.
  !! Always followed by `parallel_fatal`; the supervisor reads this frame,
  !! then observes socket EOF + non-zero exit and raises WorkerError.
  subroutine encode_error_push(code, message, header_bytes)
    character(len=*), intent(in) :: code, message
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, args
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', '')
    call c % add(root, 'op', 'ERROR')
    call c % create_object(args, 'args')
    call c % add(args, 'code', code)
    call c % add(args, 'message', message)
    call c % add(root, args)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_error_push

  subroutine encode_initialize_reply(req_id, n_global, n_ranks, header_bytes, nx, ny)
    character(len=*), intent(in) :: req_id
    integer, intent(in) :: n_global, n_ranks
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    integer, intent(in), optional :: nx, ny   !< 2D grid shape (dim=2 replies only)
    type(json_core) :: c
    type(json_value), pointer :: root, data
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'INITIALIZE.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % add(data, 'n_global', n_global)
    call c % add(data, 'n_ranks', n_ranks)
    if (present(nx)) call c % add(data, 'nx', nx)
    if (present(ny)) call c % add(data, 'ny', ny)
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_initialize_reply

  subroutine encode_advance_reply(req_id, steps_taken, finished, iteration, &
                                  sim_time, dt, residual, header_bytes)
    character(len=*), intent(in) :: req_id
    integer, intent(in) :: steps_taken, iteration
    logical, intent(in) :: finished
    real(real64), intent(in) :: sim_time, dt, residual
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'ADVANCE.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % add(data, 'steps_taken', steps_taken)
    call c % add(data, 'finished', finished)
    call c % add(data, 'iteration', iteration)
    call c % add(data, 'sim_time', sim_time)
    call c % add(data, 'dt', dt)
    call c % add(data, 'residual', residual)
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_advance_reply

  subroutine encode_run_to_end_reply(req_id, iteration, sim_time, dt, residual, header_bytes)
    character(len=*), intent(in) :: req_id
    integer, intent(in) :: iteration
    real(real64), intent(in) :: sim_time, dt, residual
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'RUN_TO_END.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % add(data, 'iteration', iteration)
    call c % add(data, 'sim_time', sim_time)
    call c % add(data, 'dt', dt)
    call c % add(data, 'residual', residual)
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_run_to_end_reply

  subroutine encode_copy_solution_reply(req_id, n, prim_names, header_bytes)
    character(len=*), intent(in) :: req_id
    integer, intent(in) :: n
    character(len=*), intent(in) :: prim_names(:)
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data, offsets
    character(len=:), allocatable :: s
    integer :: i

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'COPY_SOLUTION.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % add(data, 'n', n)
    call c % create_object(offsets, 'offsets')
    do i = 1, size(prim_names)
      call c % add(offsets, trim(prim_names(i)), (i - 1) * n * 8)
    end do
    call c % add(data, offsets)
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_copy_solution_reply

  !> 2D COPY_SOLUTION reply header: like the 1D encoder but with the grid
  !! shape (`nx`, `ny`) alongside `n = nx*ny`, so the client can reshape each
  !! offset block C-order to (ny, nx). Offsets are byte offsets in request
  !! order, one block of nx*ny little-endian float64 per primitive, packed
  !! i-fastest (x-contiguous) to match solution_gather_2d.write_global_field.
  subroutine encode_copy_solution_reply_2d(req_id, nx, ny, prim_names, header_bytes)
    character(len=*), intent(in) :: req_id
    integer, intent(in) :: nx, ny
    character(len=*), intent(in) :: prim_names(:)
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data, offsets
    character(len=:), allocatable :: s
    integer :: i, n

    n = nx * ny
    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'COPY_SOLUTION.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % add(data, 'n', n)
    call c % add(data, 'nx', nx)
    call c % add(data, 'ny', ny)
    call c % create_object(offsets, 'offsets')
    do i = 1, size(prim_names)
      call c % add(offsets, trim(prim_names(i)), (i - 1) * n * 8)
    end do
    call c % add(data, offsets)
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_copy_solution_reply_2d

  subroutine encode_write_result_reply(req_id, path, header_bytes)
    character(len=*), intent(in) :: req_id, path
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'WRITE_RESULT.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % add(data, 'path', path)
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_write_result_reply

  !> WRITE_CHECKPOINT reply header: `data.files` is the list of file paths
  !! actually written by rank 0 (protocol.md §WRITE_CHECKPOINT). v1 always
  !! writes a single file; `files` is kept as a JSON array for forward
  !! compatibility with a future multi-file checkpoint format.
  subroutine encode_write_checkpoint_reply(req_id, files, header_bytes)
    character(len=*), intent(in) :: req_id
    character(len=*), intent(in) :: files(:)
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data, arr
    character(len=:), allocatable :: s
    integer :: i

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'WRITE_CHECKPOINT.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % create_array(arr, 'files')
    do i = 1, size(files)
      call c % add(arr, '', trim(files(i)))
    end do
    call c % add(data, arr)
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_write_checkpoint_reply

  subroutine encode_get_integer_reply(req_id, value, header_bytes)
    character(len=*), intent(in) :: req_id
    integer, intent(in) :: value
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'GET_INTEGER.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % add(data, 'value', value)
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_get_integer_reply

  subroutine encode_get_real_reply(req_id, value, header_bytes)
    character(len=*), intent(in) :: req_id
    real(real64), intent(in) :: value
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'GET_REAL.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % add(data, 'value', value)
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_get_real_reply

  subroutine encode_get_logical_reply(req_id, value, header_bytes)
    character(len=*), intent(in) :: req_id
    logical, intent(in) :: value
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'GET_LOGICAL.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % add(data, 'value', value)
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_get_logical_reply

  subroutine encode_get_string_reply(req_id, value, header_bytes)
    character(len=*), intent(in) :: req_id, value
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'GET_STRING.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % add(data, 'value', trim(value))
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_get_string_reply

  subroutine encode_get_real3_reply(req_id, value, header_bytes)
    character(len=*), intent(in) :: req_id
    real(real64), intent(in) :: value(3)
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data, arr
    character(len=:), allocatable :: s
    integer :: i

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'GET_REAL3.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % create_array(arr, 'value')
    do i = 1, 3
      call c % add(arr, '', value(i))
    end do
    call c % add(data, arr)
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_get_real3_reply

  subroutine encode_get_progress_reply(req_id, iteration, sim_time, residual, header_bytes)
    character(len=*), intent(in) :: req_id
    integer, intent(in) :: iteration
    real(real64), intent(in) :: sim_time, residual
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'GET_PROGRESS.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % add(data, 'iteration', iteration)
    call c % add(data, 'sim_time', sim_time)
    call c % add(data, 'residual', residual)
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_get_progress_reply

  subroutine encode_get_point_count_reply(req_id, n_global, header_bytes)
    character(len=*), intent(in) :: req_id
    integer, intent(in) :: n_global
    integer(int8), allocatable, intent(out) :: header_bytes(:)
    type(json_core) :: c
    type(json_value), pointer :: root, data
    character(len=:), allocatable :: s

    call c % initialize(no_whitespace=.true.)
    call c % create_object(root, '')
    call c % add(root, 'v', protocol_version)
    call c % add(root, 'id', req_id)
    call c % add(root, 'op', 'GET_POINT_COUNT.reply')
    call c % add(root, 'ok', .true.)
    call c % create_object(data, 'data')
    call c % add(data, 'n_global', n_global)
    call c % add(root, data)
    call c % serialize(root, s)
    call c % destroy(root)
    call string_to_bytes(s, header_bytes)
  end subroutine encode_get_point_count_reply

  ! ----- internals ------------------------------------------------------

  !> Resolve a wire-protocol key against the configuration schema.
  !!
  !! Returns the schema-table index and expected `cfg_kind_*` token.
  !! On unknown key sets ok=.false. with a "unknown config key: ..."
  !! message that callers can forward into an error reply or push.
  subroutine schema_lookup(key, idx, expected_kind, ok, message)
    character(len=*), intent(in) :: key
    integer, intent(out) :: idx
    integer, intent(out) :: expected_kind
    logical, intent(out) :: ok
    character(len=*), intent(out) :: message
    type(config_schema_entry_t) :: entry
    logical :: found

    idx = find_config_schema_entry(trim(key))
    if (idx <= 0) then
      ok = .false.
      expected_kind = -1
      message = "unknown config key: "//trim(key)
      return
    end if
    call get_config_schema_entry(idx, entry, found)
    if (.not. found) then
      ok = .false.
      expected_kind = -1
      message = "schema entry "//trim(key)//" could not be read"
      return
    end if
    expected_kind = entry % value_kind
    ok = .true.
    message = ''
  end subroutine schema_lookup

  !> Resolve a wire key against the schema for `dim` — the single
  !! dimension-selection point for config schemas on the wire path (spec 4.3).
  !! A 3D schema plugs in here as one new case.
  subroutine schema_lookup_dim(dim, key, idx, expected_kind, ok, message)
    integer, intent(in) :: dim
    character(len=*), intent(in) :: key
    integer, intent(out) :: idx
    integer, intent(out) :: expected_kind
    logical, intent(out) :: ok
    character(len=*), intent(out) :: message

    select case (dim)
    case (2)
      call schema_lookup_2d(key, idx, expected_kind, ok, message)
    case (1)
      call schema_lookup(key, idx, expected_kind, ok, message)
    case default
      idx = 0
      expected_kind = -1
      ok = .false.
      write (message, '(a,i0)') 'unsupported dim for config schema: ', dim
    end select
  end subroutine schema_lookup_dim

  pure subroutine bytes_to_string(b, s)
    integer(int8), intent(in) :: b(:)
    character(len=:), allocatable, intent(out) :: s
    integer :: i, alloc_stat
    allocate (character(len=size(b)) :: s, stat=alloc_stat)
    if (alloc_stat /= 0) error stop 'protocol_codec: bytes_to_string: allocation failed'
    do i = 1, size(b)
      s(i:i) = achar(iand(int(b(i)), int(z'FF')))
    end do
  end subroutine bytes_to_string

  pure subroutine string_to_bytes(s, b)
    character(len=*), intent(in) :: s
    integer(int8), allocatable, intent(out) :: b(:)
    integer :: i, alloc_stat
    allocate (b(len(s)), stat=alloc_stat)
    if (alloc_stat /= 0) error stop 'protocol_codec: string_to_bytes: allocation failed'
    do i = 1, len(s)
      b(i) = int(iachar(s(i:i)), int8)
    end do
  end subroutine string_to_bytes

  pure function itoa(n) result(s)
    integer, intent(in) :: n
    character(len=:), allocatable :: s
    character(len=20) :: buf
    write (buf, '(I0)') n
    s = trim(buf)
  end function itoa

end module protocol_codec
