import io
import json
import xml.etree.ElementTree as ET
from typing import Any, AnyStr, Dict, List, Optional

from .types import *

xml_ns = "http://www.w3.org/XML/1998/namespace"
ns = "http://www.tei-c.org/ns/1.0"


def _string_to_tree(content: AnyStr) -> ET:
    if isinstance(content, str):
        return ET.parse(io.StringIO(content))
    elif isinstance(content, bytes):
        return ET.parse(io.BytesIO(content))
    if isinstance(content, io.StringIO) or isinstance(content, io.BytesIO):
        return ET.parse(content)
    elif isinstance(content, ET):
        return content
    else:
        raise TypeError(
            f"expected XML as string or bytes, got: {type(content)}")


def _parse_authors(elem: Optional[ET.Element],
                   ns: str = ns) -> List[GrobidAffiliation]:
    if not elem:
        return []
    names = []
    for author in elem.findall(f".//{{{ns}}}author"):
        pn = author.find(f"./{{{ns}}}persName")
        if not pn:
            continue
        given_name = pn.findtext(f"./{{{ns}}}forename") or None
        surname = pn.findtext(f"./{{{ns}}}surname") or None
        full_name = " ".join([t.strip() for t in pn.itertext()
                              if t.strip()]).strip()
        obj: Dict[str, Any] = dict(name=full_name)
        if given_name:
            obj["given_name"] = given_name
        if surname:
            obj["surname"] = surname
        ae = author.find(f"./{{{ns}}}affiliation")
        if ae:
            affiliation: Dict[str, Any] = dict()
            for on in ae.findall(f"./{{{ns}}}orgName"):
                on_type = on.get("type")
                if on_type:
                    affiliation[on_type] = on.text
            addr_e = ae.find(f"./{{{ns}}}address")
            if addr_e:
                address = dict()
                for t in addr_e.getchildren():
                    address[t.tag.split("}")[-1]] = t.text
                if address:
                    address['post_code'] = address.pop('postCode', None)
                    affiliation["address"] = GrobidAddress(**address)
                # previously:
                # affiliation['address'] = {
                #    'post_code': addr.findtext('./{%s}postCode' % ns) or None,
                #    'settlement': addr.findtext('./{%s}settlement' % ns) or None,
                #    'country': addr.findtext('./{%s}country' % ns) or None,
                # }
            obj["affiliation"] = GrobidAffiliation(**affiliation)
        names.append(GrobidAuthor(**obj))
    return names


def _parse_citation(elem: ET.Element, ns: str = ns) -> GrobidCitation:
    ref: Dict[str, Any] = dict()
    ref["id"] = elem.attrib.get("{http://www.w3.org/XML/1998/namespace}id")
    ref["unstructured"] = elem.findtext('.//{%s}note[@type="raw_reference"]' %
                                        ns)
    # Title stuff is messy in references...
    ref["title"] = elem.findtext(f".//{{{ns}}}analytic/{{{ns}}}title")
    other_title = elem.findtext(f".//{{{ns}}}monogr/{{{ns}}}title")
    if other_title:
        if ref["title"]:
            ref["journal"] = other_title
        else:
            ref["journal"] = None
            ref["title"] = other_title
    ref["authors"] = _parse_authors(elem, ns=ns)
    ref["publisher"] = elem.findtext(
        f".//{{{ns}}}publicationStmt/{{{ns}}}publisher")
    if not ref["publisher"]:
        ref["publisher"] = elem.findtext(
            f".//{{{ns}}}imprint/{{{ns}}}publisher")
    if ref["publisher"] == "":
        ref["publisher"] = None
    date = elem.find('.//{%s}date[@type="published"]' % ns)
    ref["date"] = (date is not None) and date.attrib.get("when")
    ref["volume"] = elem.findtext('.//{%s}biblScope[@unit="volume"]' % ns)
    ref["issue"] = elem.findtext('.//{%s}biblScope[@unit="issue"]' % ns)
    ref["doi"] = elem.findtext('.//{%s}idno[@type="DOI"]' % ns)
    ref["arxiv_id"] = elem.findtext('.//{%s}idno[@type="arXiv"]' % ns)
    if ref["arxiv_id"] and ref["arxiv_id"].startswith("arXiv:"):
        ref["arxiv_id"] = ref["arxiv_id"][6:]
    ref["pmcid"] = elem.findtext('.//{%s}idno[@type="PMCID"]' % ns)
    ref["pmid"] = elem.findtext('.//{%s}idno[@type="PMID"]' % ns)
    el = elem.find('.//{%s}biblScope[@unit="page"]' % ns)
    if el is not None:
        if el.attrib.get("from") and el.attrib.get("to"):
            ref["pages"] = "{}-{}".format(el.attrib["from"], el.attrib["to"])
        else:
            ref["pages"] = el.text
    el = elem.find(".//{%s}ptr[@target]" % ns)
    if el is not None:
        ref["url"] = el.attrib["target"]
        # Hand correction
        # TODO: move this elsewhere
        if ref["url"].endswith(".Lastaccessed"):
            ref["url"] = ref["url"].replace(".Lastaccessed", "")
        if ref["url"].startswith("<"):
            ref["url"] = ref["url"][1:]
        if ">" in ref["url"]:
            ref["url"] = ref["url"].split(">")[0]
    else:
        ref["url"] = None
    return GrobidCitation(**ref)


def _parse_journal(elem: ET.Element, ns: str = ns) -> GrobidJournal:
    journal = dict()
    journal["name"] = elem.findtext(f".//{{{ns}}}monogr/{{{ns}}}title")
    journal["publisher"] = elem.findtext(
        f".//{{{ns}}}publicationStmt/{{{ns}}}publisher")
    if journal["publisher"] == "":
        journal["publisher"] = None
    journal["issn"] = elem.findtext('.//{%s}idno[@type="ISSN"]' % ns)
    journal["eissn"] = elem.findtext('.//{%s}idno[@type="eISSN"]' % ns)
    journal["volume"] = elem.findtext('.//{%s}biblScope[@unit="volume"]' % ns)
    journal["issue"] = elem.findtext('.//{%s}biblScope[@unit="issue"]' % ns)
    journal["abbrev"] = None
    return GrobidJournal(**journal)


def _parse_header(elem: ET.Element, ns: str = ns) -> GrobidHeader:
    header = elem
    info = dict()
    info["title"] = header.findtext(f".//{{{ns}}}analytic/{{{ns}}}title")
    info["authors"] = _parse_authors(
        header.find(f".//{{{ns}}}sourceDesc/{{{ns}}}biblStruct"))
    info["journal"] = _parse_journal(header)
    date = header.find(f'.//{{{ns}}}date[@type="published"]')
    info["date"] = (date is not None) and date.attrib.get("when")
    info["doi"] = header.findtext(f'.//{{{ns}}}idno[@type="DOI"]')
    if info["doi"]:
        info["doi"] = info["doi"].lower()
    return GrobidHeader(**info)


def parse_document_xml(xml_text: AnyStr) -> GrobidDocument:
    """
    Use this function to parse TEI-XML of a full document or header processed
    by GROBID.
    
    Eg, the output of '/api/processFulltextDocument' or '/api/processHeader'
    """
    tree = _string_to_tree(xml_text)
    tei = tree.getroot()
    info = dict()

    header = tei.find(f".//{{{ns}}}teiHeader")
    if header is None:
        raise ValueError("XML does not look like TEI format")

    application_tag = header.findall(
        f".//{{{ns}}}appInfo/{{{ns}}}application")[0]

    doc = GrobidDocument(
        grobid_version=application_tag.attrib["version"].strip(),
        grobid_timestamp=application_tag.attrib["when"].strip(),
        header=_parse_header(header),
    )

    refs = []
    for (i, bs) in enumerate(
            tei.findall(f".//{{{ns}}}listBibl/{{{ns}}}biblStruct")):
        ref = _parse_citation(bs)
        ref.index = i
        refs.append(ref)
    doc.citations = refs

    text = tei.find(f".//{{{ns}}}text")
    # print(text.attrib)
    if text and text.attrib.get(f"{{{xml_ns}}}lang"):
        doc.language_code = text.attrib[f"{{{xml_ns}}}lang"]  # xml:lang

    el = tei.find(f".//{{{ns}}}profileDesc/{{{ns}}}abstract")
    doc.abstract = (el or None) and " ".join(el.itertext()).strip()
    el = tei.find(f".//{{{ns}}}text/{{{ns}}}body")
    doc.body = (el or None) and " ".join(el.itertext()).strip()
    el = tei.find(f'.//{{{ns}}}back/{{{ns}}}div[@type="acknowledgement"]')
    doc.acknowledgement = (el or None) and " ".join(el.itertext()).strip()
    el = tei.find(f'.//{{{ns}}}back/{{{ns}}}div[@type="annex"]')
    doc.annex = (el or None) and " ".join(el.itertext()).strip()

    return doc


def parse_citations_xml(xml_text: AnyStr) -> List[GrobidCitation]:
    """
    Use this function to parse TEI-XML of one or more references.
    
    Eg, the output of '/api/processReferences' or '/api/processCitation'.
    """
    # XXX: this replacement shouldn't be needed?
    xml_text = xml_text.replace('xmlns="http://www.tei-c.org/ns/1.0"', "")
    tree = _string_to_tree(xml_text)
    root = tree.getroot()

    if root.tag == 'biblStruct':
        ref = _parse_citation(root, ns='')
        ref.index = 0
        return [ref]

    refs = []
    for (i, bs) in enumerate(tree.findall(f".//biblStruct")):
        ref = _parse_citation(bs, ns='')
        ref.index = i
        refs.append(ref)
    return refs