#!/usr/bin/python
# Copyright 2012 Google Inc. All Rights Reserved.
#
"""A tool for a burnin Flash test."""

__author__ = 'dgentry@google.com (Denton Gentry)'

import ctypes
import fcntl
import os
import os.path
import subprocess
import sys
import time
import options


optspec = """
burnin-flash [options...] /path/to/write/to
--
m,mtd=  /dev/mtd# to check for errors.
p,print-interval=   Seconds between printouts [10]
r,runtime=   Number of seconds to run the test [60]
"""


class MtdEccStats(ctypes.Structure):
  """<mtd/mtd-abi.h> struct mtd_ecc_stats."""

  _fields_ = [('corrected', ctypes.c_uint32),
              ('failed', ctypes.c_uint32),
              ('badblocks', ctypes.c_uint32),
              ('bbtblocks', ctypes.c_uint32)]


def GetMtdStats(mtddev):
  """Return the MtdEccStats for the given mtd device.

  Arguments:
    mtddev: the string path to the device, ex: '/dev/mtd14'
  Raises:
    IOError: if the ioctl fails.
  Returns:
    an MtdEccStats.
  """

  # ECCGETSTATS _IOR('M', 18, struct mtd_ecc_stats)
  ECCGETSTATS = 0x40104d12  # pylint: disable-msg=C6409
  with open(mtddev, 'r') as f:
    ecc = MtdEccStats()
    if fcntl.ioctl(f, ECCGETSTATS, ctypes.addressof(ecc)) != 0:
      raise IOError('ECCGETSTATS failed')
    return ecc


def main():
  o = options.Options(optspec)
  (opt, flags, extra) = o.parse(sys.argv[1:])  #pylint: disable-msg=W0612
  if not opt.mtd:
    o.fatal('an mtd argument is required')
  if len(extra) != 1:
    o.fatal('exactly one path expected')

  filename = extra[0]
  if os.path.isdir(filename):
    filename = os.path.join(filename, 'burnin.tmp')

  f = open(filename, 'w+')

  # we still have the file open; this will ensure it gets deleted upon
  # process exit.
  os.unlink(filename)

  start_time = now = last_print_time = time.time()
  nbytes = 0
  start_ecc = GetMtdStats(opt.mtd)

  while now - start_time < opt.runtime:
    try:
      # strongly random data so as to not compress well.
      dsiz = 1048576
      randdata = subprocess.check_output(['randomdata', '0', str(dsiz)])

      # write random data to flash
      f.seek(0, os.SEEK_SET)
      f.write(randdata)
      f.flush()
      os.fsync(f.fileno())

      # read back, to notice ECC errors.
      open('/proc/sys/vm/drop_caches', 'w+').write('3')
      f.seek(0, os.SEEK_SET)
      readback = f.read(dsiz)
      if readback != randdata:
        sys.stderr.write('Read back from %s != written\n' % filename)
        return 2
      nbytes += dsiz
    except (OSError, IOError) as e:
      sys.stderr.write('%s write: %s\n' % (filename, e))
      return 2

    now = time.time()
    if opt.print_interval and now - last_print_time > opt.print_interval:
      print ('%-15s %7.2fM in %5.2fs = %6.2fM/s'
             % (filename + ':',
                nbytes / 1024. / 1024.,
                now - last_print_time,
                nbytes / 1024. / 1024. / (now - last_print_time)))
      sys.stdout.flush()
      last_print_time = now
      nbytes = 0

    ecc = GetMtdStats(opt.mtd)
    # Ignore ecc.corrected, they are uncommon but normal.
    if ecc.failed != start_ecc.failed:
      print('%s: Uncorrectable ECC error during test, %d != %d' %
            (opt.mtd, ecc.failed, start_ecc.failed))
      return 2
    if ecc.badblocks != start_ecc.badblocks:
      print('%s: bad blocks developed during test, %d != %d' %
            (opt.mtd, ecc.badblocks, start_ecc.badblocks))
      return 2
    if ecc.bbtblocks != start_ecc.bbtblocks:
      print('%s: bad block table error during test, %d != %d' %
            (opt.mtd, ecc.bbtblocks, start_ecc.bbtblocks))
      return 2

  return 0


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