#!/usr/bin/env python3
# ============================================================================
#  fix_numrecs.py  -  Khoi phuc truong 'numrecs' cho file NetCDF CLASSIC (CDF-1/2)
#                     bi ghi dang do (vd RegCM MPI_Abort truoc khi dong file) khien
#                     header bao 0 record du DU LIEU record DA nam tren dia.
#
#  Cach lam: doc layout header (recsize, offset vung record), tinh lai so record
#  thuc tu KICH THUOC FILE:  numrecs = (filesize - begin_rec) // recsize
#  Record cuoi bi CUT (abort giua chung) se bi bo (dung phep chia lay phan nguyen).
#
#  AN TOAN:
#   - Chi dung cho NetCDF classic (magic 'CDF').
#   - MAC DINH chi patch khi header dang numrecs==0 (file hong); file da co
#     numrecs>0 -> BO QUA (idempotent, chay lai nhieu lan vo hai). Dung --force
#     de ghi de du the nao.
#   - Khong --patch: chi IN ket qua tinh duoc (dry-run).
#
#  Dung:
#     python3 fix_numrecs.py FILE.nc            # dry-run: in numrecs tinh duoc
#     python3 fix_numrecs.py FILE.nc --patch    # ghi numrecs (neu header dang 0)
#     python3 fix_numrecs.py FILE.nc --patch --force
# ============================================================================
import sys, os, struct

TSZ = {1: 1, 2: 1, 3: 2, 4: 4, 5: 4, 6: 8}  # byte, char, short, int, float, double


def _ri(b, o):
    return struct.unpack('>i', b[o:o + 4])[0], o + 4


def _rstr(b, o):
    n, o = _ri(b, o)
    name = b[o:o + n]
    o += (n + 3) // 4 * 4
    return name, o


def _ratt(b, o):
    _, o = _ri(b, o)          # tag (NC_ATTRIBUTE hoac ABSENT)
    n, o = _ri(b, o)
    for _ in range(n):
        _, o = _rstr(b, o)
        t, o = _ri(b, o)
        ne, o = _ri(b, o)
        sz = ne * TSZ[t]
        o += (sz + 3) // 4 * 4
    return o


def parse(fn):
    with open(fn, 'rb') as f:
        b = f.read(1 << 20)   # header luon nho hon 1MB
    if b[:3] != b'CDF':
        raise ValueError('khong phai NetCDF classic (magic != CDF): %r' % b[:4])
    ver = b[3]
    o = 4
    numrecs, o = _ri(b, o)
    # --- dim_list ---
    _, o = _ri(b, o)          # tag
    nd, o = _ri(b, o)
    recdim = None
    for i in range(nd):
        _, o = _rstr(b, o)
        dl, o = _ri(b, o)
        if dl == 0:
            recdim = i
    o = _ratt(b, o)           # global attributes
    # --- var_list ---
    _, o = _ri(b, o)          # tag
    nv, o = _ri(b, o)
    recvars = []
    for _ in range(nv):
        nm, o = _rstr(b, o)
        ndims, o = _ri(b, o)
        dids = []
        for _ in range(ndims):
            d, o = _ri(b, o)
            dids.append(d)
        o = _ratt(b, o)       # var attributes
        _, o = _ri(b, o)      # nc_type
        vsize, o = _ri(b, o)
        if ver == 1:
            begin, o = _ri(b, o)
        else:                 # CDF-2: begin la 8 byte
            begin = struct.unpack('>q', b[o:o + 8])[0]
            o += 8
        if recdim is not None and recdim in dids:
            recvars.append((nm.decode('ascii', 'replace'), vsize, begin))
    if not recvars:
        raise ValueError('khong co bien record (khong co chieu UNLIMITED)')
    recsize = sum(v for _, v, _ in recvars)
    begin_rec = min(bg for _, _, bg in recvars)
    return ver, numrecs, recsize, begin_rec, recvars


def main(argv):
    if len(argv) < 2:
        print(__doc__)
        return 2
    fn = argv[1]
    patch = '--patch' in argv[2:]
    force = '--force' in argv[2:]
    sz = os.path.getsize(fn)
    ver, numrecs, recsize, begin_rec, recvars = parse(fn)
    nrec = (sz - begin_rec) // recsize
    rem = (sz - begin_rec) % recsize
    print('%s  CDF%d  filesize=%d  numrecs(header)=%d  recsize=%d  begin_rec=%d'
          % (os.path.basename(fn), ver, sz, numrecs, recsize, begin_rec))
    print('   => numrecs tinh tu kich thuoc = %d  (record cuoi cut: du %d byte)'
          % (nrec, rem))
    if not patch:
        print('   (dry-run: them --patch de ghi)')
        return 0
    if numrecs != 0 and not force:
        print('   BO QUA: header da co numrecs=%d (khong phai 0). Dung --force neu muon ghi de.'
              % numrecs)
        return 0
    if not (0 < nrec < 2 ** 31):
        print('   TU CHOI: numrecs=%d ngoai khoang hop le.' % nrec)
        return 1
    with open(fn, 'r+b') as f:
        f.seek(4)
        f.write(struct.pack('>i', nrec))
    print('   DA PATCH numrecs=%d.' % nrec)
    return 0


if __name__ == '__main__':
    sys.exit(main(sys.argv))