"""
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()
[((6836, 6866), 'biorun.libs.placlib.pos', 'plac.pos', (['"""data"""', '"""data names"""'], {}), "('data', 'data names')\n", (6844, 6866), True, 'import biorun.libs.placlib as plac\n'), ((6868, 6916), 'biorun.libs.placlib.flg', 'plac.flg', (['"""fetch"""', '"""download data as accessions"""'], {}), "('fetch', 'download data as accessions')\n", (6876, 6916), True, 'import biorun.libs.placlib as plac\n'), ((6918, 6963), 'biorun.libs.placlib.flg', 'plac.flg', (['"""update"""', '"""updates data in storage"""'], {}), "('update', 'updates data in storage')\n", (6926, 6963), True, 'import biorun.libs.placlib as plac\n'), ((6965, 7002), 'biorun.libs.placlib.opt', 'plac.opt', (['"""rename"""', '"""rename the data"""'], {}), "('rename', 'rename the data')\n", (6973, 7002), True, 'import biorun.libs.placlib as plac\n'), ((7004, 7056), 'biorun.libs.placlib.opt', 'plac.opt', (['"""seqid"""', '"""set the sequence id of the data"""'], {}), "('seqid', 'set the sequence id of the data')\n", (7012, 7056), True, 'import biorun.libs.placlib as plac\n'), ((7058, 7105), 'biorun.libs.placlib.flg', 'plac.flg', (['"""protein"""', '"""use the protein database"""'], {}), "('protein', 'use the protein database')\n", (7066, 7105), True, 'import biorun.libs.placlib as plac\n'), ((7107, 7146), 'biorun.libs.placlib.flg', 'plac.flg', (['"""build"""', '"""build the database"""'], {}), "('build', 'build the database')\n", (7115, 7146), True, 'import biorun.libs.placlib as plac\n'), ((7148, 7183), 'biorun.libs.placlib.flg', 'plac.flg', (['"""verbose"""', '"""verbose mode"""'], {}), "('verbose', 'verbose mode')\n", (7156, 7183), True, 'import biorun.libs.placlib as plac\n'), ((7807, 7868), 'biorun.libs.placlib.opt', 'plac.opt', (['"""delete"""', '"""deletes foo from storage"""'], {'metavar': '"""foo"""'}), "('delete', 'deletes foo from storage', metavar='foo')\n", (7815, 7868), True, 'import biorun.libs.placlib as plac\n'), ((7870, 7905), 'biorun.libs.placlib.flg', 'plac.flg', (['"""verbose"""', '"""verbose mode"""'], {}), "('verbose', 'verbose mode')\n", (7878, 7905), True, 'import biorun.libs.placlib as plac\n'), ((749, 783), 'os.path.join', 'os.path.join', (['utils.DATADIR', 'fname'], {}), '(utils.DATADIR, fname)\n', (761, 783), False, 'import sys, os, glob, re, gzip, json\n'), ((1219, 1239), 'biorun.utils.gz_read', 'utils.gz_read', (['fname'], {}), '(fname)\n', (1232, 1239), False, 'from biorun import const, utils, objects, ncbi\n'), ((1251, 1264), 'json.load', 'json.load', (['fp'], {}), '(fp)\n', (1260, 1264), False, 'import sys, os, glob, re, gzip, json\n'), ((1396, 1417), 'biorun.utils.gz_write', 'utils.gz_write', (['fname'], {}), '(fname)\n', (1410, 1417), False, 'from biorun import const, utils, objects, ncbi\n'), ((1422, 1441), 'json.dump', 'json.dump', (['data', 'fp'], {}), '(data, fp)\n', (1431, 1441), False, 'import sys, os, glob, re, gzip, json\n'), ((1619, 1644), 'os.path.isfile', 'os.path.isfile', (['json_name'], {}), '(json_name)\n', (1633, 1644), False, 'import sys, os, glob, re, gzip, json\n'), ((2039, 2054), 'biorun.ncbi.build_db', 'ncbi.build_db', ([], {}), '()\n', (2052, 2054), False, 'from biorun import const, utils, objects, ncbi\n'), ((2089, 2104), 'biorun.ncbi.get_data', 'ncbi.get_data', ([], {}), '()\n', (2102, 2104), False, 'from biorun import const, utils, objects, ncbi\n'), ((3667, 3687), 'os.path.isfile', 'os.path.isfile', (['name'], {}), '(name)\n', (3681, 3687), False, 'import sys, os, glob, re, gzip, json\n'), ((4161, 4186), 'os.path.isfile', 'os.path.isfile', (['json_name'], {}), '(json_name)\n', (4175, 4186), False, 'import sys, os, glob, re, gzip, json\n'), ((4356, 4380), 'os.path.isfile', 'os.path.isfile', (['gbk_name'], {}), '(gbk_name)\n', (4370, 4380), False, 'import sys, os, glob, re, gzip, json\n'), ((5906, 5924), 'glob.glob', 'glob.glob', (['pattern'], {}), '(pattern)\n', (5915, 5924), False, 'import sys, os, glob, re, gzip, json\n'), ((6000, 6054), 're.compile', 're.compile', (['"""(definition\\\\":\\\\s*)(?P<value>\\\\".+?\\\\")"""'], {}), '(\'(definition\\\\":\\\\s*)(?P<value>\\\\".+?\\\\")\')\n', (6010, 6054), False, 'import sys, os, glob, re, gzip, json\n'), ((7465, 7488), 'biorun.models.jsonrec.reset_counter', 'jsonrec.reset_counter', ([], {}), '()\n', (7486, 7488), False, 'from biorun.models import jsonrec\n'), ((7562, 7649), 'biorun.objects.Param', 'objects.Param', ([], {'seqid': 'seqid', 'rename': 'rename', 'start': '(1)', 'protein': 'protein', 'update': 'update'}), '(seqid=seqid, rename=rename, start=1, protein=protein, update=\n    update)\n', (7575, 7649), False, 'from biorun import const, utils, objects, ncbi\n'), ((8321, 8344), 'biorun.models.jsonrec.reset_counter', 'jsonrec.reset_counter', ([], {}), '()\n', (8342, 8344), False, 'from biorun.models import jsonrec\n'), ((560, 572), 'sys.exit', 'sys.exit', (['(-1)'], {}), '(-1)\n', (568, 572), False, 'import sys, os, glob, re, gzip, json\n'), ((957, 978), 'os.path.isfile', 'os.path.isfile', (['fname'], {}), '(fname)\n', (971, 978), False, 'import sys, os, glob, re, gzip, json\n'), ((1764, 1789), 'biorun.utils.gz_write', 'utils.gz_write', (['json_name'], {}), '(json_name)\n', (1778, 1789), False, 'from biorun import const, utils, objects, ncbi\n'), ((1798, 1817), 'json.dump', 'json.dump', (['data', 'fp'], {}), '(data, fp)\n', (1807, 1817), False, 'import sys, os, glob, re, gzip, json\n'), ((2910, 2963), 'biorun.models.jsonrec.parse_file', 'jsonrec.parse_file', ([], {'fname': 'gbk_name', 'seqid': 'param.seqid'}), '(fname=gbk_name, seqid=param.seqid)\n', (2928, 2963), False, 'from biorun.models import jsonrec\n'), ((3163, 3188), 'os.path.isfile', 'os.path.isfile', (['param.acc'], {}), '(param.acc)\n', (3177, 3188), False, 'import sys, os, glob, re, gzip, json\n'), ((4438, 4485), 'biorun.models.jsonrec.parse_file', 'jsonrec.parse_file', ([], {'fname': 'gbk_name', 'seqid': 'seqid'}), '(fname=gbk_name, seqid=seqid)\n', (4456, 4485), False, 'from biorun.models import jsonrec\n'), ((4639, 4678), 'biorun.models.jsonrec.make_jsonrec', 'jsonrec.make_jsonrec', (['name'], {'seqid': 'seqid'}), '(name, seqid=seqid)\n', (4659, 4678), False, 'from biorun.models import jsonrec\n'), ((4759, 4797), 'biorun.utils.error', 'utils.error', (['f"""data not found: {name}"""'], {}), "(f'data not found: {name}')\n", (4770, 4797), False, 'from biorun import const, utils, objects, ncbi\n'), ((5257, 5281), 'os.path.isfile', 'os.path.isfile', (['src_json'], {}), '(src_json)\n', (5271, 5281), False, 'import sys, os, glob, re, gzip, json\n'), ((5549, 5571), 'os.path.isfile', 'os.path.isfile', (['src_gb'], {}), '(src_gb)\n', (5563, 5571), False, 'import sys, os, glob, re, gzip, json\n'), ((5850, 5890), 'os.path.join', 'os.path.join', (['utils.DATADIR', '"""*.json.gz"""'], {}), "(utils.DATADIR, '*.json.gz')\n", (5862, 5890), False, 'import sys, os, glob, re, gzip, json\n'), ((6173, 6192), 'os.path.split', 'os.path.split', (['path'], {}), '(path)\n', (6186, 6192), False, 'import sys, os, glob, re, gzip, json\n'), ((992, 1008), 'os.remove', 'os.remove', (['fname'], {}), '(fname)\n', (1001, 1008), False, 'import sys, os, glob, re, gzip, json\n'), ((2642, 2737), 'biorun.ncbi.genome', 'ncbi.genome', ([], {'name': 'name', 'fname': 'gbk_name', 'update': 'param.update', 'genbank': 'genbank', 'refseq': 'refseq'}), '(name=name, fname=gbk_name, update=param.update, genbank=genbank,\n    refseq=refseq)\n', (2653, 2737), False, 'from biorun import const, utils, objects, ncbi\n'), ((2812, 2858), 'biorun.ncbi.genbank_save', 'ncbi.genbank_save', (['name'], {'db': 'db', 'fname': 'gbk_name'}), '(name, db=db, fname=gbk_name)\n', (2829, 2858), False, 'from biorun import const, utils, objects, ncbi\n'), ((3211, 3235), 'biorun.utils.gz_read', 'utils.gz_read', (['param.acc'], {}), '(param.acc)\n', (3224, 3235), False, 'from biorun import const, utils, objects, ncbi\n'), ((3249, 3272), 'os.path.isfile', 'os.path.isfile', (['altname'], {}), '(altname)\n', (3263, 3272), False, 'import sys, os, glob, re, gzip, json\n'), ((3721, 3758), 'biorun.models.jsonrec.parse_file', 'jsonrec.parse_file', (['name'], {'seqid': 'seqid'}), '(name, seqid=seqid)\n', (3739, 3758), False, 'from biorun.models import jsonrec\n'), ((5350, 5380), 'os.rename', 'os.rename', (['src_json', 'dest_json'], {}), '(src_json, dest_json)\n', (5359, 5380), False, 'import sys, os, glob, re, gzip, json\n'), ((6128, 6149), 'os.path.getsize', 'os.path.getsize', (['path'], {}), '(path)\n', (6143, 6149), False, 'import sys, os, glob, re, gzip, json\n'), ((6294, 6315), 'gzip.open', 'gzip.open', (['path', '"""rt"""'], {}), "(path, 'rt')\n", (6303, 6315), False, 'import sys, os, glob, re, gzip, json\n'), ((3295, 3317), 'biorun.utils.gz_read', 'utils.gz_read', (['altname'], {}), '(altname)\n', (3308, 3317), False, 'from biorun import const, utils, objects, ncbi\n'), ((3368, 3411), 'biorun.utils.error', 'utils.error', (['f"""data not found: {param.acc}"""'], {}), "(f'data not found: {param.acc}')\n", (3379, 3411), False, 'from biorun import const, utils, objects, ncbi\n'), ((3875, 3887), 'sys.exit', 'sys.exit', (['(-1)'], {}), '(-1)\n', (3883, 3887), False, 'import sys, os, glob, re, gzip, json\n'), ((5592, 5615), 'os.path.isfile', 'os.path.isfile', (['dest_gb'], {}), '(dest_gb)\n', (5606, 5615), False, 'import sys, os, glob, re, gzip, json\n'), ((5633, 5660), 'os.symlink', 'os.symlink', (['src_gb', 'dest_gb'], {}), '(src_gb, dest_gb)\n', (5643, 5660), False, 'import sys, os, glob, re, gzip, json\n')]