!> @file checkpoint_io.f90
!> @brief Dimension-blind checkpoint machinery: unified format, pointer files,
!!        header/grid records, and collective outcome agreement (R4, spec §6.2
!!        rev 4).
!!
!! Unified on-disk format (current version: `ckpt_version` below): every dim writes
!!   magic(42), version(`ckpt_version`), dim | iter, t, extents(1:dim), neq, dt
!!   | grid_type(char32), fingerprint(2) | payload arrays (dim-specific).
!! The dim field makes files self-identifying; cross-dim restarts are rejected
!! with a clear error. Legacy formats (1D v2/v3, 2D magic-4242 v1) are NOT
!! readable — rejected by the magic/version checks (owner-authorized break;
!! no user-generated checkpoints exist).
!!
!! MPI: `agree_outcome`/`conclude_read` wrap the collective par_lor decision
!! points. Callers own the collective SEQUENCE — every rank must reach each
!! call site uniformly.
module checkpoint_io
  use precision, only: wp
  use logger, only: log_warn
  use parallel_reductions, only: par_lor
  implicit none
  private

  integer, parameter, public :: ckpt_magic = 42   !< Unified file magic.
  !> Unified multi-dim format version. v5 (F2) adds two 1D-only trailing
  !! records after the (unchanged) v4 body — the Kahan `t_comp` compensation
  !! scalar and the boundary halo/ghost cells (see checkpoint.f90's module
  !! header) — so 1D files grow by those two records. 2D's on-disk content is
  !! byte-for-byte unchanged by this bump (checkpoint_2d.f90 gains no new
  !! fields); only the shared header's version number advances, since
  !! `ckpt_version` is one constant for both dims. v6 (Q2d) leaves 1D's record
  !! list UNCHANGED (fresh 1D files simply stamp v6) and adds 2D-only trailing
  !! records after the global solution array — the Kahan `t_comp` compensation
  !! scalar, the halo width `h` (validated on read), and the four global
  !! boundary-ghost-frame bands `gh_l`/`gh_r`/`gh_b`/`gh_t` — the 2D analogues
  !! of 1D's v5 records (see checkpoint_2d.f90's module header for the layout
  !! and the ghost-history rationale). Either way, legacy v4 files and now
  !! also v5 files (1D or 2D) are NOT readable — `read_ckpt_header`'s strict
  !! version-equality check rejects them (owner-authorized break, third in the
  !! v4→v5→v6 series; no user-generated checkpoints exist).
  integer, parameter, public :: ckpt_version = 6

  public :: ckpt_filename
  public :: update_latest_pointer, resolve_latest_pointer
  public :: write_ckpt_header, read_ckpt_header
  public :: write_grid_record, read_grid_record_and_check
  public :: agree_outcome, conclude_read

contains

  !> '<base>_<iter>.bin' (width-agnostic I0 — never overflows the field).
  function ckpt_filename(base, iter) result(fname)
    character(len=*), intent(in) :: base
    integer, intent(in) :: iter
    character(len=512) :: fname

    write (fname, '(A,A,I0,A)') trim(base), '_', iter, '.bin'
  end function ckpt_filename

  !> Write the most-recent checkpoint path to the pointer file (rank 0 only —
  !! the caller gates on rank).
  subroutine update_latest_pointer(pointer_file, fname, prefix, is_ok, message)
    character(len=*), intent(in) :: pointer_file, fname, prefix
    logical, intent(out), optional :: is_ok
    character(len=*), intent(out), optional :: message
    integer :: u, info

    if (present(is_ok)) is_ok = .true.
    if (present(message)) message = ''

    open (newunit=u, file=trim(pointer_file), status='replace', &
          action='write', form='formatted', iostat=info)
    if (info /= 0) then
      call log_warn(trim(prefix)//': cannot update '//trim(pointer_file)//' pointer file')
      if (present(is_ok)) is_ok = .false.
      if (present(message)) message = trim(prefix)//': cannot update '//trim(pointer_file)//' pointer file'
      return
    end if
    write (u, '(A)', iostat=info) trim(fname)
    if (info /= 0) then
      call log_warn(trim(prefix)//': cannot write '//trim(pointer_file)//' pointer file')
      if (present(is_ok)) is_ok = .false.
      if (present(message)) message = trim(prefix)//': cannot write '//trim(pointer_file)//' pointer file'
      close (u, iostat=info)
      return
    end if
    close (u, iostat=info)
    if (info /= 0) then
      call log_warn(trim(prefix)//': close failed for '//trim(pointer_file)//' pointer file')
      if (present(is_ok)) is_ok = .false.
      if (present(message)) message = trim(prefix)//': close failed for '//trim(pointer_file)//' pointer file'
    end if
  end subroutine update_latest_pointer

  !> If fname names the pointer file, read the real path from it; otherwise
  !! return fname unchanged. Detects >512-char truncation explicitly.
  function resolve_latest_pointer(pointer_file, fname, prefix, is_ok, message) result(actual)
    character(len=*), intent(in) :: pointer_file, fname, prefix
    logical, intent(out), optional :: is_ok
    character(len=*), intent(out), optional :: message
    character(len=512) :: actual
    character(len=4096) :: scratch
    integer :: u, info

    if (present(is_ok)) is_ok = .true.
    if (present(message)) message = ''

    if (trim(fname) == trim(pointer_file)) then
      open (newunit=u, file=trim(pointer_file), status='old', &
            action='read', form='formatted', iostat=info)
      if (info /= 0) then
        call fail_resolve(trim(prefix)//': cannot open '//trim(pointer_file)//' pointer file')
        return
      end if
      read (u, '(A)', iostat=info) scratch
      if (info /= 0) then
        close (u, iostat=info)
        if (info /= 0) &
          call log_warn(trim(prefix)//': close failed for '//trim(pointer_file)//' pointer file')
        call fail_resolve(trim(prefix)//': cannot read '//trim(pointer_file)//' pointer file')
        return
      end if
      close (u, iostat=info)
      if (info /= 0) &
        call log_warn(trim(prefix)//': close failed for '//trim(pointer_file)//' pointer file')
      if (len_trim(scratch) > 512) then
        call fail_resolve(trim(prefix)//': '//trim(pointer_file)// &
                          ' path exceeds 512 characters (truncation detected)')
        return
      end if
      actual = trim(scratch)
    else
      actual = trim(fname)
    end if

  contains

    subroutine fail_resolve(text)
      character(len=*), intent(in) :: text
      if (present(is_ok)) is_ok = .false.
      if (present(message)) message = text
      if (.not. present(is_ok) .and. .not. present(message)) error stop trim(text)
      actual = ''
    end subroutine fail_resolve
  end function resolve_latest_pointer

  !> Write the unified header (magic/`ckpt_version`/dim) + scalars. `extents` carries the dim global
  !! extents (1D: [n_pt_global]; 2D: [nx_global, ny_global]).
  subroutine write_ckpt_header(u, dim, iter, t, extents, neq, dt, info)
    integer, intent(in) :: u, dim, iter, extents(:), neq
    real(wp), intent(in) :: t, dt
    integer, intent(out) :: info

    if (size(extents) /= dim) error stop 'checkpoint_io: extents size does not match dim (write_ckpt_header)'
    write (u, iostat=info) ckpt_magic, ckpt_version, dim
    if (info == 0) write (u, iostat=info) iter, t, extents, neq, dt
  end subroutine write_ckpt_header

  !> Read + validate the unified header against the caller's expectations.
  !! On any failure ok=.false. and err carries '<prefix>: <reason>'.
  subroutine read_ckpt_header(u, prefix, expected_dim, expected_extents, expected_neq, &
                              iter, t, dt, ok, err)
    integer, intent(in) :: u, expected_dim, expected_extents(:), expected_neq
    character(len=*), intent(in) :: prefix
    integer, intent(out) :: iter
    real(wp), intent(out) :: t, dt
    logical, intent(out) :: ok
    character(len=*), intent(out) :: err

    integer :: info, magic, version, dim_ck, neq_ck
    integer :: extents_ck(size(expected_extents))
    character(len=64) :: dim_txt

    ok = .true.
    err = ''
    iter = 0
    t = 0.0_wp
    dt = 0.0_wp

    if (size(expected_extents) /= expected_dim) &
      error stop 'checkpoint_io: expected_extents size does not match expected_dim (read_ckpt_header)'

    read (u, iostat=info) magic, version, dim_ck
    if (info /= 0) then
      ok = .false.
      err = trim(prefix)//': truncated header'
      return
    end if
    if (magic /= ckpt_magic) then
      ok = .false.
      err = trim(prefix)//': bad magic number — not a valid checkpoint file '// &
            '(legacy pre-v4 checkpoints are not restartable)'
      return
    end if
    if (version /= ckpt_version) then
      ok = .false.
      block
        character(len=16) :: vtxt
        write (vtxt, '(I0)') ckpt_version
        err = trim(prefix)//': unsupported version (unified v'//trim(vtxt)// &
              ' required; older formats are not restartable)'
      end block
      return
    end if
    if (dim_ck /= expected_dim) then
      ok = .false.
      write (dim_txt, '(A,I0,A,I0,A)') ': checkpoint is a ', dim_ck, &
        'D file but this solver expected ', expected_dim, 'D'
      err = trim(prefix)//trim(dim_txt)
      return
    end if

    read (u, iostat=info) iter, t, extents_ck, neq_ck, dt
    if (info /= 0) then
      ok = .false.
      err = trim(prefix)//': truncated scalars'
      return
    end if
    if (any(extents_ck /= expected_extents)) then
      ok = .false.
      err = trim(prefix)//': global grid extents mismatch between checkpoint and current grid'
      return
    end if
    if (neq_ck /= expected_neq) then
      ok = .false.
      err = trim(prefix)//': neq mismatch in checkpoint file'
      return
    end if
  end subroutine read_ckpt_header

  !> Write the grid record: grid_type token + 2-component fingerprint.
  subroutine write_grid_record(u, grid_type, fp, info)
    integer, intent(in) :: u
    character(len=32), intent(in) :: grid_type
    real(wp), intent(in) :: fp(2)
    integer, intent(out) :: info

    write (u, iostat=info) grid_type, fp
  end subroutine write_grid_record

  !> Read the grid record and compare against the configured grid using the
  !! historical 1e-12 relative tolerance.
  subroutine read_grid_record_and_check(u, prefix, grid_type, cur_fp, ok, err)
    integer, intent(in) :: u
    character(len=*), intent(in) :: prefix
    character(len=32), intent(in) :: grid_type
    real(wp), intent(in) :: cur_fp(2)
    logical, intent(out) :: ok
    character(len=*), intent(out) :: err

    integer :: info
    character(len=32) :: ck_grid_type
    real(wp) :: ck_fp(2)

    ok = .true.
    err = ''
    read (u, iostat=info) ck_grid_type, ck_fp(1), ck_fp(2)
    if (info /= 0) then
      ok = .false.
      err = trim(prefix)//': truncated grid record'
      return
    end if
    if (trim(ck_grid_type) /= trim(grid_type) .or. &
        any(abs(ck_fp - cur_fp) > 1.0e-12_wp * max(1.0_wp, abs(cur_fp)))) then
      ok = .false.
      err = trim(prefix)//': grid mismatch — the checkpoint grid differs from '// &
            'the configured grid; restart with the original grid_type/grid_file'
    end if
  end subroutine read_grid_record_and_check

  !> Collective outcome agreement (write-side): failed = any rank not ok.
  !! On failure the optional status pair gets fail_text; on SUCCESS the pair
  !! is left exactly as the caller preset it — hence intent(inout), NOT
  !! intent(out) (which would leave the caller's values undefined on the
  !! success path). COLLECTIVE — every rank must call.
  function agree_outcome(local_ok, fail_text, is_ok, message) result(failed)
    logical, intent(in) :: local_ok
    character(len=*), intent(in) :: fail_text
    logical, intent(inout), optional :: is_ok
    character(len=*), intent(inout), optional :: message
    logical :: failed

    failed = par_lor(.not. local_ok)
    if (failed) then
      if (present(is_ok)) is_ok = .false.
      if (present(message)) message = trim(fail_text)
    end if
  end function agree_outcome

  !> Shared read epilogue: ONE collective decision point (the only par_lor of
  !! the read path's conclusion); returns the agreed failure flag so callers
  !! write `if (conclude_read(...)) return`. The failing rank reports its own
  !! err, other ranks the generic text; error-stop fallback when the caller
  !! passed no status arguments. COLLECTIVE — every rank must call.
  !! Status pair semantics as in agree_outcome: assigned only on failure,
  !! caller presets preserved on success (intent(inout)).
  function conclude_read(local_ok, err, generic_text, is_ok, message) result(failed)
    logical, intent(in) :: local_ok
    character(len=*), intent(in) :: err, generic_text
    logical, intent(inout), optional :: is_ok
    character(len=*), intent(inout), optional :: message
    logical :: failed

    failed = par_lor(.not. local_ok)
    if (.not. failed) return
    if (present(is_ok)) is_ok = .false.
    if (present(message)) then
      if (.not. local_ok) then
        message = trim(err)
      else
        message = trim(generic_text)
      end if
    end if
    if (.not. present(is_ok) .and. .not. present(message)) then
      if (.not. local_ok) then
        error stop trim(err)
      else
        error stop trim(generic_text)
      end if
    end if
  end function conclude_read

end module checkpoint_io
