import xml.etree.ElementTree as ET
from typing import Union
from pathlib import Path
from typing import List


def from_ufsac_to_tsv(
    xml_file: Union[str, Path],
    conll_file: Union[str, Path],
    datasetname: str,
    encoding: str = "utf8",
    cut_multisense: bool = True,
):
    """
    Function that converts the UFSAC format into tab separated column format in a new file.
    Parameters
    ----------
    xml_file : Union[str, Path]
        Path to the xml file.
    conll_file : Union[str, Path]
        Path for the new conll file.
    datasetname: str
        Name of the dataset from UFSAC, needed because of different handling of multi-word-spans in the datasets
    encoding : str, optional
        Encoding used in open function. The default is "utf8".
    cut_multisense : bool, optional
        Boolean that determines whether or not the wn30_key tag should be cut if it contains multiple possible senses.
        If True only the first listed sense will be used. Otherwise the whole list of senses will be detected
        as one new sense. The default is True.

    """

    def make_line(word, begin_or_inside, attributes):
        """
        Function that creates an output line from a word.
        Parameters
        ----------
        word :
            String of the actual word.
        begin_or_inside:
            Either 'B-' or 'I-'
        attributes:
            List of attributes of the word (pos, lemma, wn30_key)
        """
        line = word
        if cut_multisense:
            attributes[-1] = attributes[-1].split(";")[0]  # take only first sense

        for attrib in attributes:
            if attrib != "O":
                line = line + "\t" + begin_or_inside + attrib
            else:
                line = line + "\tO"
        line += "\n"

        return line

    def split_span(word_fields: List[str], datasetname: str):
        """
        Function that splits a word if necessary, i.e. if it is a multiple-word-span.
        Parameters
        ----------
        word_fields :
            list ['surface_form', 'lemma', 'pos', 'wn30_key'] of a word
        datasetname:
            name of corresponding dataset
        """

        span = word_fields[0]

        if datasetname in [
            "trainomatic",
            "masc",
        ]:  # splitting not sensible for these datasets
            return [span]
        elif datasetname == "omsti":
            if (
                word_fields[3] != "O" and not span == "_" and "__" not in span
            ):  # has annotation and does not consist only of '_' (still not 100% clean)
                return span.split("_")
            else:
                return [span]
        else:  # for all other datasets splitting at '_' is always sensible
            return span.split("_")

    txt_out = open(file=conll_file, mode="w", encoding=encoding)

    tree = ET.parse(xml_file)
    corpus = tree.getroot()

    number_of_docs = len(corpus.findall("document"))

    fields = ["surface_form", "lemma", "pos", "wn30_key"]
    for document in corpus:
        # Docstart
        if number_of_docs > 1:
            txt_out.write("-DOCSTART-\n\n")

        for paragraph in document:

            for sentence in paragraph:

                for word in sentence:

                    dictionary = word.attrib
                    fields_of_word = [word.attrib[field] if (field in dictionary) else "O" for field in fields]

                    chunks = split_span(fields_of_word, datasetname)

                    txt_out.write(make_line(chunks[0], "B-", fields_of_word[1:]))

                    # if there is more than one word in the chunk we write each in a separate line
                    for chunk in chunks[1:]:
                        # print(chunks)
                        txt_out.write(make_line(chunk, "I-", fields_of_word[1:]))

                # empty line after each sentence
                txt_out.write("\n")

    txt_out.close()