"""Build a relational database from the UMLS flat file release that is stored
within zipped archive. Additionally load up semantic network flat files. This
uses SQL Alchemy to give some choice of backend. However, thus far it has only
been tested on SQLite and MariaDB. The loading process might need to be adapted
for other databases.
"""

from umls_tools import __version__, __name__ as pkg_name, log, orm, common
from sqlalchemy.exc import OperationalError
from sqlalchemy.engine import Connection, Engine
from sqlalchemy.types import Float, Integer
from sqlalchemy_config import config as cfg
from zipfile import ZipFile
from collections import namedtuple
from contextlib import nullcontext
from tqdm import tqdm
import shutil
import glob
import argparse
import os
import re
import tempfile
import sys
import csv
# import pprint as pp


# deal with long fields in csv
csv.field_size_limit(sys.maxsize)

_PROG_NAME = "umls-build"
"""The name of the program that doubles as the logger name (`str`)
"""
_DESC = __doc__
"""The program description given to argparse (`str`)
"""

FileInfo = namedtuple("FileInfo", ["path", "type", "version", "date"])
"""Will hold information on the ZIP archive and the XML files extracted from
the zip file (`FileInfo`)
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():  # pragma: no cover
    """The main entry point. For API access see
    ``umls_tools.admin.build.build_umls``
    """
    parser = _init_cmd_args()
    args = _parse_cmd_args(parser)

    logger = log.init_logger(_PROG_NAME, verbose=args.verbose)
    log.log_prog_name(logger, pkg_name, __version__)
    log.log_args(logger, args)

    # Get a sessionmaker to create sessions to interact with the database
    sm = cfg.get_sessionmaker(
        args.config_section,
        common.DEFAULT_PREFIX,
        url_arg=args.dburl,
        config_arg=args.config,
        config_env=None,
        config_default=common.DEFAULT_CONFIG,
    )

    load_verbose = False
    if args.verbose > 1:
        load_verbose = True

    # Make sure the database exists (or does not exist)
    cfg.create_db(sm, exists=args.exists)

    try:
        logger.info("building the UMLS, this may take some time...")
        build_umls(
            sm,
            args.indata,
            tmpdir=args.tmp,
            verbose=load_verbose,
            commit_every=args.commit_every,
            exists=args.exists,
        )
        log.log_end(logger)
    except (OSError, FileNotFoundError):
        raise
    except (BrokenPipeError, IOError, 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)

        try:
            os.dup2(devnull, sys.stdout.fileno())
        except Exception:
            pass
        log.log_interrupt(logger)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _init_cmd_args():  # pragma: no cover
    """Use argparse to parse the command line arguments

    Returns
    -------
    parser : `argparse.ArgumentParser`
        The arguments expected on the command line.
    """
    parser = argparse.ArgumentParser(description=_DESC)

    parser.add_argument(
        "indata",
        type=str,
        help="The input zip archive downloaded from the UMLS , do not unzip",
    )
    parser.add_argument(
        "dburl",
        nargs="?",
        type=str,
        help="An SQLAlchemy connection URL or filename if "
        "using SQLite. If you do not want to put full "
        "connection parameters on the cmd-line use the "
        "config file (--config) and config section"
        "(--config-section) to supply the parameters",
    )
    parser.add_argument(
        "-c",
        "--config",
        type=str,
        default=common.get_default_config_cli_path(),
        help="The location of the config file",
    )
    parser.add_argument(
        "-s",
        "--config-section",
        type=str,
        default=common.DEFAULT_SECTION,
        help="The section name in the config file",
    )
    parser.add_argument(
        "-T",
        "--tmp",
        type=str,
        help="The location of tmp, if not provided will use the system tmp",
    )
    parser.add_argument(
        "--commit-every",
        type=int,
        default=10000,
        help="The commit to the database after every --commit-every rows",
    )
    parser.add_argument(
        "-e",
        "--exists",
        action="store_true",
        help="The database already exists, this will assume"
        " that the user has pre-created a database and will "
        "not throw errors if the DB/tables already exist ",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        default=0,
        action="count",
        help="give more output, ``--vv`` turns on progress monitoring",
    )
    return parser


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _parse_cmd_args(parser):  # pragma: no cover
    """Use argparse to parse the command line arguments.

    Parameters
    ----------
    parser : `argparse.ArgumentParser`
        The arguments expected on the command line.

    Returns
    -------
    args : `ArgumentParser.args`
        The arguments from parsing the cmd line args.
    """
    args = parser.parse_args()
    # Because the default is ~/ for Sphinx docs
    args.config = os.path.expanduser(args.config)
    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_umls(
    sm, zip_file, tmpdir=None, verbose=False, commit_every=10000, exists=False
):
    """Kick of the build process for the UMLS database. This should be the
    main API entry point for any external programs using this.

    Parameters
    ----------
    sm : `sqlalchemy.Sessionmaker`
        The SQLAchemy session maker to use to create sessions.
    zip_file : `str`
        The path to the downloaded DM+D zip archive. This should not be renamed
        or extracted. It should be exactly as downloaded from TRUD.
        It will have a format like:
        ``nhsbsa_dmd_<major version>.<minor version>.<patch>_YYYYMMDD<some other digits>.zip``
    tmpdir : `str` or `NoneType`, optional, default: `NoneType`
        The tmp directory to extract the ZIP archive into. If ``NoneType`` then
        the system tmp directory is used. Note that the zip archives will be
        extracted into sub directories within here. These will be deleted when
        this exits (either through failure of success).
    verbose : `bool`, optional, default: `False`
        Show progress of the individual table loads.
    commit_every : `int, optional, default: `10000`
        Batch size passed to ``load_table``. Rows are buffered and inserted
        via connection ``execute`` + ``commit``. If set to 0 or negative,
        rows are committed in a single final batch.
    exists : `bool`, optional, default: `False`
        Should the database be expected to already exist? If `False`, then
        existing databases will cause and error, if True then no error will be
        raised. This is to handle situations where the uer is unable to create
        databases.
    """
    # Create a tmp output directory to zip into
    zip_out = tempfile.mkdtemp(dir=tmpdir)

    try:
        # Extract the information from the zip file
        zip_info = get_zip_file_data(zip_file)

        # Now extract the zip files from the archive
        files = extract_data(zip_file, zip_info, out_dir=zip_out)

        for infile, load_opts in tqdm(
            orm.CLASS_MAP.items(),
            disable=not verbose,
            unit=" tables",
            desc="[info] loading tables",
        ):
            load_class, expected = load_opts
            try:
                load_file = files[infile]
                session = sm()
                try:
                    load_table(
                        session,
                        load_file,
                        load_class,
                        delimiter="|",
                        commit_every=commit_every,
                        verbose=verbose,
                        exists=exists,
                    )
                finally:
                    session.close()
            except KeyError:
                if expected is True:
                    raise FileNotFoundError(
                        f"can't find required file: {infile}"
                    )
    finally:
        shutil.rmtree(zip_out)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_zip_file_data(zip_file):
    """Extract the version and file type information from the file name of the
    zip file.

    Parameters
    ----------
    zip_file : `str`
        The path to the zip file.

    Returns
    -------
    file_info : `FileInfo`
        A ``FileInfo`` named tuple with the ``type`` attribute ``zip`` or
        ``bonus_zip``.

    Notes
    -----
    The expectation is that the zip file base name will have the structure
    below:

    ``umls-<year as YYYY><build AA or AB>-metathesaurus.zip``
    """
    # /data/dmd/nhsbsa_dmd_11.4.0_20211129000001.zip
    fn_regex = re.compile(
        r"""
        umls- # prefix
        (?P<YEAR>\d\d\d\d) # The major version number and delimiter
        (?P<BUILD>A(A|B)) # The minor version number
        -metathesaurus.zip
        """,
        re.VERBOSE,
    )

    base_zip = os.path.basename(zip_file)
    m = fn_regex.match(base_zip)

    try:
        version = "{0}{1}".format(m.group("YEAR"), m.group("BUILD"))
    except AttributeError as e:
        raise ValueError(f"unable to parse zip file name: {base_zip}") from e
    file_type = "zip"
    fi = FileInfo(
        path=zip_file, type=file_type, version=version, date=m.group("YEAR")
    )
    return fi


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def extract_data(inzip, zip_info, out_dir):
    """Extract all the files from the ZIP archive.

    Parameters
    ----------
    inzip : `str`
        The input ZIP archive.
    zip_info : `umls_tools.admin.build.FileInfo`
        A ``FileInfo`` named tuple with the ``type`` attribute ``zip`` or
        ``bonus_zip``.
    out_dir : `str`
        The directory to extract the ZIP archive into.

    Returns
    -------
    extracted_files : `list` or `str`
        the paths to the files extracted from the archive.
    """
    with ZipFile(inzip, "r") as zip:
        # extract all files to another directory
        zip.extractall(out_dir)
    glob_dir = os.path.join(out_dir, zip_info.version, "META")
    files = glob.glob(os.path.join(glob_dir, "*.RRF"))
    files = dict([(os.path.basename(i), i) for i in files])
    return files


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def load_table(
    session,
    infile,
    load_class,
    delimiter="|",
    commit_every=10000,
    verbose=False,
    exists=False,
    encoding="utf-8",
):
    """Load the contents of a flat file into the UMLS.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The SQLAchemy session used to interact with the database. The
        ``Session.get_bind()`` call must return an ``Engine`` or
        ``Connection``.
    infile : `str`
        The input flat file that will be inserted into the database using
        ``load_class``.
    load_class : `umls_tools.orm.Base`
        A UMLS ORM class for the table being imported.
    delimiter : `str`, optional, default: `|`
        The delimiter of the flat file.
    commit_every : `int, optional, default: `10000`
        Number of rows to buffer before executing an INSERT batch and
        committing the active database connection. If set to 0 or negative,
        rows are inserted in one final batch.
    verbose : `bool`, optional, default: `False`
        Show progress of the individual table loads.
    exists : `bool`, optional, default: `False`
        Should the table be expected to already exist? If ``False``, then
        existing table will cause and error, if ``True`` then no error will be
        raised. This is to handle situations where the user is unable to create
        databases and tables.
    encoding : `str`, optional, default: `utf-8`
        Text encoding used to read the flat file.

    Notes
    -----
    For performance purposes this avoids the ORM Session for inserts and will
    execute directly using a database connection.

    If exists is ``True``, then a count of the rows in the table will performed
    and the infile will be traversed to the next row after that and start the
    import from there (assuming there is a next row and we do not reach the end
    of the file).

    See also
    --------
    `sqlalchemy.engine.Connection.execute`
    """
    bind = session.get_bind()

    if isinstance(bind, Engine):
        connection_context = bind.connect()
    elif isinstance(bind, Connection):
        connection_context = nullcontext(bind)
    else:
        raise TypeError(
            "session.get_bind() must return an Engine or Connection"
        )

    if exists is True:
        try:
            load_class.__table__.create(bind=bind, checkfirst=True)
        except OperationalError as e:
            if re.search(r"already exists", str(e), re.IGNORECASE) is None:
                raise
    else:
        load_class.__table__.create(bind=bind, checkfirst=False)

    with connection_context as conn:
        with open(infile, "rt", encoding=encoding) as incsv:
            if exists is True:
                # See if we need to pickup the import
                goto_row(session, load_class, incsv, verbose=verbose)
            reader = csv.reader(incsv, delimiter=delimiter)
            basefile = os.path.basename(infile)
            # header = next(reader)
            # header = [orm.get_colname(i) for i in header]
            columns = list(load_class.__table__.c)[1:]
            header = [i.name for i in columns]
            load_buffer = []
            batch_size = commit_every if commit_every > 0 else None
            for idx, row in enumerate(
                tqdm(
                    reader,
                    disable=not verbose,
                    unit=" rows",
                    desc=f"[info] processing {basefile}",
                    leave=False,
                ),
                1,
            ):
                # Create a dictionary of column names => data values
                kwargs = dict(
                    [
                        (header[idx], set_column_data(i, row[idx]))
                        for idx, i in enumerate(columns)
                    ]
                )
                load_buffer.append(kwargs)
                if batch_size is not None and idx % batch_size == 0:
                    conn.execute(load_class.__table__.insert(), load_buffer)
                    conn.commit()
                    load_buffer = []
            if len(load_buffer) > 0:
                conn.execute(load_class.__table__.insert(), load_buffer)
                conn.commit()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def goto_row(session, load_class, incsv, verbose=False):
    """Shoot through the file until we reach the last row that was imported
    into the database (based on row number). The import will continue from
    that position.

    Parameters
    ----------

    """
    offset = 0
    # TODO: remove this MRCONSO fudge
    if load_class == orm.MrConso:
        offset = 1

    nrows = session.query(load_class).count()
    # If there are 0 rows in the table then just return
    if nrows == 0:
        return

    nrows += offset

    tqdm_kwargs = dict(
        unit=" rows",
        desc=f"[info] navigating to row '{nrows}'",
        total=nrows,
        leave=False,
    )
    for i in tqdm(range(nrows), **tqdm_kwargs):
        next(incsv)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def set_column_data(column, data):
    """Cast the data according to SQLAlchemy column types.

    Parameters
    ----------
    column : `sqlalchemy.Column`
        An SQLAchemy column object.
    data : `str`
        The data value to cast.

    Notes
    -----
    Integer and float casting use ``isinstance`` checks against SQLAlchemy
    type classes.
    """
    if data == "":
        data = None
    elif isinstance(column.type, Integer):
        data = int(data)
    elif isinstance(column.type, Float):
        data = float(data)
    return data


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
if __name__ == "__main__":
    main()
