import argparse
from collections import defaultdict
import csv
from dataclasses import dataclass, field
from enum import Enum, unique, auto
import os
import sys
import tempfile
import yaml
import zipfile

import gffutils
from google.protobuf import json_format
from ncbi.datasets.v1alpha1 import dataset_catalog_pb2
from ncbi.datasets.v1alpha1.reports import assembly_pb2
from ncbi.datasets.reports.report_reader import DatasetsReportReader


def retrieve_assembly_report(zip_in, catalog, assm_acc: str) -> assembly_pb2.AssemblyDataReport:
    report_files = get_catalog_files_for_assembly(catalog, dataset_catalog_pb2.File.FileType.DATA_REPORT, assm_acc)
    for path in report_files:
        yaml = zip_in.read(path)
        rpt_rdr = DatasetsReportReader()
        return rpt_rdr.assembly_report(yaml)


def retrieve_data_catalog(zip_in) -> dataset_catalog_pb2.Catalog:
    catalog_json = zip_in.read('ncbi_dataset/data/dataset_catalog.json')
    return json_format.Parse(catalog_json, dataset_catalog_pb2.Catalog())


def get_catalog_files_for_assembly(catalog: dataset_catalog_pb2.Catalog, desired_filetype: dataset_catalog_pb2.File.FileType, assm_acc: str):
    report_files = get_catalog_files(catalog, desired_filetype, assm_acc)
    filepaths = []
    for assm_acc, paths in report_files.items():
        filepaths.extend(paths)
    return filepaths


def get_catalog_files(catalog: dataset_catalog_pb2.Catalog, desired_filetype: dataset_catalog_pb2.File.FileType, assm_acc: str = None):
    files = defaultdict(list)
    for assm in catalog.assemblies:
        acc = assm.accession
        if assm_acc and assm_acc != acc:
            continue
        for f in assm.files:
            filepath = os.path.join('ncbi_dataset', 'data', f.file_path)
            if f.file_type == desired_filetype:
                files[acc].append(filepath)
    return files


def get_zip_file_for_acc(acc, path):
    fname = os.path.join(path, f'{acc}.zip')
    if os.path.isfile(fname):
        return fname
    return None


@dataclass
class Gene:
    id: str
    feat_type: str
    name: str
    chrom: str
    strand: str
    range_start: int
    range_stop: int
    protein_accession: str = ""

    def get_fields(self):
        return [self.feat_type, self.name, self.range_start, self.range_stop, self.protein_accession]

    def name_val(self):
        return self.protein_accession if self.protein_accession else self.name


def find_genes_by_loc(gff3_db, csvout, assm_acc, seq_acc, start, stop, extra_fields):
    found_genes = []
    feat_types = ('gene', 'pseudogene')
    for gene in gff3_db.region(seqid=seq_acc, start=start, end=stop, featuretype=feat_types, completely_within=False):
        gene_name = gene.attributes.get('Name', None)[0]
        prot_acc = ""
        if gene.attributes['gene_biotype'][0] == 'protein_coding':
            cds = list(gff3_db.children(gene, featuretype='CDS'))
            prot_acc = cds[0].attributes.get('protein_id', None)[0]
        geneobj = Gene(
            gene.id,
            gene.featuretype,
            gene_name,
            gene.chrom,
            gene.strand,
            gene.start,
            gene.stop,
            prot_acc,
        )
        csvout.writerow([assm_acc, seq_acc, start, stop, *extra_fields, *geneobj.get_fields()])
        found_genes.append(geneobj)
    return found_genes


class FindGenesByLoc:
    default_packages_dir = os.path.join('var', 'data', 'packages')

    def __init__(self):
        parser = argparse.ArgumentParser()
        parser.add_argument('--packages-dir', type=str, default=self.default_packages_dir,
                            help=f'root of input data directory [{self.default_packages_dir}]')
        parser.add_argument('--locs', type=str, help='file containing genomic locations')
        self.args = parser.parse_args()
        self.writer = csv.writer(sys.stdout, dialect='excel-tab')

    def read_data(self):
        for row in csv.reader(iter(sys.stdin.readline, ''), dialect='excel-tab'):
            yield row

    def run(self):
        for assm_acc, seq_acc, start, stop, *extra in self.read_data():
            self.find_all_for_location(assm_acc, seq_acc, start, stop, extra)

    def process_loc_for_gff(self, zin, gff_fname, assm_acc, seq_acc, start, stop, extra_fields):
        with tempfile.NamedTemporaryFile() as tmpfile:
            tmpfile.write(zin.read(gff_fname))
            db = gffutils.create_db(
                tmpfile.name,
                dbfn=':memory:',
                force=True,
                keep_order=True,
                merge_strategy='merge',
                sort_attribute_values=True
            )
            find_genes_by_loc(db, self.writer, assm_acc, seq_acc, start, stop, extra_fields)

    def find_all_for_location(self, assm_acc, seq_acc, start, stop, extra_fields):
        zip_file = get_zip_file_for_acc(assm_acc, self.args.packages_dir)
        try:
            with zipfile.ZipFile(zip_file, 'r') as zin:
                catalog = retrieve_data_catalog(zin)
                gff_files = get_catalog_files(catalog, dataset_catalog_pb2.File.FileType.GFF3)
                for assm_acc, gff_files in gff_files.items():
                    report = retrieve_assembly_report(zin, catalog, assm_acc)
                    for gff_fname in gff_files:
                        self.process_loc_for_gff(zin, gff_fname, assm_acc, seq_acc, start, stop, extra_fields)
        except zipfile.BadZipFile:
            print(f'{zip_file} is not a zip file')


if __name__ == '__main__':
    FindGenesByLoc().run()


[((41, 12, 41, 29), 'collections.defaultdict', 'defaultdict', ({(41, 24, 41, 28): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((54, 12, 54, 44), 'os.path.join', 'os.path.join', ({(54, 25, 54, 29): 'path', (54, 31, 54, 43): 'f"""{acc}.zip"""'}, {}), "(path, f'{acc}.zip')", False, 'import os\n'), ((55, 7, 55, 28), 'os.path.isfile', 'os.path.isfile', ({(55, 22, 55, 27): 'fname'}, {}), '(fname)', False, 'import os\n'), ((103, 27, 103, 66), 'os.path.join', 'os.path.join', ({(103, 40, 103, 45): '"""var"""', (103, 47, 103, 53): '"""data"""', (103, 55, 103, 65): '"""packages"""'}, {}), "('var', 'data', 'packages')", False, 'import os\n'), ((23, 18, 23, 40), 'ncbi.datasets.reports.report_reader.DatasetsReportReader', 'DatasetsReportReader', ({}, {}), '()', False, 'from ncbi.datasets.reports.report_reader import DatasetsReportReader\n'), ((29, 43, 29, 72), 'ncbi.datasets.v1alpha1.dataset_catalog_pb2.Catalog', 'dataset_catalog_pb2.Catalog', ({}, {}), '()', False, 'from ncbi.datasets.v1alpha1 import dataset_catalog_pb2\n'), ((106, 17, 106, 42), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((111, 22, 111, 65), 'csv.writer', 'csv.writer', (), '', False, 'import csv\n'), ((47, 23, 47, 72), 'os.path.join', 'os.path.join', ({(47, 36, 47, 50): '"""ncbi_dataset"""', (47, 52, 47, 58): '"""data"""', (47, 60, 47, 71): 'f.file_path'}, {}), "('ncbi_dataset', 'data', f.file_path)", False, 'import os\n'), ((122, 13, 122, 42), 'tempfile.NamedTemporaryFile', 'tempfile.NamedTemporaryFile', ({}, {}), '()', False, 'import tempfile\n'), ((124, 17, 131, 13), 'gffutils.create_db', 'gffutils.create_db', (), '', False, 'import gffutils\n'), ((137, 17, 137, 47), 'zipfile.ZipFile', 'zipfile.ZipFile', ({(137, 33, 137, 41): 'zip_file', (137, 43, 137, 46): '"""r"""'}, {}), "(zip_file, 'r')", False, 'import zipfile\n')]