"""
Handles functionality related to data storege.
"""
import sys, os, glob, re, gzip, json
from biorun import const, utils, objects, ncbi
from biorun.models import jsonrec
import biorun.libs.placlib as plac

# Module level logger.
logger = utils.logger

# A nicer error message on incorrect installation.
try:
    from Bio import SeqIO
except ImportError as exc:
    print(f"*** Error: {exc}", file=sys.stderr)
    print(f"*** This program requires biopython", file=sys.stderr)
    print(f"*** Install: conda install -y biopython>=1.78", file=sys.stderr)
    sys.exit(-1)

def resolve_fname(name, format='json'):
    """
    Resolve a file name given an accession number.
    """
    ext = format.lower()
    fname = f"{name}.{ext}.gz"
    fname = os.path.join(utils.DATADIR, fname)
    return fname


def delete_data(text):
    """
    Deletes data under a filename.
    """
    for name in text.split(","):
        fname = resolve_fname(name)
        if os.path.isfile(fname):
            os.remove(fname)
            logger.info(f"removed: {fname}")
        else:
            logger.info(f"file does not exist: {fname}")


def read_json_file(fname):
    """
    Returns the content of a JSON file.
    """
    fp = utils.gz_read(fname)
    data = json.load(fp)
    fp.close()
    return data


def save_json_file(fname, data):
    """
    Returns the content of a JSON file.
    """
    fp = utils.gz_write(fname)
    json.dump(data, fp)
    fp.close()
    logger.info(f"saved {fname}")
    return data


def change_seqid(json_name, seqid):
    """
    Changes the sequence id stored in a json file.
    """
    if os.path.isfile(json_name):
        data = read_json_file(json_name)
        for item in data:
            item[const.SEQID] = seqid
        fp = utils.gz_write(json_name)
        json.dump(data, fp)
        fp.close()


def fetch_data(data, param):
    """
    Obtains data from NCBI. Fills each parameter with a json field.
    """

    db = "protein" if param.protein else "nuccore"

    # Ensure json DB is built
    ncbi.build_db()

    genbank, taxon_acc, refseq = ncbi.get_data()

    for name in data:

        # Pretend no data if it is an update.
        json = None if param.update else get_json(name)

        # The data exists, nothing needs to be done.
        if json:
            continue

        # The JSON representation of the data.
        json_name = resolve_fname(name=name, format="json")

        # GenBank representation of the data.
        gbk_name = resolve_fname(name=name, format="gb")

        # Genome assembly data.
        if name.startswith("GCA") or name.startswith("GCF"):

            ncbi.genome(name=name, fname=gbk_name, update=param.update, genbank=genbank,
                        refseq=refseq)
        else:
            # Genbank data.
            ncbi.genbank_save(name, db=db, fname=gbk_name)

        # Convert Genbank to JSON.
        data = jsonrec.parse_file(fname=gbk_name, seqid=param.seqid)

        # Save JSON file.
        save_json_file(fname=json_name, data=data)


def genbank_view(params):
    for param in params:
        altname = resolve_fname(param.acc, format="gb")

        if os.path.isfile(param.acc):
            stream = utils.gz_read(param.acc)
        elif os.path.isfile(altname):
            stream = utils.gz_read(altname)
        else:
            stream = []
            utils.error(f"data not found: {param.acc}")

        for line in stream:
            print(line, end='')


def get_json(name, seqid=None, inter=False, strict=False):
    """
    Attempts to return a JSON formatted data based on a name.
    """

    # Data is an existing path to a JSON file.
    if os.path.isfile(name):
        try:
            data = jsonrec.parse_file(name, seqid=seqid)
        except Exception as exc:
            logger.error(f"JSON parsing error for file {name}: {exc}")
            sys.exit(-1)
        return data

    # The JSON representation of the data.
    json_name = resolve_fname(name=name, format="json")

    # GenBank representation of the data.
    gbk_name = resolve_fname(name=name, format="gb")

    # Found the JSON representation of the file.
    if os.path.isfile(json_name):
        logger.info(f"found {json_name}")
        data = read_json_file(json_name)
        return data

    # There is no JSON file but there is a GenBank file.
    if os.path.isfile(gbk_name):
        logger.info(f"found {gbk_name}")
        data = jsonrec.parse_file(fname=gbk_name, seqid=seqid)
        data = save_json_file(fname=json_name, data=data)
        return data

    # Interactive input, make JSON from name
    if inter:
        data = jsonrec.make_jsonrec(name, seqid=seqid)
        return data

    # Raise error if in strict mode
    if strict:
        utils.error(f"data not found: {name}")
    return None


def rename_data(data, param, newname=None):
    """
    Rename data.
    """
    # Will only rename a single data
    newnames = newname.split(",")

    for name1, name2 in zip(data, newnames):
        src_json = resolve_fname(name=name1, format="json")
        dest_json = resolve_fname(name=name2, format="json")

        src_gb = resolve_fname(name=name1, format="gb")
        dest_gb = resolve_fname(name=name2, format="gb")

        if os.path.isfile(src_json):
            logger.info(f"renamed {name1} as {name2}")
            os.rename(src_json, dest_json)
            if param.seqid:
                change_seqid(dest_json, seqid=param.seqid)
        else:
            logger.info(f"file not found: {src_json}")

        if os.path.isfile(src_gb):
            if not os.path.isfile(dest_gb):
                os.symlink(src_gb, dest_gb)
        else:
            logger.info(f"file not found: {src_gb}")


def print_data_list():
    """
    Returns a list of the files in the data directory
    """
    pattern = os.path.join(os.path.join(utils.DATADIR, '*.json.gz'))
    matched = glob.glob(pattern)

    # Extract the definition from the JSON without parsing it.
    patt = re.compile(r'(definition\":\s*)(?P<value>\".+?\")')
    collect = []

    for path in matched:
        fsize = utils.human_size(os.path.getsize(path))
        base, fname = os.path.split(path)
        fname = fname.rsplit(".", maxsplit=2)[0]

        # Parse the first N lines
        stream = gzip.open(path, 'rt') if path.endswith('gz') else open(path, 'rt')
        text = stream.read(1000)
        match = patt.search(text)

        title = match.group("value") if match else ''
        title = title.strip('", ')

        # Trim the title
        stitle = title[:100]
        stitle = stitle + "..." if len(title) != len(stitle) else stitle

        collect.append((str(fsize), f"{fname:10s}", stitle))

    collect = sorted(collect, key=lambda x: x[2])
    for row in collect:
        line = "\t".join(row)
        print(line)


@plac.pos("data", "data names")
@plac.flg('fetch', "download data as accessions")
@plac.flg('update', "updates data in storage")
@plac.opt('rename', "rename the data")
@plac.opt('seqid', "set the sequence id of the data")
@plac.flg('protein', "use the protein database")
@plac.flg('build', "build the database")
@plac.flg('verbose', "verbose mode")
def run(update=False, rename='', seqid='', protein=False, verbose=False, *data):
    """
    Fetches and manages data in storage.
    """

    # Set the verbosity
    utils.set_verbosity(logger, level=int(verbose))

    # Reset counter (needed for consistency during testing).
    jsonrec.reset_counter()

    # A simple wrapper class to represent input parameters.
    param = objects.Param(seqid=seqid, rename=rename, start=1, protein=protein, update=update)

    # Fetch the data.
    fetch_data(data, param=param)

    # Renaming after fetching.
    if rename:
        rename_data(data, param=param, newname=rename)


@plac.opt('delete', "deletes foo from storage", metavar='foo')
@plac.flg('verbose', "verbose mode")
def data(delete, verbose=False):
    """
    Shows the data in the storage.

    Usage:

        bio data                   : lists the data
        bio data --delete foo      : deletes data called foo
        bio data --delete foo,bar  : deletes multiple datasets
    """
    # Set the verbosity
    utils.set_verbosity(logger, level=int(verbose))

    # Reset counter (needed for consistency during testing).
    jsonrec.reset_counter()

    # Delete should be the first to execute.
    if delete:
        delete_data(delete)
    else:
        # Prints the data listing.
        print_data_list()
[((246, 1, 246, 31), 'biorun.libs.placlib.pos', 'plac.pos', ({(246, 10, 246, 16): '"""data"""', (246, 18, 246, 30): '"""data names"""'}, {}), "('data', 'data names')", True, 'import biorun.libs.placlib as plac\n'), ((247, 1, 247, 49), 'biorun.libs.placlib.flg', 'plac.flg', ({(247, 10, 247, 17): '"""fetch"""', (247, 19, 247, 48): '"""download data as accessions"""'}, {}), "('fetch', 'download data as accessions')", True, 'import biorun.libs.placlib as plac\n'), ((248, 1, 248, 46), 'biorun.libs.placlib.flg', 'plac.flg', ({(248, 10, 248, 18): '"""update"""', (248, 20, 248, 45): '"""updates data in storage"""'}, {}), "('update', 'updates data in storage')", True, 'import biorun.libs.placlib as plac\n'), ((249, 1, 249, 38), 'biorun.libs.placlib.opt', 'plac.opt', ({(249, 10, 249, 18): '"""rename"""', (249, 20, 249, 37): '"""rename the data"""'}, {}), "('rename', 'rename the data')", True, 'import biorun.libs.placlib as plac\n'), ((250, 1, 250, 53), 'biorun.libs.placlib.opt', 'plac.opt', ({(250, 10, 250, 17): '"""seqid"""', (250, 19, 250, 52): '"""set the sequence id of the data"""'}, {}), "('seqid', 'set the sequence id of the data')", True, 'import biorun.libs.placlib as plac\n'), ((251, 1, 251, 48), 'biorun.libs.placlib.flg', 'plac.flg', ({(251, 10, 251, 19): '"""protein"""', (251, 21, 251, 47): '"""use the protein database"""'}, {}), "('protein', 'use the protein database')", True, 'import biorun.libs.placlib as plac\n'), ((252, 1, 252, 40), 'biorun.libs.placlib.flg', 'plac.flg', ({(252, 10, 252, 17): '"""build"""', (252, 19, 252, 39): '"""build the database"""'}, {}), "('build', 'build the database')", True, 'import biorun.libs.placlib as plac\n'), ((253, 1, 253, 36), 'biorun.libs.placlib.flg', 'plac.flg', ({(253, 10, 253, 19): '"""verbose"""', (253, 21, 253, 35): '"""verbose mode"""'}, {}), "('verbose', 'verbose mode')", True, 'import biorun.libs.placlib as plac\n'), ((276, 1, 276, 62), 'biorun.libs.placlib.opt', 'plac.opt', (), '', True, 'import biorun.libs.placlib as plac\n'), ((277, 1, 277, 36), 'biorun.libs.placlib.flg', 'plac.flg', ({(277, 10, 277, 19): '"""verbose"""', (277, 21, 277, 35): '"""verbose mode"""'}, {}), "('verbose', 'verbose mode')", True, 'import biorun.libs.placlib as plac\n'), ((27, 12, 27, 46), 'os.path.join', 'os.path.join', ({(27, 25, 27, 38): 'utils.DATADIR', (27, 40, 27, 45): 'fname'}, {}), '(utils.DATADIR, fname)', False, 'import sys, os, glob, re, gzip, json\n'), ((48, 9, 48, 29), 'biorun.utils.gz_read', 'utils.gz_read', ({(48, 23, 48, 28): 'fname'}, {}), '(fname)', False, 'from biorun import const, utils, objects, ncbi\n'), ((49, 11, 49, 24), 'json.load', 'json.load', ({(49, 21, 49, 23): 'fp'}, {}), '(fp)', False, 'import sys, os, glob, re, gzip, json\n'), ((58, 9, 58, 30), 'biorun.utils.gz_write', 'utils.gz_write', ({(58, 24, 58, 29): 'fname'}, {}), '(fname)', False, 'from biorun import const, utils, objects, ncbi\n'), ((59, 4, 59, 23), 'json.dump', 'json.dump', ({(59, 14, 59, 18): 'data', (59, 20, 59, 22): 'fp'}, {}), '(data, fp)', False, 'import sys, os, glob, re, gzip, json\n'), ((69, 7, 69, 32), 'os.path.isfile', 'os.path.isfile', ({(69, 22, 69, 31): 'json_name'}, {}), '(json_name)', False, 'import sys, os, glob, re, gzip, json\n'), ((86, 4, 86, 19), 'biorun.ncbi.build_db', 'ncbi.build_db', ({}, {}), '()', False, 'from biorun import const, utils, objects, ncbi\n'), ((88, 33, 88, 48), 'biorun.ncbi.get_data', 'ncbi.get_data', ({}, {}), '()', False, 'from biorun import const, utils, objects, ncbi\n'), ((143, 7, 143, 27), 'os.path.isfile', 'os.path.isfile', ({(143, 22, 143, 26): 'name'}, {}), '(name)', False, 'import sys, os, glob, re, gzip, json\n'), ((158, 7, 158, 32), 'os.path.isfile', 'os.path.isfile', ({(158, 22, 158, 31): 'json_name'}, {}), '(json_name)', False, 'import sys, os, glob, re, gzip, json\n'), ((164, 7, 164, 31), 'os.path.isfile', 'os.path.isfile', ({(164, 22, 164, 30): 'gbk_name'}, {}), '(gbk_name)', False, 'import sys, os, glob, re, gzip, json\n'), ((215, 14, 215, 32), 'glob.glob', 'glob.glob', ({(215, 24, 215, 31): 'pattern'}, {}), '(pattern)', False, 'import sys, os, glob, re, gzip, json\n'), ((218, 11, 218, 62), 're.compile', 're.compile', ({(218, 22, 218, 61): '"""(definition\\\\":\\\\s*)(?P<value>\\\\".+?\\\\")"""'}, {}), '(\'(definition\\\\":\\\\s*)(?P<value>\\\\".+?\\\\")\')', False, 'import sys, os, glob, re, gzip, json\n'), ((263, 4, 263, 27), 'biorun.models.jsonrec.reset_counter', 'jsonrec.reset_counter', ({}, {}), '()', False, 'from biorun.models import jsonrec\n'), ((266, 12, 266, 94), 'biorun.objects.Param', 'objects.Param', (), '', False, 'from biorun import const, utils, objects, ncbi\n'), ((292, 4, 292, 27), 'biorun.models.jsonrec.reset_counter', 'jsonrec.reset_counter', ({}, {}), '()', False, 'from biorun.models import jsonrec\n'), ((19, 4, 19, 16), 'sys.exit', 'sys.exit', ({(19, 13, 19, 15): '(-1)'}, {}), '(-1)', False, 'import sys, os, glob, re, gzip, json\n'), ((37, 11, 37, 32), 'os.path.isfile', 'os.path.isfile', ({(37, 26, 37, 31): 'fname'}, {}), '(fname)', False, 'import sys, os, glob, re, gzip, json\n'), ((73, 13, 73, 38), 'biorun.utils.gz_write', 'utils.gz_write', ({(73, 28, 73, 37): 'json_name'}, {}), '(json_name)', False, 'from biorun import const, utils, objects, ncbi\n'), ((74, 8, 74, 27), 'json.dump', 'json.dump', ({(74, 18, 74, 22): 'data', (74, 24, 74, 26): 'fp'}, {}), '(data, fp)', False, 'import sys, os, glob, re, gzip, json\n'), ((115, 15, 115, 68), 'biorun.models.jsonrec.parse_file', 'jsonrec.parse_file', (), '', False, 'from biorun.models import jsonrec\n'), ((125, 11, 125, 36), 'os.path.isfile', 'os.path.isfile', ({(125, 26, 125, 35): 'param.acc'}, {}), '(param.acc)', False, 'import sys, os, glob, re, gzip, json\n'), ((166, 15, 166, 62), 'biorun.models.jsonrec.parse_file', 'jsonrec.parse_file', (), '', False, 'from biorun.models import jsonrec\n'), ((172, 15, 172, 54), 'biorun.models.jsonrec.make_jsonrec', 'jsonrec.make_jsonrec', (), '', False, 'from biorun.models import jsonrec\n'), ((177, 8, 177, 46), 'biorun.utils.error', 'utils.error', ({(177, 20, 177, 45): 'f"""data not found: {name}"""'}, {}), "(f'data not found: {name}')", False, 'from biorun import const, utils, objects, ncbi\n'), ((195, 11, 195, 35), 'os.path.isfile', 'os.path.isfile', ({(195, 26, 195, 34): 'src_json'}, {}), '(src_json)', False, 'import sys, os, glob, re, gzip, json\n'), ((203, 11, 203, 33), 'os.path.isfile', 'os.path.isfile', ({(203, 26, 203, 32): 'src_gb'}, {}), '(src_gb)', False, 'import sys, os, glob, re, gzip, json\n'), ((214, 27, 214, 67), 'os.path.join', 'os.path.join', ({(214, 40, 214, 53): 'utils.DATADIR', (214, 55, 214, 66): '"""*.json.gz"""'}, {}), "(utils.DATADIR, '*.json.gz')", False, 'import sys, os, glob, re, gzip, json\n'), ((223, 22, 223, 41), 'os.path.split', 'os.path.split', ({(223, 36, 223, 40): 'path'}, {}), '(path)', False, 'import sys, os, glob, re, gzip, json\n'), ((38, 12, 38, 28), 'os.remove', 'os.remove', ({(38, 22, 38, 27): 'fname'}, {}), '(fname)', False, 'import sys, os, glob, re, gzip, json\n'), ((108, 12, 109, 38), 'biorun.ncbi.genome', 'ncbi.genome', (), '', False, 'from biorun import const, utils, objects, ncbi\n'), ((112, 12, 112, 58), 'biorun.ncbi.genbank_save', 'ncbi.genbank_save', (), '', False, 'from biorun import const, utils, objects, ncbi\n'), ((126, 21, 126, 45), 'biorun.utils.gz_read', 'utils.gz_read', ({(126, 35, 126, 44): 'param.acc'}, {}), '(param.acc)', False, 'from biorun import const, utils, objects, ncbi\n'), ((127, 13, 127, 36), 'os.path.isfile', 'os.path.isfile', ({(127, 28, 127, 35): 'altname'}, {}), '(altname)', False, 'import sys, os, glob, re, gzip, json\n'), ((145, 19, 145, 56), 'biorun.models.jsonrec.parse_file', 'jsonrec.parse_file', (), '', False, 'from biorun.models import jsonrec\n'), ((197, 12, 197, 42), 'os.rename', 'os.rename', ({(197, 22, 197, 30): 'src_json', (197, 32, 197, 41): 'dest_json'}, {}), '(src_json, dest_json)', False, 'import sys, os, glob, re, gzip, json\n'), ((222, 33, 222, 54), 'os.path.getsize', 'os.path.getsize', ({(222, 49, 222, 53): 'path'}, {}), '(path)', False, 'import sys, os, glob, re, gzip, json\n'), ((227, 17, 227, 38), 'gzip.open', 'gzip.open', ({(227, 27, 227, 31): 'path', (227, 33, 227, 37): '"""rt"""'}, {}), "(path, 'rt')", False, 'import sys, os, glob, re, gzip, json\n'), ((128, 21, 128, 43), 'biorun.utils.gz_read', 'utils.gz_read', ({(128, 35, 128, 42): 'altname'}, {}), '(altname)', False, 'from biorun import const, utils, objects, ncbi\n'), ((131, 12, 131, 55), 'biorun.utils.error', 'utils.error', ({(131, 24, 131, 54): 'f"""data not found: {param.acc}"""'}, {}), "(f'data not found: {param.acc}')", False, 'from biorun import const, utils, objects, ncbi\n'), ((148, 12, 148, 24), 'sys.exit', 'sys.exit', ({(148, 21, 148, 23): '(-1)'}, {}), '(-1)', False, 'import sys, os, glob, re, gzip, json\n'), ((204, 19, 204, 42), 'os.path.isfile', 'os.path.isfile', ({(204, 34, 204, 41): 'dest_gb'}, {}), '(dest_gb)', False, 'import sys, os, glob, re, gzip, json\n'), ((205, 16, 205, 43), 'os.symlink', 'os.symlink', ({(205, 27, 205, 33): 'src_gb', (205, 35, 205, 42): 'dest_gb'}, {}), '(src_gb, dest_gb)', False, 'import sys, os, glob, re, gzip, json\n')]