#!/usr/bin/env python
import sys
import argparse
import logging
import uuid
from CPT_GFFParser import gffParse, gffWrite
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.SeqFeature import FeatureLocation, CompoundLocation
from gff3 import feature_lambda, feature_test_type, get_id

logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)


def main(fasta, gff3, feature_filter=None, nodesc=False):
    if feature_filter == "nice_cds":
        from gff2gb import gff3_to_genbank as cpt_Gff2Gbk

        for rec in cpt_Gff2Gbk(gff3, fasta, 11):
            seenList = {}
            if rec.seq[0] == "?":
                sys.stderr.write("Error: No Fasta ID matches GFF ID '" + rec.id + "'\n")
                exit(1)
            for feat in sorted(rec.features, key=lambda x: x.location.start):
                if feat.type != "CDS":
                    continue

                ind = 0
                if (
                    str(
                        feat.qualifiers.get("locus_tag", get_id(feat)).replace(" ", "-")
                    )
                    in seenList.keys()
                ):
                    seenList[
                        str(
                            feat.qualifiers.get("locus_tag", get_id(feat)).replace(
                                " ", "-"
                            )
                        )
                    ] += 1
                    ind = seenList[
                        str(
                            feat.qualifiers.get("locus_tag", get_id(feat)).replace(
                                " ", "-"
                            )
                        )
                    ]
                else:
                    seenList[
                        str(
                            feat.qualifiers.get("locus_tag", get_id(feat)).replace(
                                " ", "-"
                            )
                        )
                    ] = 1
                append = ""
                if ind != 0:
                    append = "_" + str(ind)

                if nodesc:
                    description = ""
                else:
                    feat.qualifiers["ID"] = [feat._ID]
                    product = feat.qualifiers.get("product", "")
                    description = (
                        "{1} [Location={0.location};ID={0.qualifiers[ID][0]}]".format(
                            feat, product
                        )
                    )
                yield [
                    SeqRecord(
                        feat.extract(rec).seq,
                        id=str(
                            feat.qualifiers.get("locus_tag", get_id(feat)).replace(
                                " ", "-"
                            )
                        )
                        + append,
                        description=description,
                    )
                ]

    elif feature_filter == "unique_cds":
        seq_dict = SeqIO.to_dict(SeqIO.parse(fasta, "fasta"))
        seen_ids = {}

        for rec in gffParse(gff3, base_dict=seq_dict):
            noMatch = True
            if "Alias" in rec.features[0].qualifiers.keys():
                lColumn = rec.features[0].qualifiers["Alias"][0]
            else:
                lColumn = ""
            for x in seq_dict:
                if x == rec.id or x == lColumn:
                    noMatch = False
            if noMatch:
                sys.stderr.write("Error: No Fasta ID matches GFF ID '" + rec.id + "'\n")
                exit(1)
            newfeats = []
            for feat in sorted(
                feature_lambda(
                    rec.features, feature_test_type, {"type": "CDS"}, subfeatures=False
                ),
                key=lambda f: f.location.start,
            ):
                nid = rec.id + "____" + feat.id
                if nid in seen_ids:
                    nid = nid + "__" + uuid.uuid4().hex
                feat.qualifiers["ID"] = [nid]
                newfeats.append(feat)
                seen_ids[nid] = True

                if nodesc:
                    description = ""
                else:
                    if feat.strand == -1:
                        important_data = {
                            "Location": FeatureLocation(
                                feat.location.start + 1,
                                feat.location.end - feat.phase,
                                feat.strand,
                            )
                        }
                    else:
                        important_data = {
                            "Location": FeatureLocation(
                                feat.location.start + 1 + feat.phase,
                                feat.location.end,
                                feat.strand,
                            )
                        }
                    if "Name" in feat.qualifiers:
                        important_data["Name"] = feat.qualifiers.get("Name", [""])[0]

                    description = "[{}]".format(
                        ";".join(
                            [
                                "{key}={value}".format(key=k, value=v)
                                for (k, v) in important_data.items()
                            ]
                        )
                    )
                # if feat.id == "CPT_Privateer_006.p01":
                # print(feat)
                # exit()

                if isinstance(feat.location, CompoundLocation):
                    finSeq = ""
                    if feat.strand == -1:
                        for x in feat.location.parts:
                            finSeq += str(
                                (
                                    rec.seq[
                                        feat.location.start : feat.location.end
                                        - feat.phase
                                    ]
                                ).reverse_complement()
                            )
                    else:
                        for x in feat.location.parts:
                            finSeq += str(
                                rec.seq[
                                    feat.location.start + feat.phase : feat.location.end
                                ]
                            )
                    yield [
                        SeqRecord(
                            finSeq,
                            id=nid.replace(" ", "-"),
                            description=description,
                        )
                    ]
                elif feat.strand == -1:
                    yield [
                        SeqRecord(
                            (
                                rec.seq[
                                    feat.location.start : feat.location.end - feat.phase
                                ]
                            ).reverse_complement(),
                            id=nid.replace(" ", "-"),
                            description=description,
                        )
                    ]
                else:
                    yield [
                        SeqRecord(
                            # feat.extract(rec).seq,
                            rec.seq[
                                feat.location.start + feat.phase : feat.location.end
                            ],
                            id=nid.replace(" ", "-"),
                            description=description,
                        )
                    ]
            rec.features = newfeats
            rec.annotations = {}
            # gffWrite([rec], sys.stdout)
    else:
        seq_dict = SeqIO.to_dict(SeqIO.parse(fasta, "fasta"))

        for rec in gffParse(gff3, base_dict=seq_dict):
            noMatch = True
            if "Alias" in rec.features[0].qualifiers.keys():
                lColumn = rec.features[0].qualifiers["Alias"][0]
            else:
                lColumn = ""
            for x in seq_dict:
                if x == rec.id or x == lColumn:
                    noMatch = False
            if noMatch:
                sys.stderr.write("Error: No Fasta ID matches GFF ID '" + rec.id + "'\n")
                exit(1)
            for feat in sorted(
                feature_lambda(
                    rec.features,
                    feature_test_type,
                    {"type": feature_filter},
                    subfeatures=True,
                ),
                key=lambda f: f.location.start,
            ):
                id = feat.id
                if len(id) == 0:
                    id = get_id(feat)

                if nodesc:
                    description = ""
                else:
                    if feat.strand == -1:
                        important_data = {
                            "Location": FeatureLocation(
                                feat.location.start + 1,
                                feat.location.end - feat.phase,
                                feat.strand,
                            )
                        }
                    else:
                        important_data = {
                            "Location": FeatureLocation(
                                feat.location.start + 1 + feat.phase,
                                feat.location.end,
                                feat.strand,
                            )
                        }
                    if "Name" in feat.qualifiers:
                        important_data["Name"] = feat.qualifiers.get("Name", [""])[0]

                    description = "[{}]".format(
                        ";".join(
                            [
                                "{key}={value}".format(key=k, value=v)
                                for (k, v) in important_data.items()
                            ]
                        )
                    )

                if isinstance(feat.location, CompoundLocation):
                    finSeq = ""
                    if feat.strand == -1:
                        for x in feat.location.parts:
                            finSeq += str(
                                (
                                    rec.seq[x.start : x.end - feat.phase]
                                ).reverse_complement()
                            )
                    else:
                        for x in feat.location.parts:
                            finSeq += str(rec.seq[x.start + feat.phase : x.end])
                    yield [
                        SeqRecord(
                            Seq(finSeq),
                            id=id.replace(" ", "-"),
                            description=description,
                        )
                    ]

                else:

                    if feat.strand == -1:
                        yield [
                            SeqRecord(
                                seq=Seq(
                                    str(
                                        rec.seq[
                                            feat.location.start : feat.location.end
                                            - feat.phase
                                        ]
                                    )
                                ).reverse_complement(),
                                id=id.replace(" ", "-"),
                                description=description,
                            )
                        ]
                    else:
                        yield [
                            SeqRecord(
                                # feat.extract(rec).seq,
                                seq=Seq(
                                    str(
                                        rec.seq[
                                            feat.location.start
                                            + feat.phase : feat.location.end
                                        ]
                                    )
                                ),
                                id=id.replace(" ", "-"),
                                description=description,
                            )
                        ]


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Export corresponding sequence in genome from GFF3", epilog=""
    )
    parser.add_argument("fasta", type=argparse.FileType("r"), help="Fasta Genome")
    parser.add_argument("gff3", type=argparse.FileType("r"), help="GFF3 File")
    parser.add_argument(
        "--feature_filter", default=None, help="Filter for specific feature types"
    )
    parser.add_argument(
        "--nodesc", action="store_true", help="Strip description field off"
    )
    args = parser.parse_args()
    for seq in main(**vars(args)):
        # if isinstance(seq, list):
        #  for x in seq:
        #    print(type(x.seq))
        #    SeqIO.write(x, sys.stdout, "fasta")
        # else:
        SeqIO.write(seq, sys.stdout, "fasta")