#!/usr/bin/env python3

import errno
import gc
import platform
import resource
import sys

import av
import av.filter


try:
 import setproctitle as _setproctitle
 setproctitle = _setproctitle.setproctitle
except ImportError:
 setproctitle = lambda s: None


def leak(file):
 c = av.open(file, "r")
 v = c.decode(video=0)

 g = av.filter.Graph()
 v_src = g.add_buffer(template=c.streams.video[0])
 f_fps = g.add("fps", "fps=%d" % 60)
 v_src.link_to(f_fps)
 v_snk = g.add("buffersink")
 f_fps.link_to(v_snk)
 
 n_frames = 0
 while True:
  try:
   vf = next(v)
   v_src.push(vf)
   del vf
   gc.collect()
   vf = v_snk.pull()
   del vf
   gc.collect()
   n_frames += 1
  except av.AVError as exc:
   if exc.errno == errno.EAGAIN:
    continue
   else:
    raise
  except StopIteration:
   break
 
 return n_frames


def main(argv):
 usage = lambda: print("Usage: %s <video> <n-iterations>" % argv[0], file=sys.stderr)
 
 if len(argv) != 3:
  usage()
  return 2
 
 try:
  v = argv[1]
  n_iter = int(argv[2])
 except ValueError:
  usage()
  return 2
 
 setproctitle("fltr-lk %d      # %s" % (n_iter, " ".join(argv)))
 
 n_frames = 0
 mem_start = resource.getrusage(resource.RUSAGE_SELF)
 for i in range(n_iter):
  n_frames += leak(v)
 mem_end = resource.getrusage(resource.RUSAGE_SELF)
 mem_delta = mem_end.ru_maxrss - mem_start.ru_maxrss
 
 print("*** DONE: %d iteration%s, %d frame%s ***"
       % (n_iter, "s" if n_iter != 1 else "",
          n_frames, "s" if n_frames != 1 else ""))
 print("*** Memory usage from run%s: %d kB (~%s) ***"
       % ("s" if n_iter != 1 else "", rusage_kb(mem_delta), format_rusage(mem_delta)))
 input("Press Enter to exit")
 return 0


def format_bytes(n):
 # from <https://github.com/mikeboers/PyAV/blob/master/examples/resource_use.py>
 order = 0
 while n > 1024:
  order += 1
  n //= 1024
 return '%d%sB' % (n, ('', 'k', 'M', 'G', 'T', 'P')[order])


# Linux's getrusage(2) uses kilobytes; macOS's uses bytes
osname = platform.system()
rusage_kb = lambda n: n if osname == "Linux" else n // 1024
format_rusage = lambda n: format_bytes(n*1024 if osname == "Linux" else n)


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