# Based on https://github.com/darwindarak/mdx_bib
import re
import string
from collections import Counter, OrderedDict
from typing import Tuple
from xml.etree import ElementTree as etree
from xml.etree.ElementTree import tostring as etree_to_string

from markdown.extensions import Extension
from markdown.inlinepatterns import Pattern
from markdown.preprocessors import Preprocessor
from mkdocs.config.config_options import Type as MkType
from mkdocs.plugins import BasePlugin
from pybtex.database.input import bibtex
from pybtex.exceptions import PybtexError

BRACKET_RE = re.compile(r"\[([^\[]+)\]")
CITE_RE = re.compile(r"@([\w_:-]+)")
DEF_RE = re.compile(r"\A {0,3}\[@([\w_:-]+)\]:\s*(.*)")
INDENT_RE = re.compile(r"\A\t| {4}(.*)")

CITATION_RE = r"(\[@(?:[\w_:-]+)(?: *, *@(?:[\w_:-]+))*\])"


class Bibliography(object):
    """Keep track of document references and citations for exporting"""

    def __init__(self, extension, plugin, bibtex_file, order):
        self.extension = extension
        self.order = order
        self.plugin = plugin

        self.citations = OrderedDict()
        self.references = dict()

        if bibtex_file:
            try:
                parser = bibtex.Parser()
                self.bibsource = parser.parse_file(bibtex_file).entries
                self.labels = {
                    id: self.formatCitation(self.bibsource[id])
                    for id in self.bibsource.keys()
                }
                for value, occurrences in Counter(self.labels.values()).items():
                    if occurrences > 1:
                        for xkey, xvalue in self.labels.items():
                            i = 0
                            if xvalue == value:
                                self.labels[
                                    xkey
                                ] = f"{xvalue}{string.ascii_lowercase[i]}"
                                i += 1

            except PybtexError:
                print("Error loading bibtex file")
                self.bibsource = dict()
                self.labels = {}
        else:
            self.bibsource = dict()

    def addCitation(self, citekey):
        self.citations[citekey] = self.citations.get(citekey, 0) + 1

    def setReference(self, citekey, reference):
        self.references[citekey] = reference

    def citationID(self, citekey):
        return "cite-" + citekey

    def referenceID(self, citekey):
        return "ref-" + citekey

    def formatAuthor(self, author):
        out = (
            author.last_names[0]
            + ((" " + author.first_names[0][0]) if author.first_names else "")
            + "."
        )
        if author.middle_names:
            out += f"{author.middle_names[0][0]}."
        return out.replace("{", "").replace("}", "")

    def formatAuthorSurname(self, author):
        out = author.last_names[0]
        return out.replace("{", "").replace("}", "")

    def formatReference(self, ref):
        author_list = list(map(self.formatAuthor, ref.persons["author"]))

        if len(author_list) == 1:
            authors = author_list[0]
        else:
            authors = ", ".join(author_list[:-1])
            authors += f" and {author_list[-1]}"

        # Harvard style
        # Surname, Initial, ... and Last_Surname,
        # Initial, Year. Title. Journal, Volume(Issue), pages. doi.

        title = ref.fields["title"].replace("{", "").replace("}", "")
        journal = ref.fields.get("journal", "")
        volume = ref.fields.get("volume", "")
        issue = ref.fields.get("issue", "")
        year = ref.fields.get("year")
        pages = ref.fields.get("pages")
        doi = ref.fields.get("doi")
        url = ref.fields.get("url")

        ref_id = self.referenceID(ref.key)
        ref = f"<p id={repr(ref_id)}>{authors}, {year}. {title}."
        if journal:
            ref += f" <i>{journal}</i>."
            if volume:
                ref += f" <i>{volume}</i>"
            if issue:
                ref += f"({issue})"
            if pages:
                ref += f", pp.{pages}"
            ref += "."
        if doi:
            ref += f' <a href="https://dx.doi.org/{doi}" target="_blank">{doi}</a>'
        elif url:
            ref += f' <a href="{url}" target="_blank">{url}</a>'
        ref += "</p>"

        return etree.fromstring(ref)

    def formatCitation(self, ref):
        author_list = list(map(self.formatAuthorSurname, ref.persons["author"]))
        year = ref.fields.get("year")

        if len(author_list) == 1:
            citation = f"{author_list[0]}"
        elif len(author_list) == 2:
            citation = f"{author_list[0]} and {author_list[1]}"
        else:
            citation = f"{author_list[0]} et al."

        citation += f", {year}"

        return citation

    def make_bibliography(self):
        if self.order == "alphabetical":
            raise (NotImplementedError)

        div = etree.Element("div")
        div.set("class", "footnote")
        div.append(etree.Element("hr"))
        ol = etree.SubElement(div, "ol")

        if not self.citations:
            return div

        # table = etree.SubElement(div, "table")
        # table.set("class", "references")
        # tbody = etree.SubElement(table, "tbody")
        etree.SubElement(div, "div")
        for id in self.citations:
            li = etree.SubElement(ol, "li")
            li.set("id", self.referenceID(id))
            # ref_id = etree.SubElement(li, "td")
            ref_txt = etree.SubElement(li, "p")
            if id in self.references:
                self.extension.parser.parseChunk(ref_txt, self.references[id])
            elif id in self.bibsource:
                ref_txt.append(self.formatReference(self.bibsource[id]))
            else:
                ref_txt.text = "Missing citation for {}".format(id)

        return div

    def clear_citations(self):
        self.citations = OrderedDict()


class CitationsPreprocessor(Preprocessor):
    """Gather reference definitions and citation keys"""

    def __init__(self, bibliography):
        self.bib = bibliography

    def subsequentIndents(self, lines, i):
        """Concatenate consecutive indented lines"""
        linesOut = []
        while i < len(lines):
            m = INDENT_RE.match(lines[i])
            if m:
                linesOut.append(m.group(1))
                i += 1
            else:
                break
        return " ".join(linesOut), i

    def run(self, lines):
        linesOut = []
        i = 0

        while i < len(lines):
            # Check to see if the line starts a reference definition
            m = DEF_RE.match(lines[i])
            if m:
                key = m.group(1)
                reference = m.group(2)
                indents, i = self.subsequentIndents(lines, i + 1)
                reference += " " + indents

                self.bib.setReference(key, reference)
                continue

            # Look for all @citekey patterns inside hard brackets
            for bracket in BRACKET_RE.findall(lines[i]):
                for c in CITE_RE.findall(bracket):
                    self.bib.addCitation(c)
            linesOut.append(lines[i])
            i += 1

        return linesOut


class CitationsPattern(Pattern):
    """Handles converting citations keys into links"""

    def __init__(self, pattern, bibliography):
        super(CitationsPattern, self).__init__(pattern)
        self.bib = bibliography

    def handleMatch(self, m):
        span = etree.Element("span")
        for cite_match in CITE_RE.finditer(m.group(2)):
            id = cite_match.group(1)
            if id in self.bib.bibsource:
                a = etree.Element("a")
                a.set("id", self.bib.citationID(id))
                a.set("href", "./#" + self.bib.referenceID(id))
                a.set("class", "citation")
                a.text = self.bib.labels[id]
                span.append(a)
            else:
                continue
        if len(span) == 0:
            return None
        return span


context_citations = None


class CitationsExtension(Extension):
    def __init__(self):
        super(CitationsExtension, self).__init__()
        self.bib = None

    def extendMarkdown(self, md):
        md.registerExtension(self)
        self.parser = md.parser
        self.md = md

        md.preprocessors.register(CitationsPreprocessor(self.bib), "mdx_bib", 15)
        md.inlinePatterns.register(
            CitationsPattern(CITATION_RE, self.bib), "mdx_bib", 175
        )


def makeExtension(*args, **kwargs):
    return CitationsExtension(*args, **kwargs)


class BibTexPlugin(BasePlugin):
    config_scheme: Tuple[Tuple[str, MkType]] = (
        ("bibtex_file", MkType(str)),  # type: ignore[assignment]
        ("order", MkType(str, default="unsorted")),  # type: ignore[assignment]
    )

    def __init__(self):
        self.citations = None

    def on_config(self, config, **kwargs):
        extension = CitationsExtension()
        self.bib = Bibliography(
            extension,
            self,
            self.config["bibtex_file"],
            self.config["order"],
        )
        extension.bib = self.bib
        config["markdown_extensions"].append(extension)

    def on_page_content(self, html, page, config, files):
        html += "\n" + etree_to_string(self.bib.make_bibliography()).decode()
        self.bib.clear_citations()
        return html