
import os
from pathlib import Path

from bs4 import BeautifulSoup

from text2story.readers import token_corpus
from text2story.readers.read import Read
import xml.etree.ElementTree as ET

import spacy
from text2story.readers.token_corpus import TokenRelation
import re
from text2story.core.utils import find_substring_match

class ReadACE(Read):

    def __init__(self):
        self.file_lst = []
        if not(spacy.util.is_package('en_core_web_lg')):
            spacy.cli.download('en_core_web_lg')

    def _get_rel_mention_data(mentions, raw_text):

        info_lst = []
        current_idx = 0
        mention_lst =  []

        for mention in mentions:
            mention_id = mention.get('ID')
            #print("-->", mention_id) CNN_CF_20030303.1900.00-E84-167
            arguments_lst = mention.findall(".//relation_mention_argument")
            args = []
            for arg in arguments_lst:
                extent = arg.find(".//extent/charseq")
                arg_id = arg.get("REFID")
                arg_role = arg.get("ROLE")
                args.append((arg_id, arg_role))

            mention_lst.append((mention_id, args))

        return mention_lst


    @staticmethod
    def _get_mention_data(mentions, raw_text, mention_regex=".//extent/charseq"):

        info_lst = []
        current_idx = 0
        mention_lst =  []

        for mention in mentions:
            mention_id = mention.get('ID')
            #print("-->", mention_id) CNN_CF_20030303.1900.00-E84-167

            extent = mention.find(mention_regex)
            start_mention = extent.get('START')
            end_mention = extent.get('END')

            mention_lst.append((mention_id, extent.text, int(start_mention), int(end_mention)))

        mention_lst.sort(key = lambda x:x[2])
        old_mention_start = None
        old_mention_end = None

        for mention in mention_lst:

            mention_id, mention_text, mention_start, mention_end = mention

            if old_mention_start !=None and old_mention_end !=None:
                if mention_start>= old_mention_start and mention_end<=old_mention_end:
                    # if the current mention is inside of the mention before
                    # the current_idx should be restarted
                    current_idx = start

            pattern = r'\b' + re.escape(mention_text) + r'\b'
            matches = re.finditer(pattern, raw_text, flags=re.IGNORECASE)
            span_idx = find_substring_match(matches, current_idx)

            if span_idx is None:
               span_idx = raw_text.find(mention_text, current_idx)

            start = span_idx
            end = span_idx + len(mention_text)

            current_idx = end
            old_mention_start = mention_start
            old_mention_end = mention_end


            info_lst.append((mention_id, mention_text, start, end))

        return info_lst

    @staticmethod
    def _get_participant_annotations(root, raw_text):

        participant_lst = []  # first the participants

        for entity in root.findall(".//entity"):
            entity_type = entity.get('TYPE')
            entity_mentions = entity.findall('entity_mention')

            mention_lst = ReadACE._get_mention_data(entity_mentions, raw_text)

            participant_lst.append(({'Participant_Type_Domain':entity_type}, mention_lst))

        return participant_lst

    @staticmethod
    def _map2isorelation(ace_relation, ace_subrelation):
        ace_relation2iso_relation = {
            "PART-WHOLE":"OLINK_partOf",
            "ORG-AFF":"OLINK_memberOf",
            "PHYS":"QSLINK_ground"
        }

        return ace_relation2iso_relation.get(ace_relation)
    @staticmethod
    def _map2isocategory(ace_category):

        # mapping generated by chatGPT
        event_type_to_category = {
            "Conflict": "Occurrence",
            "Contact": "Contact",
            "Transaction": "Transaction",
            "Movement": "Occurrence",
            "Personnel": "State",
            "Life": "Occurrence",
            "Business": "State",
            "Justice": "State",
            "Contact.Meet": "Reporting",
            "Contact.Correspondence": "Reporting",
            "Contact.Broadcast": "Reporting",
        }

        # Adding additional event types
        #event_type_to_category["Contact.Phone-Write"] = "Reporting"
        #event_type_to_category["Contact.Speech"] = "Perception"
        #event_type_to_category["Transaction.Transfer-Money"] = "Transaction"
        #event_type_to_category["Business.Merge"] = "Business"
        #event_type_to_category["Business.Declare-Bankruptcy"] = "Business"
        #event_type_to_category["Life.Marry"] = "State"

        return event_type_to_category[ace_category]

    @staticmethod
    def _get_relations_annotations(root, raw_text):
        relations_lst = []

        for relation in root.findall(".//relation"):
            relation_type = relation.get('TYPE')
            relation_subtype = relation.get('SUBTYPE')
            relation_mentions = relation.findall(".//relation_mention")
            mention_lst = ReadACE._get_rel_mention_data(relation_mentions, raw_text)

            relations_lst.append(({'Relation_Type': relation_type,'Relation_Subtype': relation_subtype}, mention_lst))

        return relations_lst

    @staticmethod
    def _get_event_annotations(root, raw_text):

        event_lst = []  # then the events

        for event in root.findall(".//event"):
            event_type = event.get('TYPE')
            event_subtype = event.get('SUBTYPE')
            modality = event.get('MODALITY')
            event_mentions = event.findall(".//event_mention")

            mention_lst = ReadACE._get_mention_data(event_mentions, raw_text, mention_regex=".//anchor/charseq")

            event_lst.append(({'Class':ReadACE._map2isocategory(event_type), \
                               'Subclass':event_subtype}, mention_lst))


        return event_lst

    @staticmethod
    def _get_timex_annotations(root, raw_text):

        timex_lst = []  # then the events
        for timex in root.findall(".//timex2"):

            timex_id = timex.get('ID')
            timex_val = timex.get('VAL')
            timex_anchor_val = timex.get('ANCHOR_VAL')
            timex_anchor_dir = timex.get('ANCHOR_DIR')

            timex_mentions = timex.findall(".//timex2_mention")

            mention_lst = ReadACE._get_mention_data(timex_mentions, raw_text)
            timex_lst.append(({'Timex_Id':timex_id, 'Timex_Val':timex_val,\
                                'Timex_Anchor_Val':timex_anchor_val, 'Timex_Anchor_Dir':timex_anchor_dir}, mention_lst))

        return timex_lst

    def _read_label_file(self, apf_file_path, raw_text):

        # Load and parse the APF XML file
        tree = ET.parse(apf_file_path)
        root = tree.getroot()

        current_idx = 0

        # Iterate through the APF XML to extract annotations
        participant_lst = self._get_participant_annotations(root, raw_text)  # first the participants

        event_lst = self._get_event_annotations(root, raw_text)  # then the events

        timex_lst = self._get_timex_annotations(root, raw_text)

        relations_lst = self._get_relations_annotations(root, raw_text)

        return participant_lst, event_lst, timex_lst, relations_lst

    @staticmethod
    def _search_subtoken(start_idx, end_idx, annotations):

        ann_lst = []
        for category in annotations:

            for label_lst in annotations[category]:
                if label_lst == []:
                    continue

                mention_lst = label_lst[-1]

                for  mention in mention_lst:
                    mention_id, mention_text, mention_start, mention_end = mention

                    # the given token is inside this mention
                    if  start_idx >= mention_start and end_idx <= mention_end:
                        mention_annotations = label_lst[0]
                        if mention_annotations != {}:
                            ann_lst.append((category, mention_annotations, mention))
        return ann_lst

    @staticmethod
    def _read_text_file(data_file):
        with open(data_file, "r") as fd:
            file_content = fd.read()
            soup = BeautifulSoup(file_content, 'html.parser')

        # 'bs4.element.ResultSet'
        body = soup.doc.body
        text = ""

        for child in body.children:
            if child.name == 'headline' or child.name == 'text':
                text = text + child.text

        return text.strip()

    def process_file(self, data_file, lang="en"):
        """
        Process an ACE file format (.sgm extension) and return a tokenCorpus list.

        It assumes that in the same path of .sgm file there is an .apf.xml file with the
        same name.
        @param data_file:
        @return:
        """

        label_file = os.path.splitext(data_file)[0] + ".apf.xml"
        text = ReadACE._read_text_file(data_file)

        # get the labels
        participant_lst, event_lst, timex_lst, relations_lst = self._read_label_file(label_file, text)

        annotations = {'Participant':participant_lst,\
                       'Event':event_lst,\
                       'Time':timex_lst}

        # tokenization of the raw text and serialization to a list of TokenCorpus data type
        # if I use spacy to tokenization, then I have idx
        nlp = spacy.load('en_core_web_lg')

        doc = nlp(text)
        tok_lst = []

        sent_id = -1
        ref_count = 1

        # a map between referecens in ACE dataset and new ids in TokenCorpus format
        mention2ref = {}
        for tok_count, tok in enumerate(doc):

            if tok.is_sent_start:
                sent_id += 1

            mytok = token_corpus.TokenCorpus()
            mytok.id = tok_count
            mytok.text = tok.text
            mytok.lemma = tok.lemma_
            mytok.pos = tok.pos_
            mytok.dep = tok.dep_
            mytok.head = tok.head.text
            mytok.head_pos = tok.head.pos_
            mytok.head_lemma = tok.head.lemma_
            mytok.offset = tok.idx
            mytok.sent_id = sent_id

            ans_sub = ReadACE._search_subtoken(tok.idx, tok.idx + len(tok.text),\
                                           annotations)
            for ans in ans_sub:
                category, ann, mention = ans
                mention_id, _ , _ , _ = mention

                # this mention already exists as TokenCorpus Annotation?
                if mention_id in mention2ref:
                    # yes!
                    ref_tok, _ = mention2ref[mention_id]
                    mytok.id_ann.append(ref_tok)
                else:
                    # otherwise create another ref
                    mytok.id_ann.append(f'T{ref_count}')
                    mention2ref[mention_id] = (f'T{ref_count}', mytok)
                    ref_count += 1

                mytok.attr.append((category, ann))

            tok_lst.append(mytok)

        # add relations that this token may be part of
        rel_count = 1
        for rel in relations_lst:

            attr, mention_lst = rel
            relation_type = attr['Relation_Type']
            relation_subtype = attr['Relation_Subtype']

            for mention in mention_lst:
                _, args = mention # args it is a list (arg_id, arg_role)

                tok_arg_lst = []
                not_mapped_flag = False
                for (arg_id, arg_role) in args:
                    if arg_id not in mention2ref:
                        print(f"Warning: {arg_id} is not a mapped element.")
                        not_mapped_flag = True
                        continue
                    ref_tok, tok = mention2ref[arg_id]
                    tok_arg_lst.append((tok, arg_role))

                if not_mapped_flag:
                    continue
                # now update each token in the arguments with the relation
                iso_relation = ReadACE._map2isorelation(relation_type, relation_subtype)
                # if the iso_relation is None, there is no possible map for the internal representation of an ISO relation
                if iso_relation is not None:
                    for tok, arg_role in tok_arg_lst:
                        arg_id = arg_role.lower().replace("-","") # in this case, arg_id is the number of argument arg1 or arg2
                        tok.relations.append(TokenRelation(f'R{rel_count}', [t for t, _ in tok_arg_lst], iso_relation, arg_id=arg_id))
                    rel_count += 1

        # now we can add the OLINK_objIdentity of the same entity mentions
        for mention_id1 in mention2ref:
            for mention_id2 in mention2ref:
                if mention_id1 != mention_id2:
                    stem_id1 = mention_id1.split('-')
                    stem_id1 = "-".join(stem_id1[:-1])



        return tok_lst

    def process(self, data_dir, output_dir=None, split=None):
        """
        It reads a set of files of annotations and text, then returns
        the processed tokens as TokenCorpus type

        @param string: path of data to gather and process
        @param string: if an output directory is provided, then the text files will be put there
        @param string: a file name that contains a list of the files in data dir that
        should be processed by this reader


        @return [[TokenCorpus]]: a list of lists of tokens
        """

        data_tokens = []

        for dirpath, dirnames, filenames in os.walk(data_dir):
            for f in filenames:
                if f.endswith(".sgm"):

                    p = Path(f)
                    data_file = os.path.join(dirpath, f)
                    print(f"Processing {data_file}")
                    token_lst = self.process_file(data_file)

                    if output_dir != None:
                        output_txt = os.path.join(output_dir, p.stem + ".txt")
                    else:
                        output_txt = os.path.join(dirpath, p.stem + ".txt")

                    text = ReadACE._read_text_file(data_file)
                    with open(output_txt, "w") as fd:
                        fd.write(text)


                    self.file_lst.append(output_txt)

                    if len(token_lst) > 0:
                        data_tokens.append(token_lst)
        return data_tokens

