!> @file cortex_run_attached.f90
!> @brief Phase G: rank-0 dispatch loop + peer command loop for the
!!        attached (Cortex-driven) execution mode.
!!
!! Spec §11.1, §11.2, §11.3. Identical dispatch on every rank so
!! collective MPI ops in solver_session_* line up. Only rank 0 owns
!! the socket and writes replies.

module cortex_run_attached_mod
  use, intrinsic :: iso_fortran_env, only: int8, int32, int64, real64
  use precision, only: wp
  use mpi_runtime, only: my_rank, n_ranks, mpi_world, parallel_fatal
  use logger, only: log_warn
#ifdef CFD_SOLVER_LEGACY_MPI
  use mpi, only: MPI_Bcast, MPI_INTEGER, MPI_BYTE
#else
  use mpi_f08, only: MPI_Bcast, MPI_INTEGER, MPI_BYTE
#endif
  use cortex_link, only: cortex_open, cortex_send_frame, cortex_recv_frame, cortex_close_fd, &
                         MAX_SEND_FRAME_BYTES
  use protocol_codec, only: decode_frame_header, decode_load_namelist, decode_advance, &
                            decode_copy_solution, decode_write_result, decode_write_checkpoint, &
                            decode_keyed_request, decode_run_to_end, &
                            encode_reply_ok, encode_reply_error, encode_push, encode_hello_push, &
                            encode_initialize_reply, encode_advance_reply, &
                            encode_run_to_end_reply, &
                            encode_copy_solution_reply, encode_copy_solution_reply_2d, &
                            encode_write_result_reply, encode_write_checkpoint_reply, &
                            encode_get_integer_reply, encode_get_real_reply, &
                            encode_get_logical_reply, encode_get_string_reply, &
                            encode_get_real3_reply, &
                            encode_get_progress_reply, encode_get_point_count_reply, &
                            encode_progress_tick, &
                            wire_value_t, wire_entry_t, &
                            decode_set_value, decode_load_config_inline, &
                            schema_lookup
  use config_schema, only: cfg_kind_int, cfg_kind_real, cfg_kind_logical, &
                           cfg_kind_string, cfg_kind_choice, cfg_kind_real3
  use config_schema_2d, only: schema_lookup_2d
  use solver_session, only: solver_session_t, solver_progress_t, &
                            solver_session_create, solver_session_destroy, &
                            solver_session_load_namelist, solver_session_initialize, &
                            solver_session_advance, solver_session_run_to_end, &
                            solver_session_get_progress, solver_session_get_point_count, &
                            solver_session_get_global_point_count, &
                            solver_session_copy_solution, solver_session_copy_global_solution, &
                            solver_session_get_global_grid_shape, &
                            solver_session_copy_global_solution_2d, &
                            solver_session_write_result, &
                            solver_session_write_checkpoint, solver_status_ok, &
                            solver_status_invalid_argument, solver_status_invalid_state, &
                            solver_status_config_error, solver_status_io_error, &
                            solver_session_get_integer, solver_session_get_real, &
                            solver_session_get_logical, solver_session_get_string, &
                            solver_session_get_real3, &
                            solver_session_set_integer, solver_session_set_real, &
                            solver_session_set_logical, solver_session_set_string, &
                            solver_session_set_real3, &
                            solver_session_prepare_inline, solver_session_resolve_case_paths
  use version_info, only: solver_build
  implicit none
  private

  public :: cortex_run_attached

  integer, save :: g_fd = -1

contains

  subroutine cortex_run_attached(uri, worker_name)
    character(len=*), intent(in) :: uri
    character(len=*), intent(in), optional :: worker_name

    character(len=32) :: wname
    type(solver_session_t) :: session
    integer(int8), allocatable :: header(:), payload(:)
    integer(int8), allocatable :: reply_header(:), reply_payload(:)
    integer :: ierr
    character(len=64) :: op, id
    logical :: done

    ! Suppress unused-variable warning for ierr (used inside bcast_bytes).
    ierr = 0

    call solver_session_create(session)
    g_fd = -1
    wname = 'euler_1d'
    if (present(worker_name)) wname = worker_name
    if (my_rank() == 0) then
      call cortex_open(uri, g_fd)
      call send_hello(g_fd, trim(wname))
    end if

    done = .false.
    do while (.not. done)
      if (my_rank() == 0) then
        call cortex_recv_frame(g_fd, header, payload)
      end if
      call bcast_bytes(header)
      call bcast_bytes(payload)
      call decode_frame_header(header, op, id)

      ! The request payload frame is read and broadcast above so the wire
      ! stream stays in sync, but no current opcode consumes a request
      ! payload (all request args travel in the header), so it is not passed
      ! to the dispatcher. Re-add it here and below if a future opcode needs it.
      call dispatch_opcode(session, op, id, header, &
                           reply_header, reply_payload, done)

      if (my_rank() == 0 .and. allocated(reply_header)) then
        if (.not. allocated(reply_payload)) then
          allocate (reply_payload(0), stat=ierr)
          if (ierr /= 0) error stop 'cortex_run_attached: empty reply payload allocation failed'
        end if
        call cortex_send_frame(g_fd, reply_header, reply_payload)
      end if

      if (allocated(reply_header)) deallocate (reply_header, stat=ierr)
      if (allocated(reply_payload)) deallocate (reply_payload, stat=ierr)
    end do

    if (my_rank() == 0 .and. g_fd >= 0) call cortex_close_fd(g_fd)
    call solver_session_destroy(session)
  end subroutine cortex_run_attached

  subroutine dispatch_opcode(session, op, id, header, &
                             reply_header, reply_payload, done)
    type(solver_session_t), intent(inout) :: session
    character(len=*), intent(in) :: op, id
    integer(int8), intent(in) :: header(:)
    integer(int8), allocatable, intent(out) :: reply_header(:), reply_payload(:)
    logical, intent(inout) :: done

    character(len=512) :: path, message
    character(len=512) :: written_path
    character(len=512) :: case_dir
    integer :: status, max_steps, steps_taken, n_global
    logical :: finished
    type(solver_progress_t) :: progress

    select case (trim(op))
    case ("LOAD_NAMELIST")
      call decode_load_namelist(header, path, case_dir)
      call solver_session_load_namelist(session, trim(path), status, message, &
                                        case_dir=trim(case_dir))
      if (status /= solver_status_ok) then
        if (my_rank() == 0) call encode_reply_error(id, op, "CONFIG_ERROR", trim(message), reply_header)
      else
        if (my_rank() == 0) call encode_reply_ok(id, op, reply_header)
      end if

    case ("INITIALIZE")
      call solver_session_initialize(session, status, message)
      if (op_failed(id, op, status, message, reply_header)) return
      ! Report the *global* point count, not this rank's local slice;
      ! the GUI sizes whole-domain buffers from this value.
      call solver_session_get_global_point_count(session, n_global, status, message)
      if (op_failed(id, op, status, message, reply_header)) return
      if (session % dim == 2) then
        block
          integer :: nx_g, ny_g
          call solver_session_get_global_grid_shape(session, nx_g, ny_g, status, message)
          if (op_failed(id, op, status, message, reply_header)) return
          if (my_rank() == 0) call encode_initialize_reply(id, n_global=n_global, &
                                                           n_ranks=n_ranks(), &
                                                           header_bytes=reply_header, &
                                                           nx=nx_g, ny=ny_g)
        end block
      else
        if (my_rank() == 0) call encode_initialize_reply(id, n_global=n_global, &
                                                         n_ranks=n_ranks(), &
                                                         header_bytes=reply_header)
      end if

    case ("ADVANCE")
      call decode_advance(header, max_steps)
      call solver_session_advance(session, max_steps, steps_taken, finished, status, message)
      if (op_failed(id, op, status, message, reply_header)) return
      call solver_session_get_progress(session, progress)
      if (my_rank() == 0) call encode_advance_reply(id, steps_taken=steps_taken, &
                                                    finished=finished, &
                                                    iteration=progress % iteration, &
                                                    sim_time=progress % sim_time, &
                                                    dt=progress % dt, &
                                                    residual=progress % residual, &
                                                    header_bytes=reply_header)

    case ("RUN_TO_END")
      block
        integer :: tick_every_steps
        real(wp) :: tick_every_seconds
        call decode_run_to_end(header, tick_every_steps, tick_every_seconds)
        if (my_rank() == 0) then
          call solver_session_run_to_end(session, status, message, &
                                         on_progress=emit_progress_tick, &
                                         every_steps=tick_every_steps, &
                                         every_seconds=tick_every_seconds)
        else
          call solver_session_run_to_end(session, status, message)
        end if
        if (op_failed(id, op, status, message, reply_header)) return
        call solver_session_get_progress(session, progress)
        if (my_rank() == 0) call encode_run_to_end_reply(id, &
                                                         iteration=progress % iteration, &
                                                         sim_time=progress % sim_time, &
                                                         dt=progress % dt, &
                                                         residual=progress % residual, &
                                                         header_bytes=reply_header)
      end block

    case ("WRITE_RESULT")
      call decode_write_result(header, path)
      if (len_trim(path) == 0) then
        if (session % dim == 2) then
          path = "result_2d.dat"
        else
          path = "result.dat"
        end if
      end if
      call solver_session_write_result(session, filename=trim(path), status=status, message=message)
      if (status /= solver_status_ok) then
        if (my_rank() == 0) call encode_reply_error(id, op, "IO_ERROR", trim(message), reply_header)
      else
        if (my_rank() == 0) call encode_write_result_reply(id, trim(path), reply_header)
      end if

    case ("COPY_SOLUTION")
      call handle_copy_solution(session, id, op, header, reply_header, reply_payload, status, message)

    case ("WRITE_CHECKPOINT")
      call decode_write_checkpoint(header, path)
      if (len_trim(path) == 0) path = "checkpoint"
      call solver_session_write_checkpoint(session, trim(path), status, message, written_path)
      if (status /= solver_status_ok) then
        if (my_rank() == 0) call encode_reply_error(id, op, "IO_ERROR", trim(message), reply_header)
      else
        if (my_rank() == 0) call encode_write_checkpoint_reply(id, [written_path], reply_header)
      end if

    case ("SET_INTEGER")
      block
        character(len=128) :: key
        integer :: idx, expected_kind
        logical :: dec_ok
        character(len=256) :: lookup_msg
        type(wire_value_t) :: wv
        character(len=64) :: code
        character(len=256) :: dec_msg
        call decode_keyed_request(header, key)
        call session_schema_lookup(session, trim(key), idx, expected_kind, dec_ok, lookup_msg)
        if (.not. dec_ok) then
          if (my_rank() == 0) call encode_reply_error(id, op, "unknown_key", trim(lookup_msg), reply_header)
        else if (reject_kind_mismatch(op, id, trim(key), cfg_kind_int, expected_kind, reply_header)) then
          ! type_mismatch reply already encoded by reject_kind_mismatch
        else
          call decode_set_value(header, expected_kind, wv, dec_ok, code, dec_msg)
          if (.not. dec_ok) then
            if (my_rank() == 0) call encode_reply_error(id, op, trim(code), trim(dec_msg), reply_header)
          else
            call solver_session_set_integer(session, trim(key), wv % int_value, status, message)
            if (status /= solver_status_ok) then
              if (my_rank() == 0) call encode_reply_error(id, op, "value_out_of_range", trim(message), reply_header)
            else
              if (my_rank() == 0) call encode_reply_ok(id, op, reply_header)
            end if
          end if
        end if
      end block

    case ("SET_REAL")
      block
        character(len=128) :: key
        integer :: idx, expected_kind
        logical :: dec_ok
        character(len=256) :: lookup_msg
        type(wire_value_t) :: wv
        character(len=64) :: code
        character(len=256) :: dec_msg
        call decode_keyed_request(header, key)
        call session_schema_lookup(session, trim(key), idx, expected_kind, dec_ok, lookup_msg)
        if (.not. dec_ok) then
          if (my_rank() == 0) call encode_reply_error(id, op, "unknown_key", trim(lookup_msg), reply_header)
        else if (reject_kind_mismatch(op, id, trim(key), cfg_kind_real, expected_kind, reply_header)) then
          ! type_mismatch reply already encoded by reject_kind_mismatch
        else
          call decode_set_value(header, expected_kind, wv, dec_ok, code, dec_msg)
          if (.not. dec_ok) then
            if (my_rank() == 0) call encode_reply_error(id, op, trim(code), trim(dec_msg), reply_header)
          else
            call solver_session_set_real(session, trim(key), wv % real_value, status, message)
            if (status /= solver_status_ok) then
              if (my_rank() == 0) call encode_reply_error(id, op, "value_out_of_range", trim(message), reply_header)
            else
              if (my_rank() == 0) call encode_reply_ok(id, op, reply_header)
            end if
          end if
        end if
      end block

    case ("SET_LOGICAL")
      block
        character(len=128) :: key
        integer :: idx, expected_kind
        logical :: dec_ok
        character(len=256) :: lookup_msg
        type(wire_value_t) :: wv
        character(len=64) :: code
        character(len=256) :: dec_msg
        call decode_keyed_request(header, key)
        call session_schema_lookup(session, trim(key), idx, expected_kind, dec_ok, lookup_msg)
        if (.not. dec_ok) then
          if (my_rank() == 0) call encode_reply_error(id, op, "unknown_key", trim(lookup_msg), reply_header)
        else if (reject_kind_mismatch(op, id, trim(key), cfg_kind_logical, expected_kind, reply_header)) then
          ! type_mismatch reply already encoded by reject_kind_mismatch
        else
          call decode_set_value(header, expected_kind, wv, dec_ok, code, dec_msg)
          if (.not. dec_ok) then
            if (my_rank() == 0) call encode_reply_error(id, op, trim(code), trim(dec_msg), reply_header)
          else
            call solver_session_set_logical(session, trim(key), wv % log_value, status, message)
            if (status /= solver_status_ok) then
              if (my_rank() == 0) call encode_reply_error(id, op, "value_out_of_range", trim(message), reply_header)
            else
              if (my_rank() == 0) call encode_reply_ok(id, op, reply_header)
            end if
          end if
        end if
      end block

    case ("SET_STRING")
      block
        character(len=128) :: key
        integer :: idx, expected_kind
        logical :: dec_ok
        character(len=256) :: lookup_msg
        type(wire_value_t) :: wv
        character(len=64) :: code
        character(len=256) :: dec_msg
        call decode_keyed_request(header, key)
        call session_schema_lookup(session, trim(key), idx, expected_kind, dec_ok, lookup_msg)
        if (.not. dec_ok) then
          if (my_rank() == 0) call encode_reply_error(id, op, "unknown_key", trim(lookup_msg), reply_header)
        else if (reject_kind_mismatch(op, id, trim(key), cfg_kind_string, expected_kind, reply_header)) then
          ! type_mismatch reply already encoded by reject_kind_mismatch
        else
          call decode_set_value(header, expected_kind, wv, dec_ok, code, dec_msg)
          if (.not. dec_ok) then
            if (my_rank() == 0) call encode_reply_error(id, op, trim(code), trim(dec_msg), reply_header)
          else
            call solver_session_set_string(session, trim(key), trim(wv % str_value), status, message)
            if (status /= solver_status_ok) then
              if (my_rank() == 0) call encode_reply_error(id, op, "value_out_of_range", trim(message), reply_header)
            else
              if (my_rank() == 0) call encode_reply_ok(id, op, reply_header)
            end if
          end if
        end if
      end block

    case ("SET_REAL3")
      block
        character(len=128) :: key
        integer :: idx, expected_kind
        logical :: dec_ok
        character(len=256) :: lookup_msg
        type(wire_value_t) :: wv
        character(len=64) :: code
        character(len=256) :: dec_msg
        call decode_keyed_request(header, key)
        call session_schema_lookup(session, trim(key), idx, expected_kind, dec_ok, lookup_msg)
        if (.not. dec_ok) then
          if (my_rank() == 0) call encode_reply_error(id, op, "unknown_key", trim(lookup_msg), reply_header)
        else if (reject_kind_mismatch(op, id, trim(key), cfg_kind_real3, expected_kind, reply_header)) then
          ! type_mismatch reply already encoded by reject_kind_mismatch
        else
          call decode_set_value(header, expected_kind, wv, dec_ok, code, dec_msg)
          if (.not. dec_ok) then
            if (my_rank() == 0) call encode_reply_error(id, op, trim(code), trim(dec_msg), reply_header)
          else
            call solver_session_set_real3(session, trim(key), wv % vec_value, status, message)
            if (status /= solver_status_ok) then
              if (my_rank() == 0) call encode_reply_error(id, op, "value_out_of_range", trim(message), reply_header)
            else
              if (my_rank() == 0) call encode_reply_ok(id, op, reply_header)
            end if
          end if
        end if
      end block

    case ("LOAD_CONFIG_INLINE")
      block
        type(wire_entry_t), allocatable :: entries(:)
        integer :: i, n_entries, dstat, cfg_dim
        logical :: dec_ok
        character(len=64) :: code
        character(len=256) :: dec_msg
        character(len=512) :: cfg_case_dir
        call decode_load_config_inline(header, entries, n_entries, dec_ok, code, dec_msg, &
                                       dim=cfg_dim, case_dir=cfg_case_dir)
        if (.not. dec_ok) then
          if (my_rank() == 0) call encode_reply_error(id, op, trim(code), trim(dec_msg), reply_header)
        else
          ! Fix the target dimension and case anchor BEFORE applying entries, so
          ! the dim-aware SET_* setters land on the 1D/2D config the keys came
          ! from. Collective-safe: every rank decoded the same dim/case_dir.
          call solver_session_prepare_inline(session, cfg_dim, trim(cfg_case_dir))
          status = solver_status_ok
          do i = 1, n_entries
            select case (entries(i) % value % kind)
            case (cfg_kind_int)
              call solver_session_set_integer(session, trim(entries(i) % key), entries(i) % value % int_value, status, message)
            case (cfg_kind_real)
              call solver_session_set_real(session, trim(entries(i) % key), entries(i) % value % real_value, status, message)
            case (cfg_kind_logical)
              call solver_session_set_logical(session, trim(entries(i) % key), entries(i) % value % log_value, status, message)
            case (cfg_kind_string, cfg_kind_choice)
              call solver_session_set_string(session, trim(entries(i) % key), trim(entries(i) % value % str_value), status, message)
            case (cfg_kind_real3)
              call solver_session_set_real3(session, trim(entries(i) % key), entries(i) % value % vec_value, status, message)
            end select
            if (status /= solver_status_ok) exit
          end do
          if (status /= solver_status_ok) then
            if (my_rank() == 0) call encode_reply_error(id, op, "value_out_of_range", trim(message), reply_header)
          else
            ! Anchor the relative grid_file against the case dir now the entries
            ! have landed; the inline path bypasses read_config where this
            ! resolution otherwise lives.
            call solver_session_resolve_case_paths(session)
            if (my_rank() == 0) call encode_reply_ok(id, op, reply_header)
          end if
        end if
        if (allocated(entries)) deallocate (entries, stat=dstat)
      end block

    case ("GET_INTEGER")
      if (my_rank() == 0) then
        block
          character(len=128) :: key
          integer :: int_value, idx, expected_kind
          logical :: dec_ok
          character(len=256) :: lookup_msg
          call decode_keyed_request(header, key)
          call session_schema_lookup(session, trim(key), idx, expected_kind, dec_ok, lookup_msg)
          if (.not. dec_ok) then
            call encode_keyed_failure(op, id, "unknown_key", trim(lookup_msg), reply_header)
          else if (reject_kind_mismatch(op, id, trim(key), cfg_kind_int, expected_kind, reply_header)) then
            ! type_mismatch reply already encoded by reject_kind_mismatch
          else
            call solver_session_get_integer(session, trim(key), int_value, status, message)
            if (status /= solver_status_ok) then
              call encode_keyed_failure(op, id, "unknown_key", trim(message), reply_header)
            else
              call encode_get_integer_reply(id, int_value, reply_header)
            end if
          end if
        end block
      end if

    case ("GET_REAL")
      if (my_rank() == 0) then
        block
          character(len=128) :: key
          real(real64) :: real_value
          integer :: idx, expected_kind
          logical :: dec_ok
          character(len=256) :: lookup_msg
          call decode_keyed_request(header, key)
          call session_schema_lookup(session, trim(key), idx, expected_kind, dec_ok, lookup_msg)
          if (.not. dec_ok) then
            call encode_keyed_failure(op, id, "unknown_key", trim(lookup_msg), reply_header)
          else if (reject_kind_mismatch(op, id, trim(key), cfg_kind_real, expected_kind, reply_header)) then
            ! type_mismatch reply already encoded by reject_kind_mismatch
          else
            call solver_session_get_real(session, trim(key), real_value, status, message)
            if (status /= solver_status_ok) then
              call encode_keyed_failure(op, id, "unknown_key", trim(message), reply_header)
            else
              call encode_get_real_reply(id, real_value, reply_header)
            end if
          end if
        end block
      end if

    case ("GET_LOGICAL")
      if (my_rank() == 0) then
        block
          character(len=128) :: key
          logical :: log_value
          integer :: idx, expected_kind
          logical :: dec_ok
          character(len=256) :: lookup_msg
          call decode_keyed_request(header, key)
          call session_schema_lookup(session, trim(key), idx, expected_kind, dec_ok, lookup_msg)
          if (.not. dec_ok) then
            call encode_keyed_failure(op, id, "unknown_key", trim(lookup_msg), reply_header)
          else if (reject_kind_mismatch(op, id, trim(key), cfg_kind_logical, expected_kind, reply_header)) then
            ! type_mismatch reply already encoded by reject_kind_mismatch
          else
            call solver_session_get_logical(session, trim(key), log_value, status, message)
            if (status /= solver_status_ok) then
              call encode_keyed_failure(op, id, "unknown_key", trim(message), reply_header)
            else
              call encode_get_logical_reply(id, log_value, reply_header)
            end if
          end if
        end block
      end if

    case ("GET_STRING")
      if (my_rank() == 0) then
        block
          character(len=128) :: key
          character(len=256) :: str_value
          integer :: idx, expected_kind
          logical :: dec_ok
          character(len=256) :: lookup_msg
          call decode_keyed_request(header, key)
          call session_schema_lookup(session, trim(key), idx, expected_kind, dec_ok, lookup_msg)
          if (.not. dec_ok) then
            call encode_keyed_failure(op, id, "unknown_key", trim(lookup_msg), reply_header)
          else if (reject_kind_mismatch(op, id, trim(key), cfg_kind_string, expected_kind, reply_header)) then
            ! type_mismatch reply already encoded by reject_kind_mismatch
          else
            call solver_session_get_string(session, trim(key), str_value, status, message)
            if (status /= solver_status_ok) then
              call encode_keyed_failure(op, id, "unknown_key", trim(message), reply_header)
            else
              call encode_get_string_reply(id, trim(str_value), reply_header)
            end if
          end if
        end block
      end if

    case ("GET_REAL3")
      if (my_rank() == 0) then
        block
          character(len=128) :: key
          real(real64) :: vec_value(3)
          integer :: idx, expected_kind
          logical :: dec_ok
          character(len=256) :: lookup_msg
          call decode_keyed_request(header, key)
          call session_schema_lookup(session, trim(key), idx, expected_kind, dec_ok, lookup_msg)
          if (.not. dec_ok) then
            call encode_keyed_failure(op, id, "unknown_key", trim(lookup_msg), reply_header)
          else if (reject_kind_mismatch(op, id, trim(key), cfg_kind_real3, expected_kind, reply_header)) then
            ! type_mismatch reply already encoded by reject_kind_mismatch
          else
            call solver_session_get_real3(session, trim(key), vec_value, status, message)
            if (status /= solver_status_ok) then
              call encode_keyed_failure(op, id, "unknown_key", trim(message), reply_header)
            else
              call encode_get_real3_reply(id, vec_value, reply_header)
            end if
          end if
        end block
      end if

    case ("GET_PROGRESS")
      if (my_rank() == 0) then
        call solver_session_get_progress(session, progress)
        call encode_get_progress_reply(id, &
                                       iteration=progress % iteration, &
                                       sim_time=progress % sim_time, &
                                       residual=progress % residual, &
                                       header_bytes=reply_header)
      end if

    case ("GET_POINT_COUNT")
      if (my_rank() == 0) then
        call solver_session_get_global_point_count(session, n_global, status, message)
        if (status /= solver_status_ok) then
          call encode_keyed_failure(op, id, "INVALID_STATE", trim(message), reply_header)
        else
          call encode_get_point_count_reply(id, n_global, reply_header)
        end if
      end if

    case ("DESTROY")
      ! solver_session_destroy is called outside the loop; just reply.
      if (my_rank() == 0) call encode_reply_ok(id, op, reply_header)
      done = .true.

    case default
      if (my_rank() == 0) call encode_reply_error(id, op, "INVALID_ARGUMENT", &
                                                  "unknown opcode "//trim(op), reply_header)
    end select
  end subroutine dispatch_opcode

  !> Service a COPY_SOLUTION request.
  !!
  !! Decode the requested primitive list, copy current solution from
  !! the session into per-primitive buffers, pack them as a binary
  !! payload (little-endian real64, primitives concatenated in the
  !! decoded order), and emit the reply header that names the field
  !! offsets. Only rank 0 owns the socket and emits the reply.
  subroutine handle_copy_solution(session, id, op, header, &
                                  reply_header, reply_payload, status, message)
    type(solver_session_t), intent(inout) :: session
    character(len=*), intent(in) :: id, op
    integer(int8), intent(in) :: header(:)
    integer(int8), allocatable, intent(out) :: reply_header(:), reply_payload(:)
    integer, intent(out) :: status
    character(len=*), intent(out) :: message

    character(len=16), allocatable :: prim_names(:)
    integer :: n_prims, n_global, n_rank0, i, astat
    integer(int64) :: nbytes, stride, lo, hi
    real(real64), allocatable :: x(:), rho(:), u(:), p(:)
    integer(int8), allocatable :: chunk(:)

    call decode_copy_solution(header, n_prims, prim_names)
    if (session % dim == 2) then
      call handle_copy_solution_2d(session, id, op, n_prims, prim_names, &
                                   reply_header, reply_payload, status, message)
      return
    end if
    call solver_session_get_global_point_count(session, n_global, status, message)
    if (op_failed(id, op, status, message, reply_header)) return

    ! Refuse a reply that cannot be represented on the wire BEFORE any
    ! allocation or collective work: the reply header's per-primitive byte
    ! offsets are int32 JSON values and the payload frame length prefix is
    ! u32 (audit 2026-07-06 N1) — crashing the worker over a too-large
    ! reply is never acceptable. n_global and n_prims are identical on
    ! every rank, so all ranks take this branch in lockstep.
    stride = int(n_global, int64) * 8_int64
    if (int(n_prims - 1, int64) * stride > int(huge(0_int32), int64) .or. &
        int(n_prims, int64) * stride > MAX_SEND_FRAME_BYTES) then
      status = solver_status_invalid_argument
      message = 'COPY_SOLUTION reply too large for a single frame; request fewer primitives'
      if (op_failed(id, op, status, message, reply_header)) return
    end if

    ! Allocate the whole-domain buffer only on rank 0; other ranks pass a
    ! zero-length buffer that solver_session_copy_global_solution ignores.
    ! The call itself is collective — every rank must enter the gather, so
    ! an allocation failure must tear the job down collectively
    ! (parallel_fatal), never a rank-divergent error stop that leaves the
    ! peers blocked in the gather (audit 2026-07-06 N3, same class as L3).
    if (my_rank() == 0) then
      n_rank0 = n_global
    else
      n_rank0 = 0
    end if
    allocate (x(n_rank0), rho(n_rank0), u(n_rank0), p(n_rank0), stat=astat)
    if (astat /= 0) call parallel_fatal('cortex_run_attached: copy-solution buffer allocation failed')
    call solver_session_copy_global_solution(session, x, rho, u, p, status, message)
    if (op_failed(id, op, status, message, reply_header)) return

    if (my_rank() == 0) then
      ! Size and index the reply payload in int64. n_prims is capped at
      ! MAX_PRIMITIVES by decode_copy_solution and n_global is the grid
      ! point count, but their product * 8 can exceed int32 on a large
      ! grid, so all of the byte arithmetic is done in int64.
      nbytes = int(n_prims, int64) * stride
      allocate (reply_payload(nbytes), stat=astat)
      if (astat /= 0) call parallel_fatal('cortex_run_attached: copy-solution reply payload allocation failed')
      do i = 1, n_prims
        select case (trim(prim_names(i)))
        case ("x")
          chunk = transfer(x, chunk, stride)
        case ("rho")
          chunk = transfer(rho, chunk, stride)
        case ("u")
          chunk = transfer(u, chunk, stride)
        case ("p")
          chunk = transfer(p, chunk, stride)
        case default
          ! Lenient unknown-name handling is part of the COPY_SOLUTION
          ! wire-protocol contract with Cortex/the GUI; copy x rather
          ! than reject. Log it for observability (non-breaking).
          ! NOTE: encode_copy_solution_reply (below) echoes the *requested*
          ! primitive label back to the caller over the x data — the label
          ! and the data are therefore mismatched for unknown primitives.
          ! This is intentional wire-compat leniency (WONTFIX): rejecting
          ! unknown names would be a breaking change to the Cortex protocol.
          call log_warn('COPY_SOLUTION: unknown primitive "'//trim(prim_names(i))// &
                        '", copying x')
          chunk = transfer(x, chunk, stride)
        end select
        lo = int(i - 1, int64) * stride + 1_int64
        hi = int(i, int64) * stride
        reply_payload(lo:hi) = chunk
      end do
      call encode_copy_solution_reply(id, n_global, prim_names, reply_header)
    end if
    deallocate (x, rho, u, p, stat=astat)
  end subroutine handle_copy_solution

  !> 2D COPY_SOLUTION: session packs one i-fastest float64 block per requested
  !! primitive (x,y,rho,u,v,p) into a flat buffer; rank 0 transfers it to the
  !! binary payload and emits the nx/ny-bearing reply header. Collective.
  subroutine handle_copy_solution_2d(session, id, op, n_prims, prim_names, &
                                     reply_header, reply_payload, status, message)
    type(solver_session_t), intent(inout) :: session
    character(len=*), intent(in) :: id, op
    integer, intent(in) :: n_prims
    character(len=16), intent(in) :: prim_names(:)
    integer(int8), allocatable, intent(out) :: reply_header(:), reply_payload(:)
    integer, intent(out) :: status
    character(len=*), intent(out) :: message

    integer :: nx_g, ny_g, astat
    integer(int64) :: nbuf, nbytes, stride
    real(real64), allocatable :: buf(:)
    integer(int64), allocatable :: offsets(:)

    call solver_session_get_global_grid_shape(session, nx_g, ny_g, status, message)
    if (op_failed(id, op, status, message, reply_header)) return

    ! Same wire-representability pre-check as the 1D handler: int32 reply
    ! header offsets, u32 frame length prefix (audit 2026-07-06 N1). The
    ! grid shape and n_prims are identical on every rank — symmetric,
    ! collective-safe.
    stride = int(nx_g, int64) * int(ny_g, int64) * 8_int64
    if (int(n_prims - 1, int64) * stride > int(huge(0_int32), int64) .or. &
        int(n_prims, int64) * stride > MAX_SEND_FRAME_BYTES) then
      status = solver_status_invalid_argument
      message = 'COPY_SOLUTION reply too large for a single frame; request fewer primitives'
      if (op_failed(id, op, status, message, reply_header)) return
    end if

    if (my_rank() == 0) then
      nbuf = int(n_prims, int64) * int(nx_g, int64) * int(ny_g, int64)
    else
      nbuf = 0_int64
    end if
    ! The gather below is collective — a rank-divergent error stop here
    ! would leave the peers blocked in it (audit 2026-07-06 N3).
    allocate (buf(nbuf), offsets(n_prims), stat=astat)
    if (astat /= 0) call parallel_fatal('cortex_run_attached: 2D copy-solution buffer allocation failed')

    call solver_session_copy_global_solution_2d(session, prim_names(1:n_prims), buf, offsets, &
                                                nx_g, ny_g, status, message)
    if (op_failed(id, op, status, message, reply_header)) return

    if (my_rank() == 0) then
      nbytes = nbuf * 8_int64
      allocate (reply_payload(nbytes), stat=astat)
      if (astat /= 0) call parallel_fatal('cortex_run_attached: 2D copy-solution payload allocation failed')
      reply_payload = transfer(buf, reply_payload, nbytes)
      call encode_copy_solution_reply_2d(id, nx_g, ny_g, prim_names(1:n_prims), reply_header)
    end if
    deallocate (buf, offsets, stat=astat)
  end subroutine handle_copy_solution_2d

  subroutine send_hello(fd, worker_name)
    integer, intent(in) :: fd
    character(len=*), intent(in) :: worker_name
    integer(int8), allocatable :: header(:), empty_payload(:)
    integer :: astat

    call encode_hello_push(build=trim(worker_name)//"-"//solver_build(), n_ranks=n_ranks(), &
                           header_bytes=header)
    allocate (empty_payload(0), stat=astat)
    if (astat /= 0) error stop 'cortex_run_attached: hello empty payload allocation failed'
    call cortex_send_frame(fd, header, empty_payload)
    deallocate (header, empty_payload, stat=astat)
    ! Receive HELLO.reply and discard (rank 0 trusts Cortex's v=1 here;
    ! the version check fires in Cortex if mismatched, before we ever
    ! get a frame).
    block
      integer(int8), allocatable :: rh(:), rp(:)
      call cortex_recv_frame(fd, rh, rp)
      if (allocated(rh)) deallocate (rh, stat=astat)
      if (allocated(rp)) deallocate (rp, stat=astat)
    end block
  end subroutine send_hello

  subroutine bcast_bytes(buf)
    integer(int8), allocatable, intent(inout) :: buf(:)
    integer :: n, ierr

    if (my_rank() == 0) then
      n = size(buf)
    end if
    call MPI_Bcast(n, 1, MPI_INTEGER, 0, mpi_world(), ierr)
    if (ierr /= 0) call parallel_fatal('cortex_run_attached: bcast length failed')
    if (my_rank() /= 0) then
      if (allocated(buf)) then
        deallocate (buf, stat=ierr)
        if (ierr /= 0) error stop 'cortex_run_attached: bcast buffer deallocate failed'
      end if
      allocate (buf(n), stat=ierr)
      if (ierr /= 0) error stop 'cortex_run_attached: bcast buffer allocation failed'
    end if
    if (n > 0) then
      call MPI_Bcast(buf, n, MPI_BYTE, 0, mpi_world(), ierr)
      if (ierr /= 0) call parallel_fatal('cortex_run_attached: bcast bytes failed')
    end if
  end subroutine bcast_bytes

  !> Encode a rank-0-only ok=false reply for a config-key read/write.
  !! Centralised so the same code shape is used by every SET/GET arm.
  subroutine encode_keyed_failure(op, id, code, message, reply_header)
    character(len=*), intent(in) :: op, id, code, message
    integer(int8), allocatable, intent(out) :: reply_header(:)
    if (my_rank() /= 0) return
    call encode_reply_error(id, op, code, message, reply_header)
  end subroutine encode_keyed_failure

  !> Resolve a config key against the schema matching the session's dim.
  !! dim=1 behaviour is byte-identical to the old direct schema_lookup call.
  subroutine session_schema_lookup(session, key, idx, expected_kind, ok, message)
    type(solver_session_t), intent(in) :: session
    character(len=*), intent(in) :: key
    integer, intent(out) :: idx, expected_kind
    logical, intent(out) :: ok
    character(len=*), intent(out) :: message

    if (session % dim == 2) then
      call schema_lookup_2d(key, idx, expected_kind, ok, message)
    else
      call schema_lookup(key, idx, expected_kind, ok, message)
    end if
  end subroutine session_schema_lookup

  !> True when a SET_*/GET_* opcode's own wire kind is compatible with the
  !! schema's `expected_kind` for the target key. Choice-typed keys (e.g.
  !! `flux_scheme`) travel as JSON strings on the wire and are accepted by
  !! SET_STRING/GET_STRING, mirroring `decode_set_value`'s existing
  !! wire-compatible grouping of `cfg_kind_string` with `cfg_kind_choice`.
  pure function opcode_kind_matches(opcode_kind, expected_kind) result(matches)
    integer, intent(in) :: opcode_kind, expected_kind
    logical :: matches

    matches = (opcode_kind == expected_kind) .or. &
              (opcode_kind == cfg_kind_string .and. expected_kind == cfg_kind_choice)
  end function opcode_kind_matches

  !> Human-readable label for a `cfg_kind_*` tag, used only in
  !! `type_mismatch` messages (never as a wire code).
  pure function kind_label(kind) result(label)
    integer, intent(in) :: kind
    character(len=16) :: label

    select case (kind)
    case (cfg_kind_int)
      label = "integer"
    case (cfg_kind_real)
      label = "real"
    case (cfg_kind_logical)
      label = "logical"
    case (cfg_kind_choice)
      label = "choice"
    case (cfg_kind_string)
      label = "string"
    case (cfg_kind_real3)
      label = "real3"
    case default
      label = "unknown"
    end select
  end function kind_label

  !> Rank-0-only ok=false `type_mismatch` reply for a SET_*/GET_* opcode
  !! whose own kind does not match the schema's `expected_kind` for `key`.
  !! Callers must check `session_schema_lookup`'s `ok` first and keep
  !! reporting `unknown_key` when the key itself is not found; this only
  !! covers a key that *does* exist but was addressed by the wrong-kind
  !! opcode (e.g. SET_INTEGER on the real-valued key `cfl`), aligning the
  !! real worker with the cortex stub's stricter behaviour.
  function reject_kind_mismatch(op, id, key, opcode_kind, expected_kind, reply_header) result(mismatch)
    character(len=*), intent(in) :: op, id, key
    integer, intent(in) :: opcode_kind, expected_kind
    integer(int8), allocatable, intent(out) :: reply_header(:)
    logical :: mismatch

    mismatch = .not. opcode_kind_matches(opcode_kind, expected_kind)
    if (mismatch .and. my_rank() == 0) then
      call encode_reply_error(id, op, "type_mismatch", &
                              trim(key)//": expected "//trim(kind_label(expected_kind))//", requested via "// &
                              trim(op)//" ("//trim(kind_label(opcode_kind))//")", reply_header)
    end if
  end function reject_kind_mismatch

  !> Turn a failed session operation into a proper reply-level error and tell
  !! the caller to stop handling the current opcode. Returns .true. on failure.
  !!
  !! Recoverable session errors (an invalid config, a missing/unreadable
  !! grid_file, an over-fine MPI decomposition, …) must NOT tear the worker
  !! down: the previous behaviour emitted an out-of-band error *push* and then
  !! `parallel_fatal`, which (a) crashed the worker mid-request so the client saw
  !! a bare connection reset, and (b) raced the frame so the real message was
  !! usually lost ("unknown error"). Instead we encode a normal reply-error
  !! (with the request id + message, like LOAD_NAMELIST) into `reply_header`;
  !! the frame loop sends it, the client surfaces the actual reason, and the
  !! session stays alive to serve the client's follow-up (retry or DESTROY).
  !!
  !! Collective-safe: session errors are deterministic in the (broadcast)
  !! config, so every rank computes the same `ok` and returns together, keeping
  !! the frame loop in lockstep (only rank 0 encodes the reply).
  logical function op_failed(id, opcode, status, message, reply_header) result(failed)
    character(len=*), intent(in) :: id, opcode, message
    integer, intent(in) :: status
    integer(int8), allocatable, intent(inout) :: reply_header(:)
    character(len=16) :: code
    failed = status /= solver_status_ok
    if (failed .and. my_rank() == 0) then
      ! Map the session status to the protocol's established ErrorCode set so
      ! the client classifies it exactly as the LOAD_*/WRITE_* handlers do,
      ! rather than inventing a new code the GUI treats as unknown.
      select case (status)
      case (solver_status_config_error)
        code = "CONFIG_ERROR"
      case (solver_status_io_error)
        code = "IO_ERROR"
      case (solver_status_invalid_state)
        code = "INVALID_STATE"
      case (solver_status_invalid_argument)
        code = "INVALID_ARGUMENT"
      case default
        code = "INTERNAL_ERROR"
      end select
      call encode_reply_error(id, opcode, trim(code), trim(message), reply_header)
    end if
  end function op_failed

  !> Tick handler bound on rank 0 only. Encodes a PROGRESS_TICK and
  !! writes it to the existing attached-mode socket fd.
  subroutine emit_progress_tick(iter, sim_time, dt, residual, wallclock_s)
    integer, intent(in) :: iter
    real(wp), intent(in) :: sim_time, dt, residual, wallclock_s
    integer(int8), allocatable :: header_bytes(:), empty_payload(:)
    integer :: astat
    if (my_rank() /= 0) return
    call encode_progress_tick(iter, sim_time, dt, residual, wallclock_s, header_bytes)
    allocate (empty_payload(0), stat=astat)
    if (astat /= 0) error stop 'cortex_run_attached: progress tick empty payload allocation failed'
    call cortex_send_frame(g_fd, header_bytes, empty_payload)
    deallocate (header_bytes, empty_payload, stat=astat)
  end subroutine emit_progress_tick

end module cortex_run_attached_mod