from io import StringIO


def get_q2_comment_lines(md_file_loc):
    """Returns a list of line numbers in the file that start with "#q2:".

    These lines should be skipped when parsing the file outside of Q2 (i.e.
    in pandas). I guess we could also ostensibly use these lines' types here
    eventually, but for now we just skip them.

    Notes:
     -The line numbers are 0-indexed (so they can easily be thrown in to
      pandas.read_csv() as the skiprows parameter)
     -This doesn't check check the first line of the file (assumed to be the
      header)
     -This stops checking lines once it gets to the first non-header line
      that doesn't start with "#q2:". Currently, "#q2:types" is the only Q2
      "comment directive" available, but ostensibly this could detect future
      Q2 comment directives.
     -This checks if md_file_loc is of type StringIO. If so, this will
      handle it properly (iterating over it directly); otherwise, this
      assumes that md_file_loc is an actual filename, and this will open
      it using open().
      (I realize that ideally this wouldn't have to do any type checking,
      but it's either this or do a bunch of weird refactoring to get my test
      code working.)
    """

    def iterate_over_file_obj_lines(file_obj):
        q2_lines = []
        line_num = 0
        for line in file_obj:
            # Don't check for a #q2: comment on the first line of the file,
            # since the first line (should) define the file header.
            if line_num > 0:
                if line.startswith("#q2:"):
                    q2_lines.append(line_num)
                else:
                    # We assume that all #q2: lines will occur at the start of
                    # the file. Once we've reached a line that doesn't start
                    # with "#q2:", we stop checking.
                    break
            line_num += 1
        return q2_lines

    if type(md_file_loc) == StringIO:
        q2_lines = iterate_over_file_obj_lines(md_file_loc)
        # HACK: Allow us to read through this StringIO again --
        # https://stackoverflow.com/a/27261215/10730311
        # Note that we're only ever bothering with StringIOs here during test
        # code, so this weirdness should be ignored during normal operation of
        # Qurro.
        md_file_loc.seek(0)
        return q2_lines
    else:
        with open(md_file_loc, "r") as md_file_obj:
            return iterate_over_file_obj_lines(md_file_obj)