"""Command line script and API function to output a specific UMLS SAB to file"""

from umls_tools import __version__, __name__ as pkg_name, log, orm as o, common
from sqlalchemy_config import config as cfg
from sqlalchemy_utils.functions import database_exists
from tqdm import tqdm
from sqlalchemy import and_
from sqlalchemy.exc import ObjectNotExecutableError
from sqlalchemy.orm.exc import NoResultFound, MultipleResultsFound
from texttable import Texttable
import stdopen
import argparse
import os
import sys
import tempfile
import csv
import warnings
# import pprint as pp

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


_PROG_NAME = "umls-output-sab"
"""The name of the program that doubles as the logger name (`str`)
"""
_DESC = __doc__
"""The program description given to argparse (`str`)
"""
_SUB_DELIM = "|"
"""A sub-delimiter for the semantic types in the output (`str`)
"""
_TEMP_DELIM = "\t"
"""A delimiter to use in the temp file.
"""
_LIST_KEYWORD = "list-sabs"
"""The command line keyword to list all of the available SABs (`str`)
"""

_OUT_COLS = [
    o.MrConso.AUI,
    o.MrConso.CUI,
    o.MrConso.LUI,
    o.MrConso.SUI,
    o.MrConso.TTY,
    o.MrConso.SAB,
    o.MrConso.CODE,
    o.MrConso.STR,
    o.MrSty.TUI,
    o.MrSty.STY,
    o.MrRank.MRRANK_RANK,
]
_OUT_HEADER = [i.name for i in _OUT_COLS]
"""The header row of the SAB output (`list` of `str`)
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():  # pragma: no cover
    """The main entry point. For API access see
    ``umls_tools.admin.unique_cui.create_unique_cui``
    """
    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)

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

    # 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,
    )

    try:
        # Make sure the database exists (or does not exist)
        if database_exists(sm().bind.url) is False:
            raise FileNotFoundError("database does not exist")
    except ObjectNotExecutableError:
        warnings.warn(
            "unable to check if database exists: sqlalchemy-utils does"
            " not work well with SQLAlchemy >= 1.4"
        )

    tmpobj, tmpfile = tempfile.mkstemp(dir=args.tmpdir)
    os.close(tmpobj)

    if args.sab == _LIST_KEYWORD:
        session = sm()
        try:
            list_sabs(sm())
        finally:
            session.close()
        sys.exit(0)

    tqdm_kwargs = dict(
        unit=" rows", desc=f"[info] extracting {args.sab}", disable=not verbose
    )

    try:
        session = sm()
        logger.info(f"extracting {args.sab}...")

        with stdopen.open(
            args.outfile, mode="wt", use_temp=True, tmpdir=args.tmpdir
        ) as outfile:
            writer = csv.DictWriter(
                outfile,
                _OUT_HEADER,
                delimiter=args.delimiter,
                lineterminator=os.linesep,
            )
            writer.writeheader()
            for row in tqdm(output_sab(session, args.sab), **tqdm_kwargs):
                writer.writerow(row)
        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)
    finally:
        session.close()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
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(
        "sab",
        type=str,
        help="The UMLS source abbreviation of the vocabulary to output. To see"
        " a list of all the SABs in your UMLS database use the keyword "
        f"'{_LIST_KEYWORD}' here",
    )
    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",
        "--tmpdir",
        type=str,
        help="The output is written to tmp first before being moved to the"
        " final location (unless outputting to STDOUT)",
    )
    parser.add_argument(
        "-O",
        "--outfile",
        type=str,
        help="A output file name, if not supplied output will be to STDOUT",
    )
    parser.add_argument(
        "-d",
        "--delimiter",
        type=str,
        default="\t",
        help="The delimiter of the output",
    )
    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 list_sabs(session):
    """List all the SABs in the UMLS version to STDOUT.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The SQLAchemy session used to interact with the database.
    """
    sql = session.query(o.MrSab.RSAB, o.MrSab.VSAB, o.MrSab.SON).order_by(
        o.MrSab.RSAB
    )

    tt = Texttable()
    tt.set_cols_align(["l", "c", "r"])
    sabs = [["SAB", "version (VSAB)", "description"]]
    for row in sql:
        sabs.append([row.RSAB, row.VSAB, row.SON])
    tt.add_rows(sabs, header=True)
    print(tt.draw())


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def output_sab(session, sab):
    """The main API entry point to output an SAB.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The SQLAchemy session used to interact with the database. Must return
        the engine from a ``Session.get_bind()`` call.
    sab : `str`
        The source vocabulary that you want to output.

    Yields
    ------
    sab_row : `dict`
        A rows from the source vocabulary.

    Raises
    ------
    ValueError
        If the SAB does not exist in the UMLS.
    """
    # First make sure the sab exists
    try:
        session.query(o.MrSab.RSAB).filter(o.MrSab.RSAB == sab).one()
    except NoResultFound as e:
        raise ValueError(f"unknown SAB: {sab}") from e
    except MultipleResultsFound:
        pass

    sql = (
        session.query(*_OUT_COLS)
        .join(o.MrSty, o.MrConso.CUI == o.MrSty.CUI)
        .join(
            o.MrRank,
            and_(o.MrConso.SAB == o.MrRank.SAB, o.MrConso.TTY == o.MrRank.TTY),
        )
        .filter(o.MrConso.SAB == sab)
        .order_by(o.MrConso.CUI, o.MrConso.CODE)
    )

    for row in sql:
        yield row._asdict()


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