import jinja2
import pathogenprofiler as pp
from .models import ProfileResult, Lineage
from typing import List

def lineagejson2text(lineages: List[Lineage],sep: str = "\t") -> str:
    textlines = []
    for l in lineages:
        row = [l.lineage,l.fraction,l.info["family"],l.info["spoligotype"],l.info["rd"]]
        textlines.append(sep.join(row))
    return "\n".join(textlines)

default_template = """
TBProfiler report
=================

The following report has been generated by TBProfiler.

Summary
-------
ID{{d['sep']}}{{d['id']}}
Date{{d['sep']}}{{d['timestamp']}}
Strain{{d['sep']}}{{d['strain']}}
Drug-resistance{{d['sep']}}{{d['drtype']}}
Median Depth{{d['sep']}}{{d['median_depth']}}

{% if 'notes' in d %}
Notes
--------------------
{{d['notes']}}
{% endif %}

Lineage report
--------------
{{d['lineage_report']}}
{% if 'spacers' in d %}
Spoligotype report
------------------
Binary{{d['sep']}}{{d['binary']}}
Octal{{d['sep']}}{{d['octal']}}
Family{{d['sep']}}{{d['family']}}
SIT{{d['sep']}}{{d['SIT']}}
{% endif %}
{% if 'close_samples' in d %}
Close samples report
--------------------
{{d['close_samples']}}
{% endif %}
{% if 'tree' in d %}
Neighbour joining tree
--------------------
{{d['tree']}}
{% endif %}
Resistance report
-----------------
{{d['dr_report']}}

Resistance variants report
-----------------
{{d['dr_var_report']}}

Other variants report
---------------------
{{d['other_var_report']}}

QC failed variants report
-------------------------

{{d['qc_fail_var_report']}}

Coverage report
---------------------
{{d['coverage_report']}}

Missing positions report
---------------------
{{d['missing_report']}}
{% if 'spacers' in d %}
Spoligotype spacers
-------------------
{{d['spacers']}}
{% endif %}
Analysis pipeline specifications
--------------------------------
Software version{{d['sep']}}{{d['version']}}
Database version{{d['sep']}}{{d['db_version']}}

{{d['software']}}

Citation
--------
Coll, F. et al. Rapid determination of anti-tuberculosis drug resistance from
whole-genome sequences. Genome Medicine 7, 51. 2015

Phelan, JE. et al. Integrating informatics tools and portable sequencing 
technology for rapid detection of resistance to anti-tuberculous drugs. 
Genome Medicine 11, 41. 2019
"""

def load_text(text_strings: str,template: str = None,file_template: str = None):
    
    if file_template:
        template = open(file_template).read()

    t =  jinja2.Template(template)
    return t.render(d=text_strings)

def stringify_annotations(annotation):
    annotations = []
    for ann in annotation:
        annotations.append("|".join([f'{key}={val}' for key,val in ann.items()]))
    return ";".join(annotations)

def write_text(
        result: ProfileResult,
        conf: dict,
        outfile: str,
        sep="\t",
        template_file=None
    ):

    template_string = default_template
    text_strings = {}
    text_strings["id"] = result.id
    text_strings["timestamp"] = result.timestamp

    summary_table = pp.get_dr_summary(result.dr_variants,conf)
    text_strings["notes"] = "\n".join(result.notes)
    text_strings["dr_report"] = pp.dict_list2text(summary_table,sep=sep)
    text_strings["dr_var_report"] = pp.object_list2text(result.dr_variants,mappings={"pos":"Genome Position","gene_id":"Locus Tag",'gene_name':'Gene name',"type":"Variant type","change":"Change","freq":"Estimated fraction","drugs.drug":"Drug","drugs.confidence":"Confidence","drugs.comment":"Comment"},sep=sep)
    text_strings["other_var_report"] = pp.object_list2text(result.other_variants,mappings={"pos":"Genome Position","gene_id":"Locus Tag",'gene_name':'Gene name',"type":"Variant type","change":"Change","freq":"Estimated fraction","annotation.drug":"Gene associated drug","annotation.confidence":"Confidence","annotation.comment":"Comment"},sep=sep)
    text_strings["qc_fail_var_report"] = pp.object_list2text(result.qc_fail_variants,mappings={"pos":"Genome Position","gene_id":"Locus Tag",'gene_name':'Gene name',"type":"Variant type","change":"Change","freq":"Estimated fraction","annotation.drug":"Gene associated drug","annotation.confidence":"Confidence","annotation.comment":"Comment"},sep=sep)
    text_strings["coverage_report"] = result.get_qc()
    text_strings['lineage_report'] = pp.object_list2text(result.lineage,mappings={"lineage":"Lineage","fraction":"Fraction","family":"Family","rd":"rd"},sep=sep)#result.lineage)
    text_strings['strain'] = result.sub_lineage
    text_strings['drtype'] = result.drtype
    text_strings['median_depth'] = result.qc.get_target_median_depth()
    text_strings['missing_report'] = result.get_missing_pos()
    text_strings['version'] = result.pipeline.software_version
    text_strings['db_version'] = result.pipeline.db_version['commit']
    text_strings['software'] = pp.dict_list2text(result.pipeline.software,sep=sep)

    if result.spoligotype:
        text_strings.update(result.spoligotype.model_dump())
        text_strings["spacers"] = pp.object_list2text(result.spoligotype.spacers,mappings={"name":"Spacer","seq":"Sequence","count":"Count"},sep=sep)

    if sep=="\t":
        text_strings["sep"] = ": "
    else:
        text_strings["sep"] = ","

    with open(outfile,"w") as O:
        O.write(load_text(text_strings,template_string,template_file))
    
