"""Flat file header extraction utilities.
"""
# Importing the version number (in __init__.py) and the package name
from pyaddons import (
    utils
)
from io import UnsupportedOperation
import argparse
import sys
import os
import gzip
import csv
import stdopen
# import pprint as pp

_SCRIPT_NAME = "header"
"""The name of the script (`str`)
"""

# Use the module docstring for argparse
_DESC = __doc__
"""The program description (`str`)
"""

CSV_KWARGS = ['delimiter', 'doublequote', 'escapechar', 'lineterminator',
              'quotechar', 'quoting', 'skipinitialspace']
"""The keyword (dialect) arguments for the csv module (`list` or `str`)
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():
    """The main entry point for the script. For API use see
    ``something else``.
    """
    # Initialise and parse the command line arguments
    parser = _init_cmd_args()
    args = _parse_cmd_args(parser)

    try:
        open_method = open
        if args.infile is not None and utils.is_gzip(args.infile) is True:
            open_method = gzip.open

        with stdopen.open(args.infile, 'rt', method=open_method,
                          encoding=args.encoding) as infile:
            with stdopen.open(args.outfile, 'wt') as outfile:
                reader = csv.reader(infile, delimiter=args.delimiter)
                header, header_idx = get_header_from_reader(
                    reader, skiplines=args.skip_lines, comment=args.comment
                )
                for idx, i in enumerate(header, int(not args.zero)):
                    outfile.write(f"{idx}\t{i}\n")
    except (BrokenPipeError, KeyboardInterrupt):
        # Python flushes standard streams on exit; redirect remaining
        # output to devnull to avoid another BrokenPipeError at shutdown
        devnull = os.open(os.devnull, os.O_WRONLY)
        os.dup2(devnull, sys.stdout.fileno())


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _init_cmd_args():
    """Initialise the command line arguments and return the parser.

    Returns
    -------
    args : `argparse.ArgumentParser`
        The argparse parser object with arguments added
    """
    parser = argparse.ArgumentParser(
        description=_DESC
    )

    parser.add_argument(
        'infile', type=str, nargs='?',
        help="The path to the input file (can be gzip compressed), if not "
        "given input will be from STDIN"
    )
    parser.add_argument(
        '-o', '--outfile', type=str,
        help="The path to the output file (if not given will default to"
        " STDOUT)"
    )
    parser.add_argument(
        '-d', '--delimiter', type=str, default="\t",
        help="The delimiter of the input file"
    )
    parser.add_argument(
        '-e', '--encoding', type=str, default="UTF8",
        help="The encoding of the input file"
    )
    parser.add_argument(
        '-s', '--skip-lines', type=int, default=0,
        help="The skip this many lines before processing the header"
    )
    parser.add_argument(
        '-c', '--comment', type=str, default="#",
        help="Skip lines that start with this string"
    )
    parser.add_argument(
        '-z', '--zero', action="store_true",
        help="Use zero based number in the output"
    )

    return parser


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _parse_cmd_args(parser):
    """Parse the command line arguments.

    Parameters
    ----------
    parser : `argparse.ArgumentParser`
        The argparse parser object with arguments added.

    Returns
    -------
    args : `argparse.Namespace`
        The argparse namespace object containing the arguments
    """
    args = parser.parse_args()

    # TODO: You can perform checks on the arguments here if you want
    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def detect_header(infile, open_method=None, encoding='utf8', skiplines=0,
                  comment=None, delimiters=None, sniff_bytes=8192,
                  **csv_kwargs):
    """Attempt to detect if the file has a header based on the defined
    arguments.

    Parameters
    ----------
    infile : `str`
        The input file to check.
    open_method : `function`, optional, default: `NoneType`
        A method to use to open the file. If not given then an attempt is made
        to get the open method (i.e. ``builtins.open`` or ``gzip.open``) If
        neither of those are appropriate then the ``open_method`` should be
        given.
    encoding : `str`, optional, default: `utf-8`
        The encoding to use to open the file. If the open method does not take
        an encoding parameter then set this to ``NoneType``.
    skiplines : `int`, optional, default: `0`
        A fixed number of rows to skip before looking for the header.
        White space stripped.
    comment : `str`, optional, default: `#`
        A comment character, if this string starts at the beginning of a line
        before the header then it is skipped. Note that skiplines do not check
        for comment strings.
    delimiters : `str`, optional, default: `NoneType`
        Any potential delimiters that might be used in the file. This string
        can contain more than one. Note, that this is only used if
        ``csv_kwargs`` is not used.
    sniff_bytes : `int`, optional, default: `8192`
        This is the number of bytes that is read in from the file when
        attempting to detect the delimiter and header.
    **csv_kwargs
        Keyword arguments to open the csv file. If provided it will be used to
        extract the header row. If not provided then an attempt will be made to
        sniff the dialect.

    Returns
    -------
    first_row : `list` of `str`
        The first row from the file, this may or may not be the header.
    dialect : `csv.Dialect`
        The dialect detected by the ``csv.Sniffer``.
    has_header : `bool`
        `True` if the sniffer thinks we have a header `False` if not.
    skiplines : `int`
        The total number of lines that were skipped to get to the header line.

    Notes
    -----
    This uses the Python csv module but with some adjustments. If the user has
    supplied a fixed number of lines to skip, then these are skipped first.
    Next if they have supplied a comment character then any lines starting with
    this are skipped. Finally, the csv.Sniffer is deployed to detect the
    dialect and header line.

    See also
    --------
    csv.Sniffer
    """
    if open_method is None:
        open_method = utils.get_open_method(infile)

    kwargs = dict()
    if encoding is not None:
        kwargs['encoding'] = encoding

    # Open the file being agnostic for gzip
    with open_method(infile, 'rt', **kwargs) as csvfile:
        skiprows(csvfile, skiplines)
        last_tell = csvfile.tell()

        if comment is not None:
            # Sort the header
            while True:
                first_row = csvfile.readline().strip()
                if not first_row.startswith(comment):
                    csvfile.seek(last_tell)
                    break
                skiplines += 1
                last_tell = csvfile.tell()

            # print("BACK TO:")
            # print(csvfile.readline().strip())
            csvfile.seek(last_tell)

        # have a sniff for the dialect so I can use this later to return
        # the first row (that may well be the header)
        dialect = None
        if delimiters is not None:
            dialect = csv.Sniffer().sniff(
                csvfile.read(sniff_bytes), delimiters=delimiters
            )
        else:
            dialect = csv.Sniffer().sniff(csvfile.read(sniff_bytes))

        # Move back to the begining and attempt to sniff out the header
        csvfile.seek(last_tell)
        has_header = csv.Sniffer().has_header(
            csvfile.read(sniff_bytes)
        )

        # Now move back to the begining and attempt to get the first row
        csvfile.seek(last_tell)

        reader = None
        if len(csv_kwargs) > 0:
            for i in CSV_KWARGS:
                if i not in csv_kwargs:
                    csv_kwargs[i] = getattr(dialect, i)

            reader = csv.reader(csvfile, **csv_kwargs)
        else:
            reader = csv.reader(csvfile, dialect=dialect)
        first_row = next(reader)

    # return the first row (which may or maynot be a header) and a
    # bool indicating if the sniffer thinks we have a header
    return first_row, dialect, has_header, skiplines


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_header_from_file(fobj, skiplines=0, comment='#'):
    """Extract the first data row from a flat file, this may occur after a
    fixed number of lines and/or comment lines.

    Parameters
    ----------
    fobj : `File`
        A file object.
    skiplines : `int`, optional, default: `0`
        A fixed number of rows to skip before looking for the header.
        White space stripped.
    comment : `str` or `bytes` or `NoneType`, optional, default: `#`
        A comment character, if this string starts at the beginning of a line
        before the header then it is skipped. Note that skiplines do not check
        for comment strings. If there is no comment character set to
        ``NoneType``. If the file has been opened in binary mode make sure the
        the comment character is also binary.


    Returns
    -------
    header : `str`
        The header row.
    header_idx : `int`
        The zero based row index of the header.
    """
    comment = check_comment_char(comment, allow_none=True)
    skiprows(fobj, skiplines)
    first_row = fobj.readline().strip()

    # Sort the header
    if comment is not None:
        while True:
            if not first_row.startswith(comment):
                break
            first_row = fobj.readline().strip()
            skiplines += 1

    return first_row, skiplines


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_header_from_reader(reader, skiplines=0, comment='#'):
    """Extract the first data row from a flat file, this may occur after a
    fixed number of lines and/or comment lines.

    Parameters
    ----------
    reader : `csv.reader`
        A ``csv.reader``.
    skiplines : `int`, optional, default: `0`
        A fixed number of rows to skip before looking for the header.
    comment : `str`, optional, default: `#`
        A comment character, if this string starts at the beginning of a line
        before the header then it is skipped. Note that skiplines do not check
        for comment strings. If there is no comment character set to
        ``NoneType``.

    Returns
    -------
    header : `str` or `list` of `str`
        A string will be returned unless ``fobj`` is a ``csv.reader`` class.
        Each element is white space stripped.
    header_idx : `int`
        The zero based row index of the header.
    """
    comment = check_comment_char(comment, allow_none=True)
    skiprows(reader, skiplines)
    first_row = next(reader)

    test_field = 0

    # Sort the header
    if comment is not None:
        while True:
            if not first_row[test_field].strip().startswith(comment):
                break
            first_row = next(reader)
            skiplines += 1
    return first_row, skiplines


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def move_to_header(fobj, skiplines=0, comment='#', return_skiplines=False):
    """Read through all the fixed skiplines and comment characters so that the
    following row should be the header row but do not read the header row.

    Parameters
    ----------
    fobj : `File`
        A file object.
    skiplines : `int`, optional, default: `0`
        A fixed number of rows to skip before looking for the header.
        White space stripped.
    comment : `str` or `bytes` or `NoneType`, optional, default: `#`
        A comment character, if this string starts at the beginning of a line
        before the header then it is skipped. Note that skiplines do not check
        for comment strings. If there is no comment character set to
        ``NoneType``. If the file has been opened in binary mode make sure the
        the comment character is also binary.
    return_skiplines : `bool`, optional, default: `False`
        The default is to return the number of skipped lines. Set this to True
        to return the actual lines skipped before the header was encountered.

    Returns
    -------
    header_idx : `int`
        The zero based row index of the header.
    """
    comment = check_comment_char(comment, allow_none=True)
    skipped_lines = skiplines
    skiprows(fobj, skiplines)

    last_tell = fobj.tell()
    first_row = fobj.readline().strip()

    # Sort the header
    if comment is not None:
        while True:
            if not first_row.startswith(comment):
                break
            skipped_lines += 1
            last_tell = fobj.tell()
            first_row = fobj.readline().strip()

    fobj.seek(last_tell)
    return skipped_lines


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def skip_to_header(fobj, skiplines=0, comment='#', include_header=False):
    """Read through all the fixed skiplines and comment characters so that the
    following row should be the header row but do not read the header row.

    Parameters
    ----------
    fobj : `File`
        A file object.
    skiplines : `int`, optional, default: `0`
        A fixed number of rows to skip before looking for the header.
        White space stripped.
    comment : `str` or `bytes` or `NoneType`, optional, default: `#`
        A comment character, if this string starts at the beginning of a line
        before the header then it is skipped. Note that skiplines do not check
        for comment strings. If there is no comment character set to
        ``NoneType``. If the file has been opened in binary mode make sure the
        the comment character is also binary.
    include_header : `bool`, optional, default: `False`
        Include the header row in the returned rows. This is to allow for the
        fobj being input from STDIN, where tell and seek are not available.
        So can't roll make to the header row.

    Returns
    -------
    skipped_rows : `list`
        The rows that were skipped.

    Notes
    -----
    This differs from `pyaddons.flat_files.header.move_to_header`, in that the
    actual skipped rows are returned and not the number of skipped rows.
    """
    comment = check_comment_char(comment, allow_none=True)
    skipped_lines = []
    nlines = 0
    while nlines < skiplines:
        row = fobj.readline().strip()
        skipped_lines.append(row)
        nlines += 1

    try:
        last_tell = fobj.tell()
    except UnsupportedOperation:
        if include_header is False:
            raise

    first_row = fobj.readline().strip()

    # Sort the header
    if comment is not None:
        while True:
            if not first_row.startswith(comment):
                break
            skipped_lines.append(first_row)
            try:
                last_tell = fobj.tell()
            except UnsupportedOperation:
                if include_header is False:
                    raise
            first_row = fobj.readline().strip()

    if include_header is False:
        fobj.seek(last_tell)
    else:
        skipped_lines.append(first_row)

    return skipped_lines


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def check_comment_char(comment, allow_none=False):
    """Check that a comment character is valid. This strips it and makes sure
    it is not '' or NoneType (unless allow_none is True).

    Parameters
    ----------
    comment : `str`
        The comment character, this can also be  `NoneType` if allow_none is
        `True`.
    allow_none : `bool`, optional, default: `False`
        Allow `NoneType` to be treated as a valid comment character.

    Returns
    -------
    comment : `str` or `NoneType`
        The validated comment character, if this is a string then str.lstrip()
        will have been applied to it.

    Raises
    ------
    ValueError
        If the comment character is not valid, the invalid comment character
        can be found at `ValueError.args[1]`.
    """
    try:
        proc_comment = comment.lstrip()
    except AttributeError:
        if comment is None and allow_none is True:
            return comment
        raise ValueError(f"Bad comment character: '{comment}'", comment)

    if len(proc_comment) == 0:
        raise ValueError(f"Bad comment character: '{comment}'", comment)

    return proc_comment


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def skiprows(fobj, skiplines):
    """Skip ``skiplines`` rows from the file object.

    Parameters
    ----------
    fobj : `File` or `csv.reader`
        A file object or csv reader.

    Returns
    -------
    next_row : `str` or `list` of `str` or `NoneType`
        The first non skipped row. A string will be returned unless
        ``fobj`` is a ``csv.reader`` class. NoneType will be returned if
        ``skiplines`` is 0.

    Raises
    ------
    IndexError
        If there are not enough rows in the file to skip.
    """
    nlines = 0
    row = None
    while nlines < skiplines:
        try:
            try:
                row = fobj.readline()
            except AttributeError:
                # Could be a reader, so try next
                row = next(fobj)
            nlines += 1
        except StopIteration:
            raise IndexError("skipping too many rows for file")
    return row


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def make_unique_header(header, suffix='_repeat'):
    """Make sure that the column names in header are unique. If they are not
    add a suffix and potentially a number.

    Parameters
    ----------
    header : `list` of `str`
        The header to ensure unique entries.
    suffix : `str`, optional, default: `_repeat`
        A suffix to add to the end of repeated column names.

    Returns
    -------
    header : `list` of `str`
        The header with unique entries, the return is not strictly necessary as
        the list is updated in-place.
    """
    for idx in range(len(header)):
        template = header[idx]
        inc_str = ""
        inc = 1
        while header[idx] in header[idx+1:]:
            header[idx] = f"{template}{suffix}{inc_str}"
            inc_str = str(inc)
            inc += 1
    return header


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def compare_unique_header(header, compare_header, suffix='_repeat'):
    """Make sure that the column names in header are unique and not
    represented in compare_header.

    If they are add a suffix and potentially a number.

    Parameters
    ----------
    header : `list` of `str`
        The header to ensure unique entries.
    compare_header : `list` of `str`
        The header to compare against
    suffix : `str`, optional, default: `_repeat`
        A suffix to add to the end of repeated column names.

    Returns
    -------
    header : `list` of `str`
        The header with unique entries, the return is not strictly necessary as
        the list is updated in-place.
    """
    for idx in range(len(header)):
        template = header[idx]
        inc_str = ""
        inc = 1
        while header[idx] in compare_header:
            header[idx] = f"{template}{suffix}{inc_str}"
            inc_str = str(inc)
            inc += 1

    return header


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def duplicated_columns(header, compare_header):
    """Return a list of repeated column names and numbers between two headers.

    Parameters
    ----------
    header : `list` of `str`
        The header to ensure unique entries.
    compare_header : `list` of `str`
        The header to compare against

    Returns
    -------
    repeat_cols : `list` of (`str`, `int`)
        Repeated columns. The first element of the nested tuples is the column
        name and the second is the column number.
    """
    return [
        (i, compare_header.index(i)) for i in header if i in compare_header
    ]
