"""Biosample parser

Implemented as an iterator


>>> import omicidx.biosample as b
>>> for bios in b.BioSampleParser(gzip.open('biosample_set.xml.gz', 'rb')):
>>>     print(bios.json())
>>>     print(bios.dict())
"""

import os.path
import xml.etree.ElementTree as ET
from xml.etree.ElementTree import Element
import gzip
import json
import typing
import logging
import click
import subprocess
import re
import pydantic
import datetime


class IdRecs(pydantic.BaseModel):
    db: typing.Optional[str]
    label: typing.Optional[str]
    id: typing.Optional[str]


class AttrRecs(pydantic.BaseModel):
    attribute_name: typing.Optional[str]
    display_name: typing.Optional[str]
    harmonized_name: typing.Optional[str]
    value: typing.Optional[str]
    unit: typing.Optional[str]


class BioSample(pydantic.BaseModel):
    accession: str
    id: str
    title: typing.Optional[str]
    description: typing.Optional[str]
    taxonomy_name: str
    taxon_id: int
    attribute_recs: list[AttrRecs]
    attributes: list[str]
    model: typing.Optional[str]
    id_recs: list[IdRecs]
    ids: list[str]
    sra_sample: typing.Optional[str]
    dbgap: typing.Optional[str]
    gsm: typing.Optional[str]
    publication_date: datetime.datetime
    last_update: datetime.datetime
    submission_date: datetime.datetime
    access: typing.Optional[str]


class SimplePublication(pydantic.BaseModel):
    db: typing.Optional[str]
    id: typing.Optional[str]
    pubdate: typing.Optional[datetime.datetime]


class ExternalLink(pydantic.BaseModel):
    url: typing.Optional[str]
    label: typing.Optional[str]
    category: typing.Optional[str]


class BioProject(pydantic.BaseModel):
    data_types: list[str] = []
    description: typing.Optional[str]
    accession: str
    name: typing.Optional[str]
    publications: list[SimplePublication] = []
    title: typing.Optional[str]
    external_links: list[ExternalLink] = []


class BioSampleParser(object):
    """Parse a BioSample xml file.

    This is a generator that yields BioSample objects.
    The BioSample objects are pydantic models.
    Access to fields is via dot notation.
    Access to the underlying dict is via the dict() method.
    JSON serialization is via the json() method."""

    def __init__(self, fh: typing.IO):
        """Initialize a new parser

        Args:
            fh (typing.IO): An open file-like object containing Biosample XML records.
        """
        self.fhandle = fh
        self.context = ET.iterparse(self.fhandle, events=("end",))
        event, self.root = next(self.context)

    def __iter__(self):
        return self

    def __next__(self):
        for event, elem in self.context:
            if elem.tag == "BioSample":
                bios = dict()
                bios["is_reference"] = None
                for k, v in elem.items():
                    bios[k] = v
                bios["id_recs"] = []
                bios["ids"] = []
                bios["sra_sample"] = None
                bios["dbgap"] = None
                bios["gsm"] = None
                for id in elem.iterfind(".//Id"):
                    idrec = {
                        "db": id.get("db"),
                        "label": id.get("db_label"),
                        "id": id.text,
                    }
                    bios["ids"].append(idrec["id"])
                    bios["id_recs"].append(idrec)
                    # add xref fields for SRA, dbGaP, and GEO

                    if id.get("db") == "SRA":
                        bios["sra_sample"] = id.text
                    if id.get("db") == "dbGaP":
                        bios["dbgap"] = id.text
                    if id.get("db") == "GEO":
                        bios["gsm"] = id.text
                bios["title"] = elem.findtext(".//Description/Title")
                bios["description"] = elem.findtext(".//Description/Comment/Paragraph")
                organism = elem.find(".//Organism")
                bios["taxonomy_name"] = organism.get("taxonomy_name")
                bios["taxon_id"] = int(organism.get("taxonomy_id"))
                bios["attribute_recs"] = []
                bios["attributes"] = []
                for attribute in elem.findall(".//Attribute"):
                    attrec = attribute.attrib
                    attrec["value"] = attribute.text
                    bios["attribute_recs"].append(attrec)
                    try:
                        bios["attributes"].append(attrec["harmonized_name"])
                    except:
                        bios["attributes"].append(attrec["attribute_name"])

                bios["model"] = elem.findtext(".//Model")
                # print(json.dumps(bios))
                # res = es.index(index="bioes", doc_type='biosample', id=bios['id'], body=bios)
                elem.clear()
                return BioSample(**bios)
        raise StopIteration


def parse_bioproject_xml_element(element: Element) -> dict:
    """Parse a BioProject xml element

    Args:
        element (Element): An lxml.etree.Element

    Returns:
        dict: A BioProject dict.
    """
    projtop = element.find("./Project")
    d2 = {}
    d2["title"] = projtop.findtext("./Project/ProjectDescr/Title")
    d2["description"] = projtop.findtext("./Project/ProjectDescr/Description")
    d2["name"] = projtop.findtext("./Project/ProjectDescr/Name")
    archive_id = projtop.find("./Project/ProjectID/ArchiveID")
    d2["accession"] = archive_id.attrib["accession"]
    pubs = []
    for pub in projtop.findall(".//Publication"):
        pubs.append(
            {
                "pubdate": pub.get("date", None),
                "id": pub.get("id", None),
                "db": re.sub("^e", "", pub.findtext("./DbType").replace("^e", "", 1)),
            }
        )
    d2["publications"] = pubs
    ext_links = []
    for link in projtop.findall(".//ExternalLink"):
        ext_links.append(
            {
                "category": link.get("category", None),
                "label": link.get("label", None),
                "url": link.findtext("./URL"),
            }
        )
    data_types = []
    for datatype in projtop.findall(".//ProjectDataTypeSet"):
        data_types.append(datatype.findtext("./DataType"))
    d2["data_types"] = data_types
    d2["external_links"] = ext_links
    return d2


class BioProjectParser(typing.Iterable):
    """Parse a BioProject xml file.

    This is a generator that yields BioProject objects.
    The BioProject object is a pydantic model. Access the fields using dot notation.
    Access the raw data using the .dict() method.
    Access json using the .json() method."""

    def __init__(self, fh: typing.IO):
        """Initialize a BioProjectParser

        Args:
            fh (typing.IO): An open file-like object containing BioProject records.
        """
        self.fhandle = fh
        self.context = ET.iterparse(self.fhandle, events=("end",))
        event, self.root = next(self.context)

    def __iter__(self):
        return self

    def __next__(self):
        for event, elem in self.context:
            if elem.tag == "Package":
                results = parse_bioproject_xml_element(elem)
                elem.clear()
                return BioProject(**results)

        raise StopIteration
