#!/usr/bin/env python
import sys, os, urllib.request, urllib.error, argparse, re, html, textwrap, time, requests
from xml.etree import ElementTree as ET

def eutilsFetch(url):
    """ Fetch an E-utilities URL and return the parsed XML.  NCBI allows three requests a
    second and answers 429 above that, and we make two requests per citation, so back off
    and retry a few times instead of failing a long list of ids halfway through. """
    sys.stderr.write("Accessing %s\n" % url)
    for tryCount in range(5):
        try:
            return ET.XML(urllib.request.urlopen(url).read())
        except urllib.error.HTTPError as e:
            if e.code != 429 or tryCount == 4:
                # never print this on stdout, it would end up in the middle of the references
                sys.exit("error: cannot fetch %s: %s" % (url, e))
            time.sleep(2 ** tryCount)

def parsePubmed(doc, id):
    infoDict = dict()
    infoDict['url'] = "https://www.ncbi.nlm.nih.gov/pubmed/%s" % id
    attribList = ['PubDate', 'Source', 'Title', 'Volume', 'Issue', 'Pages', 'SO', 'CollectiveName']
    for element in doc:
        if element.tag != "DocSum":
            continue
        items = element.findall("Item")
        for i in items:
            if i.attrib['Name'] == 'AuthorList':
                infoDict['Authors'] = list()
                for j in i:
                    infoDict['Authors'].append(j.text)
                continue
            if i.attrib['Name'] == "ArticleIds":
                for j in i:
                    if (j.attrib['Name'] == 'pubmed'):
                        infoDict[j.attrib['Name']] = j.text
                    if (j.attrib['Name'] == 'pmc'):
                        infoDict[j.attrib['Name']] = j.text
                    if (j.attrib['Name'] == 'doi'):
                        infoDict[j.attrib['Name']] = j.text
                continue
            if i.attrib['Name'] in attribList:
                infoDict[i.attrib['Name']] = i.text

    return infoDict

def parsePmc(doc, id):
    for node in doc.iter("ArticleId"):
        foundPubMedId = 0
        for child in node:
            if child.tag == "IdType" and child.text == "pmid":
                foundPubMedId = 1
            if foundPubMedId == 1 and child.tag == "Value":
                return parseInit(child.text) 
            
    sys.stderr.write("Unable to find pubmed id for pubmed central id: %s\n" % id)
    sys.exit()

def parseInit(id):
    urlbase = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi?"
    db = "Pubmed"
    url = urlbase + "db=%s" % db + "&id=%s" % id
    if re.match("^PMC", id):
        db = "PMC"
        id = re.sub("PMC", "", id)
        url = urlbase + "db=%s" % db + "&id=%s&version=2.0" % id
    doc = eutilsFetch(url)
    if db == "Pubmed":
        infoDict = parsePubmed(doc, id)
    elif db == "PMC":
        infoDict = parsePmc(doc, id)

    return infoDict

def resolveDoi(url):
    """ NCBI hands out doi.org links for many publishers.  Follow the redirect, so that the
    citation points straight at the article page, the way it did when we still scraped the
    links out of the PubMed page.  Keep the doi.org link if the publisher does not answer. """
    if not re.match(r"https?://(dx\.)?doi\.org/", url):
        return url
    try:
        return requests.head(url, allow_redirects=True, timeout=60).url
    except requests.RequestException:
        return url

def fetchFullTextUrl(id):
    """ Return the publisher's URL for a PubMed ID, or None if NCBI lists no full text
    provider for it.  Uses the E-utilities elink/prlinks API.  We used to scrape the "full
    text links" box out of the PubMed web page, but NCBI now answers non-browser clients
    with a Javascript cookie challenge, so that page no longer contains any links. """
    url = ("https://eutils.ncbi.nlm.nih.gov/entrez/eutils/elink.fcgi?"
           "dbfrom=pubmed&cmd=prlinks&retmode=xml&id=%s" % id)
    doc = eutilsFetch(url)

    # elink returns one ObjUrl per LinkOut provider, publisher first, in the same order as
    # the "full text links" box on the PubMed page, from which we take the first one
    for objUrl in doc.iter("ObjUrl"):
        link = objUrl.find("Url")
        if link is not None and link.text:
            return resolveDoi(link.text.strip())

    return None

def htmlEscape(str):
    return html.escape(str)

def makeHtml(infoDict, plain, verbose, doi):
    authors = list()
    authcount = 0
    etal = 0
    if 'CollectiveName' in infoDict:
        authors.append(infoDict['CollectiveName'])
        authcount = 1
    for i in infoDict['Authors']:
        if authcount == 10 and not verbose:
            etal = 1
            break
        authors.append(i)
        authcount = authcount + 1
    sep = ", "
    authStr = sep.join(authors)
    if (not plain):
        authStr = htmlEscape(authStr)
    if etal and not plain:
        authStr = authStr + (" <em>et al</em>")
    if etal and plain:
       authStr = authStr + (" et al")
    authStr = authStr + (".")
    title = re.sub(r"\.$", "", infoDict['Title'])
    if 'Source' in infoDict:
        journal = infoDict['Source']
    elif 'Journal' in infoDict:
        journal = infoDict['Journal']
    if 'SO' in infoDict:
        dateStr = re.sub(";:", ":", infoDict['SO'])
    else:
        dateStr1 = infoDict['PubDate']
        dateStr = ""
        if 'Volume' in infoDict:
            dateStr = dateStr + ";%s" % infoDict['Volume']
        if 'Issue' in infoDict:
            dateStr = dateStr + "(%s)" % infoDict['Issue']
        if 'Pagination' in infoDict:
            dateStr = dateStr + "%s" % infoDict['Pagination']
        elif 'Pages' in infoDict:
            dateStr = dateStr + ":%s" %infoDict['Pages']
        dateStr = re.sub(r"\s+","", dateStr)
        dateStr = dateStr1 + dateStr
    if not re.search(r"\.$", dateStr):
        dateStr = dateStr + "."
    # construct hyperlinks for PMID and PMCID (if it exists)
    if (not plain):
        idStr = "PMID: <a href=\"%s\" target=\"_blank\">%s</a>" % (htmlEscape(infoDict['url']), infoDict['pubmed'])
        if 'pmc' in infoDict:
            idStr = idStr + "; PMC: <a href=\"https://www.ncbi.nlm.nih.gov/pmc/articles/%s/\" target=\"_blank\">%s</a>" % (infoDict['pmc'], infoDict['pmc'])
        if doi and 'doi' in infoDict:
            idStr = ("DOI: <a href=\"https://doi.org/%s\" target=\"_blank\">%s</a>; " % (htmlEscape(infoDict['doi']), infoDict['doi'] ) ) + idStr
    else:
        idStr = "PMID: <a href=\"%s\" target=\"_blank\">%s</a>" % (infoDict['url'], infoDict['pubmed'])
        if 'pmc' in infoDict:
            idStr = idStr + "; PMC: <a href=\"https://www.ncbi.nlm.nih.gov/pmc/articles/%s/\" target=\"_blank\">%s</a>" % (infoDict['pmc'], infoDict['pmc'])
        if doi and 'doi' in infoDict:
            idStr = ("DOI: <a href=\"https://doi.org/%s\" target=\"_blank\">%s</a>; " % (infoDict['doi'], infoDict['doi'] ) ) + idStr

    # now that the pubmed link has been constructed, we can overwrite the url in infoDict with the original article URL
    origComment = ""
    origUrl = infoDict['url']
    fullTextUrl = fetchFullTextUrl(infoDict['pubmed'])
    if fullTextUrl:
        infoDict['url'] = fullTextUrl
    else:
        # not an error: plenty of older papers have no full text provider at NCBI at all.
        # Keep the PubMed link, but say so on stderr, never on stdout, where the message
        # would end up in the middle of the references.
        sys.stderr.write("warning: NCBI lists no full text link for PMID %s, "
                         "linking to PubMed instead\n" % infoDict['pubmed'])

    htmlLines = list()
    htmlLines.append("<p>")
    htmlLines.append("%s" % authStr)
    if (not plain):
        htmlLines.append("<a href=\"%s\" target=\"_blank\">" % htmlEscape(infoDict['url']))
        htmlLines.append("%s</a>." % htmlEscape(title))
        htmlLines.append("<em>%s</em>. %s" % (htmlEscape(journal), dateStr))
    else:
        htmlLines.append("<a href=\"%s\" target=\"_blank\">" % infoDict['url'])
        htmlLines.append("%s</a>." % title)
        htmlLines.append("<em>%s</em>. %s" % (journal, dateStr))
    htmlLines.append("%s" % idStr)
    htmlLines.append("</p>")
    if plain:
        htmlLines = list()
        idStr = "PMID: %s" % infoDict['pubmed'];
        if 'pmc' in infoDict:
            idStr = idStr + "; PMC: %s" % infoDict['pmc']
        if 'doi' in infoDict:
            idStr = ("DOI: %s; " % infoDict['doi']) + idStr
        htmlLines.append("%s %s. %s. %s %s. %s" % (authStr, title, journal, dateStr, idStr, infoDict['url']))
    if not plain and origComment:
        htmlLines.append("%s" %  origComment)
    sep = "\n"
    space = " "

    processLines = list()

    for i in htmlLines:
        processLines.append(textwrap.fill(i, 100))
    htmlStr = sep.join(processLines)
    return htmlStr, authStr

def main():
    parser = argparse.ArgumentParser(
        description='Turns PubMed Ids and PubMedCentral Ids into GB formatted citations in html',
        epilog='example: getTrackReferences PMC3039671 21347206'
        )
    parser.add_argument('ids', metavar='IDs', type=str, nargs='+', help='The list of PubMed and PubMedCentral Ids')
    parser.add_argument('-p', '--plain', action="store_true", default=0, help="Output the references in plain-text instead of html.")
    parser.add_argument('-v', '--verbose', action="store_true", default=0, help="Output the full author list instead of truncating after the first 10.")
    parser.add_argument('-d', '--doi', action="store_true", default=0, help="Include a DOI link.")
    args = parser.parse_args()
    ids = args.ids
    references = dict()

    for i in ids:
        infoDict = parseInit(i)
        html, authStr = makeHtml(infoDict, args.plain, args.verbose, args.doi)
        references[authStr] = html
    print("\n")
    for i in sorted(references.keys()):
        print("%s\n" % references[i])

if __name__ == "__main__":
    main()
