#!/usr/bin/env python3
import os,re,sys
from datetime import datetime

from debug import DebugChannel
#dc=DebugChannel(False)

def intdivmod(a,b):
  "Like regular divmod(), but returns integers."

  a,b=divmod(a,b)
  return (int(a),int(b))

# Get the character size of this terminal when this module is loaded.
try:
  _,tty_columns=[int(x) for x in os.popen('stty size','r').read().split()]
except:
  try:
    tty_columns=int(os.environ.get('COLUMNS','80'))
  except:
    tty_columns=80

class ProgressError(Exception):
  "This is the voice of this module when it is displeased."

  def __init__(msg):
    super(ProgressError,self).__init__(msg)

class Progress(object):
  """An object of this class can produce a progress bar as a string when
  str() is applied to it. It must be instantiated with an upper limit,
  and after that, the current position from 0 to this limit may be set
  by calling the object as a function.

  Example:

  finish=17500
  progress=progress.Progress(finish)
  for i in range(finish+1):
    if i%100==0:
      sys.stdout.write('\\r'+str(progress(i)'))
  print '\\r'+str(progress(i))
  """

  #@dc
  def __init__(self,limit,fmt='({progress})',trailer='=',mark='>',leader='-',end='|',width=tty_columns,started=None,strict=True):
    """Initialize this object with its limit and any other optional
    arguments the caller might provide.

    Arguments:
        limit     The numeric upper limit of our progress, e.g. the
                  number of bytes in a file to be uploaded.
        fmt       Format string for the numeric value of this progress
                  bar. The values "progress" and "limit" are available
                  to be formatted.
        trailer   Character to fill in the bar to the left of (i.e.
                  trailing behind) the current position.
        mark      Character that marks the porgress bar's current
                  progress position.
        leader    Character to fill in the bar to the right of (leading)
                  the current position.
        end       Character that marks the right side of the bar.
        width     The number of characters to fit the progress bar into.
                  (default: terminal's character width)
        started   An optional datetime instance marking when whatever
                  we're monitoring started. If given, this also appends
                  an estimated time until completion to the end of our
                  progress output.
        strict    If True, the current progres value is always forced
                  into the interval [0-limit].
    """

    self.limit=limit
    self.fmt=fmt
    self.trailer=trailer
    self.mark=mark
    self.leader=leader
    self.end=end
    self.width=width
    self.started=started
    self.strict=strict
    self.progress=0 # Our start position is 0.

  def scale(self,progress):
    """Compute the scaled value of the given position, and return a
    (scaled_position,scaled_limit) tuple. This base implementation
    returns these values without modification, but see PercentProgress
    for how this might be used."""

    # This is the default implementation.
    return progress,self.limit

  def __call__(self,progress):
    """Accept a new position for this Progress object and return
    this object."""

    self.progress=self.enforceRange(progress)
    return self

  def enforceRange(self,progress):
    "If we're being strict, keep position within [0-limit]."

    if self.strict:
      if progress<0:
        progress=0
      elif progress>self.limit:
        progress=self.limit
    return progress

  def progressString(self,progress=None):
    """Return the the string-formatted value of the given progress (or
    by default, current the progress of this instance) of this
    Progress object. Our scale() method is called to get the progress
    and limit values to be formatted in the string."""

    if progress is None:
      progress=self.progress

    # Format our position string.
    if self.fmt:
      sp,sl=self.scale(progress) # Get our scaled position and limit.
      try:
        ps=self.fmt.format(**{'progress':sp,'limit':sl})
      except TypeError as e:
        if str(e).startswith('not all arguments'):
          ps=self.fmt
        else:
          raise
    else:
      ps=''

    return ps

  def timeRemaining(self,progress=None):
    """Return a (seconds, eta_string) tuple containing the the estimated
    number of seconds remaining until progress is complete and that same
    value formatted as a DHMS string (e.g. "3m20s"). The DHMS algorithm
    used here limits the string to the two most significant components.
    So it will return "4d5h" rather than "4d5h23m57s"."""

    if progress is None:
      progress=self.progress

    # Format our ETA.
    if self.started:
      # Get our elapsed time in seconds.
      #dc(f"{self.started=} {progress=} {self.limit=}")
      et=datetime.now()-self.started
      et=int(et.days*86400+et.seconds)
      #dc(f"{et} elapsed seconds")

      # Compute our time remaining from the rate of our progress so far.
      r=et/progress if progress else 1
      #dc(f"{r} completion rate")
      s=seconds_remaining=round(r*(self.limit-progress))
      #dc(f"{s} estimated seconds remaining")

      # Compose our DHMS string.
      days,r=intdivmod(s,86400) # Days
      eta=f"{days}d" if days else ''
      hours,r=intdivmod(r,3600) # Hours
      eta+=f"{hours}h" if hours else ''
      if not days:
        # Minutes
        minutes,r=intdivmod(r,60)
        eta+=f"{minutes}m" if minutes else ''
        if not hours:
          # Seconds are left in r.
          eta+=f"{r}s" if r else ''
    else:
      seconds_remaining=None
      eta=''

    return (seconds_remaining,eta)

  def __str__(self):
    "Output a string describing the state of this Progress object."

    # Get the string version of our position and remember its length.
    pos_str=self.progressString()
    pos_len=len(pos_str)
    rseconds,eta=self.timeRemaining()
    eta_len=len(eta)

    # Compute the ideal position of our marker in the width set up for this
    # object, and use this to start building our return value.
    progress=int((self.progress/self.limit)*(self.width-eta_len))
    # Crate our trailer position string (if any).
    if pos_str:
      if progress>pos_len:
        s=self.trailer*(progress-pos_len)
      else:
        s=''
      s+=pos_str
    else:
      s=self.trailer*progress
    # Append our position marker if possible.
    if self.mark:
      if not s:
        # Just start with the position marker with no trailer.
        s=self.mark
      else:
        if s[0]==self.trailer:
          # Remove the first leader character and append the position marker.
          s=s[1:]+self.mark
    # Append the leader.
    s+=self.leader*(self.width-eta_len-len(s))
    # Append our end mark if possible.
    if self.end and s[-1] in (self.leader,self.mark):
      s=s[:-1]+self.end
    s+=eta

    # Return our position string.
    return s

class PercentProgress(Progress):
  """This is just like the Progress class, but the default format is
  '(%0.0f%)' and the position in the progress string is expressed as a
  percentage (from 0 to 100) of the limit value."""

  def __init__(self,limit,fmt='({progress:0.0f}%)',**kwargs):
    fmt=kwargs.get('fmt',None)
    if fmt is None:
      kwargs['fmt']=kwargs.get('fmt','({progress:0.0f}%)')
    super().__init__(limit,**kwargs)

  def scale(self,progress):
    "Scale our position as a percentage of limit (which becomes 100)."

    return (progress/self.limit*100,100)

class ByteSizeProgress(Progress):
  """This is just like the Progress class, but the default format is
  '(%0.2f %s)' and the position in the progress string is expressed in
  terms of B, KB, MB, GB, TB, EB, ZB, or YB."""

  units=('B','KB','MB','GB','TB','EB','ZB','YB')

  def __init__(self,limit,**kwargs):
    """In addition to requiring a limit argument and accepting all the
    optional arguments of Progress.__init__(), an additional optional
    argument is also accepted:

    divisor   This is the value you divide a number by to convert from
              one unit (like KB) to the next higher unit (like GB). It
              defaults to 1000, but you can set it to 1024 if you're
              counting memory bytes rather than filesystem bytes. These
              are the only two permitted values for the divisor
              artument.
    """

    self.divisor=kwargs.get('divisor',1000)
    if self.divisor not in (1000,1024):
      raise ProgressError("Divisor MUST be either 1000 or 1024.")
#   kwargs['fmt']=kwargs.get('fmt','(%0.2f %s)')
    kwargs['fmt']=kwargs.get('fmt','({progress:0.2f} {units})')
    super().__init__(limit,**kwargs)
#   super(ByteSizeProgress,self).__init__(limit,**kwargs)

  def progressString(self,progress=None):
    if progress is None:
      progress=self.progress

    #rseconds,eta=self.timeRemaining(progress)

    if self.fmt:
      # Figure out which units to use.
      for units in self.units:
        if progress<self.divisor:
          break
        progress/=self.divisor
      else:
        progress*=self.divisor # Because we're at the end of our units.

      try:
        s=self.fmt.format(**{'progress':progress,'units':units,'limit':self.limit})
      except TypeError as e:
        if str(e).startswith('not all arguments'):
          s=self.fmt
        else:
          raise
    else:
      s=''
    return s

if __name__=='__main__':
  import argparse,os,sys,textwrap,time
  from datetime import datetime
  #from ast import literal_eval

  prog=os.path.basename(sys.argv[0])

  def numeric(val):
    "Convert a string value to an int or float (or terminate in error)."

    if isinstance(val,(int,float)):
      # Our work is already done.
      return val

    if isinstance(val,str):
      try:
        v=eval(str)
        if isinstance(v,int):
          return v
        if isinstance(v,float):
          return int(v) if int(v)==v else v
      except:
        pass
      print(f"{prog}: Error: Cannot convert {val!r} to a numeric value.",file=sys.stderr)

  try: # This outer "try" is just for KeyboardInterrupt exceptions.
    try:
      term_width=os.environ.get('COLUMNS','80')
      term_width=int(term_width)
    except ValueError:
      print(f"{sys.argv[0]}: Error: {term_width!r} is not a valid integer.",file=sys.stderr)
      sys.exit(1)
    ap=argparse.ArgumentParser(
      add_help=False,
      formatter_class=argparse.RawDescriptionHelpFormatter,
      description='\n'.join(textwrap.wrap("Output a progress bar that fills the character-width of the current terminal and terminate with 0 (true) if PROGRESS < LIMIT. Otherwise, return 1 (false).",width=term_width)),
      epilog='\n'.join(textwrap.wrap("""\
START TIME FORMATS:

  The --started (-s) option takes a start time value in one of two formats. This time is used to estimate when whatever progress is being measured will be complete, and this ETA is suffixed to each line of output.

  [[YY]YY-MM-DD ]HH:MM[:SS]

  I don't think this format needs much explanation. It can look like any of these:

    2023-06-13 22:33:44
    2023-06-13 22:33
    23-06-13 22:33:44
    23-06-13 22:33
    22:33:44
    22:33

  abbr_day abbr_month day HH:MM:SS tz YYYY

  This is a very common output format of the `date` command, e.g. "Tue Jun 6 22:33:44 EDT 2023". This simplifies using --started (-s) to append ETA values to the progress output.

    T0=$(date);
    cat some-file |\\
    while read x; do
      line-processor "$x"
      progress -fnps "$T0" "^$x$" x
    done
""",width=term_width,replace_whitespace=False,break_on_hyphens=False))
    )
    ap.add_argument('--file','-f',action='store_true',help="""Use the PROGRESS and LIMIT arguments as a regular expression and filename, respectively. The file position of the line following the first line matching the regular expression will be used for PROGRESS, and the size of the named file will be used for LIMIT.""")
    ap.add_argument('--file-last',action='store_true',help="""Just like --file, but the last occurance of the pattern in the file is used for the position.""")
    ap.add_argument('--percent','-p',action='store_true',help="Express PROGRESS as a percentage of LIMIT rather than as an absolute value.")
    ap.add_argument('--bytes',action='store_true',help="Express PROGRESS as a number of bytes in human-friendly units (B, KB, MB, GB, TB, or EB) according to the number of bytes in question.")
    ap.add_argument('--divisor',metavar='D',action='store',type=int,choices=(1000,1024),default=1000,help="For use with --bytes. D is the value you divide a number by to convert from one unit (like KB) to the next higher unit (like GB). It defaults to 1000, but you can set it to 1024 if you're counting memory bytes rather than filesystem bytes. These are the only two permitted values for D.")
    ap.add_argument('--started','-s',metavar='T',action='store',help="""T is the time matching the "[[YY]YY-MM-DD ]HH:MM[:SS]" pattern when the process being monitored was started. The presence of this option suffixes the output value with an estimated time to completion.""")
    ap.add_argument('--fmt',action='store',help="""Format string for the numeric value of this progress bar. This option is ignored if --bytes is given. The default format is usually "({progress})", but it's "({progres}%%)" if --percent is present. The "limit" variable is also available for formatting, so for example, you could use "({progress} / {limit}" to indicate what value progress is working toward.""")
    ap.add_argument('--mark',metavar='CHAR',action='store',default='>',help="Character that marks the porgress bar's current position. (default: %(default)s)")
    ap.add_argument('--trailer',metavar='CHAR',action='store',default='=',help="Character to fill in the bar to the left of (trailing behind) the current position. (default: %(default)s)")
    ap.add_argument('--leader',metavar='CHAR',action='store',default='-',help="Character to fill in the bar to the right of (leading) the current position. (default: %(default)s)")
    ap.add_argument('--end',metavar='CHAR',action='store',default='|',help="Character that marks the right side of the bar. (default: %(default)s)")
    ap.add_argument('--newline','-n',action='store_true',help="Write a newline character to standard output after every progress bar. By default, only a carriage return is output, returning the cursor to the left-most column of the terminal.")
    ap.add_argument('--width','-w',metavar='W',action='store',type=int,default=tty_columns-1,help="The number of characters to fit the progress bar into. (default: %(default)s)")
    ap.add_argument('--test',metavar='S',action='store',type=float,default=0,help="Write all progress bar values to standard output and exit. The value of this option is the number of seconds (which may be a floating point value) to delay between each iteration.")
    ap.add_argument('--help','-h',action='help',help="Show this help message and exit.")
    ap.add_argument('progress',metavar='PROGRESS',action='store',help="The position you want this progress bar to express.")
    ap.add_argument('limit',metavar='LIMIT',action='store',help="The upper limit of the progress argument, e.g. the size in bytes of data to be processed in some way.")
    opt=ap.parse_args()

    time_patterns=[
      # [[[YY]YY-]MM-DD ]HH:MM[:SS]
      re.compile(r'^\s*(((?P<year>\d{2,4})[-/])?(?P<month>\d{1,2})[-/](?P<day>\d{1,2})(\s+|[Tt]))?(?P<hour>\d{1,2}):(?P<min>\d{1,2})(:(?P<sec>\d{1,2}))?\s*$'),
      # Sat Dec 24 14:44:43 EST 2022
      re.compile(r'^\w+\s+(?P<month>\w+)\s+(?P<day>\d{1,2})\s+(?P<hour>\d{1,2}):(?P<min>\d{1,2})(:(?P<sec>\d{1,2}))?\s+\w+\s*(?P<year>\d{2,4})?\s*$'),
    ]

    # If --started is given, parse its argument to a datetime value.
    if opt.started:
      if opt.started=='file':
        # Use the modification time of "the file" for the start time.
        if not opt.file:
          print(f"{prog}: Error: \"--started file\" can be used only with \"--file\".",file=sys.stderr)
          sys.exit(2)
        try:
          opt.started=datetime.fromtimestamp(os.stat('x').st_mtime)
        except FileNotFoundError:
          print(f"{prog}: File not found: \"{opt.limit}\"",file=sys.stderr)
          sys.exit(2)
      else:
        # Parse the date and time in opt.started.
        now=datetime.now()
        for p in time_patterns:
          m=p.match(opt.started)
          if m:
            break
        else:
          print(f"{prog}: Error: Can't parse a date and time from {opt.started!r}",file=sys.stderr)
          sys.exit(2)
        # Get our time components, make them integers (except for month), and put them into an object.
        t=m.groupdict()
        t={k:(None if t[k] is None else int(t[k]) if k!='month' else t[k]) for k in t.keys()}
        t=type('',(),t)

        # Fill in amy missing pieces of our start date and time.
        if t.year is None: t.year=now.year
        elif t.year<100: t.year+=(now.year//100)*100
        elif t.year<1000: t.year+=(now.year//1000)*1000
        if t.month is None: t.month=now.month
        if t.day is None: t.day=now.day
        if t.sec is None: t.sec=0

        # Tolerate different types of month identifiers.
        ts=f"{t.year}-{t.month}-{t.day} {t.hour}:{t.min}:{t.sec}"
        #dc(f"{ts=}")
        for m in ('%m','%B','%b'):
          parser=f'%Y-{m}-%d %H:%M:%S'
          #dc(f"{parser=}")
          try:
            pt=datetime.strptime(ts,parser)
            break
          except ValueError:
            pass # Try again using the next month type.
        else:
          print(f"{prog}: Error: Can't parse month from {opt.started!r}",file=sys.stderr)
          sys.exit(2)
        opt.started=pt # Use our parsed time as --started's value.

    # If --file is given, find the position of the first occurance of the
    # opt.progress RE in the file named by opt.limit.
    if opt.file or opt.file_last:
      if opt.file and opt.file_last:
        print(f"{prog}: Error: You can't use both --file and --file-last. Pick one.",file=sys.stderr)
        sys.exit(2)
      try:
        pat=re.compile(opt.progress)
        filename,opt.limit=opt.limit,os.path.getsize(opt.limit)
        found=0
        with open(filename) as f:
          line=f.readline()
          while line:
            m=pat.search(line)
            if m:
              found=f.tell()
              if opt.file:
                break
            line=f.readline()
        opt.progress=found
      except OSError as e:
        print(f"{prog}: {e}",file=sys.stderr)
        sys.exit(2)
      except Exception as e:
        print(f"{prog}: Error: {e}",file=sys.stderr)
        sys.exit(2)

    try:
      if isinstance(opt.progress,str):
        opt.progress=float(opt.progress)
      if isinstance(opt.limit,str):
        opt.limit=float(opt.limit)
    except ValueError as e:
      print(f"{prog}: Error: {e}",file=sys.stderr)
      sys.exit(2)

    #if dc:
    #  dc(f"{opt.file=}")
    #  dc(f"{opt.file_last=}")
    #  dc(f"{opt.newline=}")
    #  dc(f"{opt.fmt=}")
    #  dc(f"{opt.mark=}")
    #  dc(f"{opt.trailer=}")
    #  dc(f"{opt.leader=}")
    #  dc(f"{opt.end=}")
    #  dc(f"{opt.width=}")
    #  dc(f"{opt.test=}")
    #  dc(f"{opt.percent=}")
    #  dc(f"{opt.divisor=}")
    #  dc(f"{opt.started=}")

    # Decide which Progress class the user has called for.
    kwargs={}
    if opt.fmt is not None:
      kwargs['fmt']=opt.fmt
    kwargs.update({k:v for k,v in opt.__dict__.items() if k in ('limit mark trailer leader end width started'.split())})
    #dc(f"{kwargs=}")
    if opt.percent:
      p=PercentProgress(**kwargs)
    elif opt.bytes:
      p=ByteSizeProgress(**kwargs)
    else:
      p=Progress(**kwargs)

    # Run all our test in one call of this command and terminate.
    if opt.test>0:
      for i in range(int(opt.limit)+1):
        time.sleep(opt.test)
        sys.stdout.write(str(p(i)))
        if opt.newline:
          sys.stdout.write('\n')
        else:
          sys.stdout.write('\r')
        sys.stdout.flush()
      if not opt.newline:
        sys.stdout.write('\n')
      sys.exit(0)

    # Write our progress string to standard output.
    sys.stdout.write(str(p(opt.progress)))

    # Write either a \r or a \n after our progress string, depending on
    # whether -n (--newline) was used.
    if opt.newline:
      sys.stdout.write('\n')
    else:
      sys.stdout.write('\r')
  except KeyboardInterrupt:
    sys.exit(1)

  # For the benefit of the shell script calling this script, exit with a
  # value of 0 (true) if the progress we've made is incomplete.
  # Otherwise, return 1 (false).
  sys.exit(0 if opt.progress<opt.limit else 1)
