import xml.etree.ElementTree as ET
from Bio import SeqIO
import re
import gzip
import os
import argparse


def get_alias(identifier):
    parts = identifier.split("/")
    assert len(parts) == 5 and parts[0] == "SARS-CoV-2" and parts[1] == "human"
    return parts[3]


def main(
    results,
    ena_samples,
    output,
    program,
    platform="Illumina",
    assembly_type="COVID-19 outbreak",
):
    """Generate ENA assembly submission manifest files
    ...
    Parameters
    ----------
    results :   str
        Path to results directory generated by sarscov2-variation workflow.
    ena_samples :   str
        ENA samples xml.
    output  :   str
        Output directory.
    program :   str
        Sarscov2-variation workflow doi.
    platform    :   str, default='Illumina'
        Sequencing platform type.
    assembly_type   :   str, default='COVID-19 outbreak'
        Assembly type.
    """

    # Parse ENA xml
    tree = ET.parse(ena_samples)
    root = tree.getroot()

    study = {}
    for x in root.iter("SAMPLE"):
        SUBMITTER_ID = x.find("IDENTIFIERS/SUBMITTER_ID").text
        ENA_STUDY = x.findall('SAMPLE_LINKS/SAMPLE_LINK/XREF_LINK/[DB="ENA-STUDY"]')[0][
            1
        ].text
        ENA_RUN = x.findall('SAMPLE_LINKS/SAMPLE_LINK/XREF_LINK/[DB="ENA-RUN"]')[0][
            1
        ].text
        PRIMARY_ID = x.find("IDENTIFIERS/PRIMARY_ID").text
        with open(os.path.join(results, SUBMITTER_ID, "covstats.txt"), "r") as covstats:
            COVERAGE = round(float(covstats.readlines()[1].split()[1]))
        study.update(
            {
                SUBMITTER_ID: {
                    "STUDY": ENA_STUDY,
                    "SAMPLE": PRIMARY_ID,
                    "RUN_REF": ENA_RUN,
                    "ASSEMBLY_TYPE": assembly_type,
                    "PROGRAM": program,
                    "PLATFORM": platform,
                    "COVERAGE": COVERAGE,
                }
            }
        )

    os.makedirs(output, exist_ok=True)

    # Parse sequences
    seq_dict = SeqIO.index(
        os.path.join(results, "consensus_masked.fa"), "fasta", key_function=get_alias
    )
    fasta = {}
    for sample in list(study.keys()):
        record = seq_dict[sample]
        splits = record.id.split("/")
        record.id = f"{splits[3]}_{splits[2]}_{splits[4]}"
        record.description = ""
        record.seq = record.seq.strip("N")
        FASTA = os.path.join(output, f"{sample}.fasta.gz")
        CHR_LIST = os.path.join(output, f"{sample}.chrom.gz")
        with gzip.open(FASTA, "wt") as h:
            h.write(record.format("fasta"))
        with gzip.open(CHR_LIST, "wt") as h:
            h.write(f"{record.id}   1   Chromosome")
        fasta.update(
            {
                sample: {
                    "FASTA": os.path.basename(FASTA),
                    "NAME": record.id,
                    "CHROMOSOME_LIST": os.path.basename(CHR_LIST),
                }
            }
        )

    manifest = {}
    for key in set().union(study, fasta):
        if key in study:
            manifest.setdefault(key, {}).update(study[key])
        if key in fasta:
            manifest.setdefault(key, {}).update(fasta[key])

    for s, m in manifest.items():
        with open(os.path.join(output, f"{s}_manifest.txt"), "w") as h:
            for k, v in m.items():
                h.write(f"{k}   {v}\n")


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--results",
        "-R",
        dest="results_dir",
        required=True,
        help="Path to results directory generated by sarscov2-variation workflow.",
    )
    parser.add_argument(
        "--ena_samples",
        "-X",
        dest="xml",
        required=True,
        help="ENA samples xml.",
    )
    parser.add_argument(
        "--output",
        "-O",
        dest="output",
        required=True,
        help="Output directory. Will be created if missing.",
    )
    parser.add_argument(
        "--program",
        "-P",
        dest="program",
        required=True,
        help="Sarscov2-variation workflow doi. Look for doi on <https://github.com/avilab/sarscov2-variation>",
    )
    parser.add_argument(
        "--platform",
        "-S",
        dest="platform",
        required=False,
        default="Illumina",
        help="Sequencing platform type.",
    )
    parser.add_argument(
        "--assembly_type",
        "-A",
        dest="assembly_type",
        required=False,
        default="COVID-19 outbreak",
        help="Assembly type.",
    )
    args = parser.parse_args()

    main(
        results=args.results_dir,
        ena_samples=args.xml,
        output=args.output,
        program=args.program,
        platform=args.platform,
        assembly_type=args.assembly_type,
    )
