from typing import List, Dict, Tuple
import django, argparse, json, zipfile, re
django.setup()
from tqdm import tqdm
from dataclasses import dataclass
from sefaria.model import *

VERSION_NOTES = """This text was digitized and released into the public domain by <a href="https://dicta.org.il">Dicta: The Israel Center for Text Analysis</a>, using a state-of-the-art OCR pipeline leveraging a custom-built transformer-based language model for Rabbinic Hebrew. Nikud (Vocalization) marks were auto-generated by <a href="https://nakdanlive.dicta.org.il">Dicta's Nakdan system</a>."""
DICTA_PARALLELS_INDEX_NAME = "Dicta Parallels"
DICTA_PARALLELS_INDEX_NAME_HE = "המקבילות של דיקטה"


def create_simple_index(en_title, he_title, categories, section_names, address_types=None, collective=None, base_text_titles=None):
    root = JaggedArrayNode()
    root.add_primary_titles(en_title, he_title)
    root.add_structure(section_names, address_types=address_types)
    root.key = en_title
    root.validate()
    if collective:
        create_term(collective['en'], collective['he'])
    index = {
        "title": en_title,
        "categories": categories,
        "schema": root.serialize(),
    }
    if collective:
        index['collective_title'] = collective['en']
    if base_text_titles:
        index['base_text_titles'] = base_text_titles
    return index


def create_term(en_title, he_title, scheme='toc_categories'):
    term_dict = {
        'name': en_title,
        'scheme': scheme,
        'titles': [{'lang': 'en', 'text': en_title, 'primary': True}, {'lang': 'he', 'text': he_title, 'primary': True}]
    }
    save_term(term_dict)


def create_category(path, shared_title_en, shared_title_he):
    cat_dict = {
        "path": path,
        "sharedTitle": shared_title_en,
    }
    create_term(shared_title_en, shared_title_he)
    save_category(cat_dict)


def save_index(index):
    try:
        Index().load({"title": index['title']}).delete()
    except:
        pass
    Index(index).save()


def save_text(text):
    try:
        Version().load({"title": text['title'], "versionTitle": text['versionTitle'], "language": text['language']}).delete()
    except:
        pass
    Version(text).save()


def save_links(links):
    parallel_reg = re.compile(fr'^{DICTA_PARALLELS_INDEX_NAME}')
    LinkSet({"refs": parallel_reg}).delete()
    for l in tqdm(links, desc='links'):
        Link(l).save()


def save_term(term):
    try:
        Term().load({"name": term['name']}).delete()
    except:
        pass
    Term(term).save()


def save_category(cat_dict):
    try:
        Category().load({"path": cat_dict['path']}).delete(force=True)
    except:
        pass
    Category(cat_dict).save()


class DictaLibraryManager:

    def __init__(self, lib_path):
        self.lib_path = lib_path
        self.books_by_name = {}
        self.parallels = []  # list of all parallels in all parsed books
        with open(f"{self.lib_path}/books.json", "r") as fin:
            self.books_list = [DictaBook(lib_path=lib_path, **book_dict) for book_dict in json.load(fin)]
        for book in self.books_list:
            self.books_by_name[book.displayName] = book

    def get_book(self, name):
        return self.books_by_name[name]

    def parse_and_save_book(self, book_name: 'DictaBook'):
        book = self.get_book(book_name)
        self.parallels += book.parse()
        book.save()

    def create_and_save_parallels_commentary(self, book: 'DictaBook'):
        en_title = DICTA_PARALLELS_INDEX_NAME + " on " + book.displayNameEnglish
        he_title = DICTA_PARALLELS_INDEX_NAME_HE + " על " + book.displayName
        index = create_simple_index(en_title, he_title, ["Reference", "Dicta Parallels"],
                                    ["Chapter", "Paragraph"], ["Integer", "Integer"],
                                    collective={"en": DICTA_PARALLELS_INDEX_NAME, "he": DICTA_PARALLELS_INDEX_NAME_HE},
                                    base_text_titles=[book.displayNameEnglish]
                                    )
        links = []
        parallels_text = []
        for ipage, page_parallels in enumerate(self.parallels):
            parallels_text += [[]]
            page_parallels.sort(key=lambda x: x.dataOrder)
            for iparallel, parallel in enumerate(page_parallels):
                parallels_text[-1] += [parallel.serialize()]
                parallel_tref = f'{index["title"]} {ipage+1}:{iparallel+1}'
                links += [{
                    "refs": [parallel.base_tref, parallel_tref],
                    "type": "commentary",
                    "auto": True,
                    "generated_by": "dicta library parse",
                    "inline_reference": {
                        "data-order": parallel.dataOrder,
                        "data-commentator": DICTA_PARALLELS_INDEX_NAME,
                        "data-label": parallel.dataLabel,
                    }
                }]
        version = {
            "title": en_title,
            "chapter": parallels_text,
            "versionTitle": "Dicta Library",
            "versionSource": "https://library.dicta.org.il",
            "language": "he",
            "versionNotes": VERSION_NOTES
        }
        print(f"Saving: {en_title}")
        create_category(index['categories'], DICTA_PARALLELS_INDEX_NAME, DICTA_PARALLELS_INDEX_NAME_HE)
        save_index(index)
        save_text(version)
        save_links(links)


@dataclass
class DictaParallel:
    baseTextLength: int
    baseStartChar: int
    baseMatchedText: str
    sources: list
    tool: str
    baseStartToken: int = None
    baseEndToken: int = None
    dataOrder: int = None
    dataLabel: str = None
    base_tref: str = None

    def serialize(self):
        base_text = self.baseMatchedText
        sorted_comps = sorted(self.sources, key=lambda x: len(x.get('compMatchedText', '')), reverse=True)
        sorted_comps = sorted_comps[:1]
        return f"<b>{base_text[:1000]}</b><br>{'<br><br>'.join([self.__serialize_comp(comp) for comp in sorted_comps])}"

    @staticmethod
    def __serialize_comp(comp):
        comp_text = comp['compDisplayText']
        comp_name = comp.get('compNameHe', comp.get('verseDispHe', ''))
        return f"""<small>{comp_text[:1000]}<br>({comp_name})</small>"""

@dataclass
class DictaPage:
    displayName: str
    fileName: str
    nakdanResponseFile: str

    def parse(self, root_dir, book_title, with_nikud=False):
        jin = self.__get_json_content(root_dir)
        tokens = self.__get_tokens(jin)
        index = self.__get_zero_based_index()
        parallels = self.__get_parallels(jin)
        nikud = self.__get_nikud(jin)
        paragraphs, token2seg, parallel_id2token, parallel_id2_data_order = self.serialize(tokens, nikud, with_nikud)
        self.__add_link_info_to_parallels(parallels, token2seg, parallel_id2token, parallel_id2_data_order, book_title, index)
        return paragraphs, index, parallels

    @staticmethod
    def __add_link_info_to_parallels(parallels, token2seg, parallel_id2token: dict, parallel_id2_data_order, book_title, page_index):
        # HACK: try to fix missing parallel IDs
        new_parallel_id2token = {}
        new_parallel_id2data_order = {}
        parallel_ids = sorted(parallel_id2token.keys())
        prev_pid = -1
        for pid in parallel_ids:
            new_pid = (prev_pid + 1) if (pid > prev_pid + 1) else pid
            new_parallel_id2token[new_pid] = parallel_id2token[pid]
            new_parallel_id2data_order[new_pid] = parallel_id2_data_order[pid]
            prev_pid = new_pid
        parallel_id2token = new_parallel_id2token
        parallel_id2_data_order = new_parallel_id2data_order
        # END HACK

        for parallel_id, parallel in enumerate(parallels):
            data_order = parallel_id2_data_order[parallel_id]
            itoken = parallel_id2token[parallel_id]
            iseg = token2seg[itoken]

            # modify parallel
            parallel.dataOrder = data_order
            parallel.dataLabel = str(data_order)
            parallel.base_tref = f"{book_title} {page_index+1}:{iseg}"

    @staticmethod
    def serialize(tokens, nikud, with_nikud=False) -> Tuple[List[str], Dict[int, int], Dict[int, int], Dict[int, int]]:
        """
        Serialize data into list of strings which represents a section in the DB
        :param tokens: token dicts
        :return:
        """
        text = ""
        parallel_ids_seen = set()
        token2seg = {}
        parallel_id2token = {}
        parallel_id2_data_order = {}
        curr_seg = 1
        parallel_data_order = 1
        for itoken, token in enumerate(tokens):
            # parallel markers
            token2seg[itoken] = curr_seg
            if token['str'] == '\n': curr_seg += 1
            curr_parallel_ids = set(token.get('sourcesPostProcessedIDs', []))
            unseen_parallel_ids = curr_parallel_ids - parallel_ids_seen
            for new_parallel_id in unseen_parallel_ids:
                # parallel start
                parallel_id2_data_order[new_parallel_id] = parallel_data_order
                parallel_id2token[new_parallel_id] = itoken
                marker = f"""<i data-commentator="{DICTA_PARALLELS_INDEX_NAME}" data-label="{parallel_data_order}" data-order="{parallel_data_order}"></i>"""
                parallel_data_order += 1
                text += marker
            parallel_ids_seen |= unseen_parallel_ids
            token_str = nikud[token['nikudID']] if (with_nikud and 'nikudID' in token) else token['str']
            text += token_str
        return text.split('\n'), token2seg, parallel_id2token, parallel_id2_data_order

    def __get_zero_based_index(self):
        m = re.search(r'\d+$', self.displayName)
        return int(m.group(0)) - 1

    def __get_json_content(self, root_dir):
        json_fname = self.fileName.replace('.zip', '.json')
        with zipfile.ZipFile(f"{root_dir}/{self.fileName}") as zin:
            with zin.open(json_fname, mode='r') as fin:
                return json.load(fin)

    @staticmethod
    def __get_tokens(jin):
        return jin['tokens']

    @staticmethod
    def __get_parallels(jin: dict) -> List[DictaParallel]:
        parallels = filter(lambda x: x and x['tool'] == 'parallels', jin['data']['postProcessedSources'])
        # TODO relying on existence of 'baseStartToken' which doesn't exist in 'citation's for some reason
        return [DictaParallel(**parallel_dict) for parallel_dict in parallels]

    @staticmethod
    def __get_nikud(jin: dict):
        return [x['options'][0]['w'] for x in jin['data']['nikudResults']]


@dataclass
class DictaBook:
    lib_path: str
    displayName: str
    fileName: str
    printYear: int
    printLocation: str
    author: str
    category: str
    categoryEnglish: str = None
    authorEnglish: str = None
    printLocationEnglish: str = None
    displayNameEnglish: str = None
    source: str = None
    firstpage: dict = None  # seems to be a non-standard field
    pages: List[DictaPage] = None

    def __post_init__(self):
        self._root_path = f"{self.lib_path}/{self.fileName}"
        self._parsed_pages = []
        self.__load_pages()

    def __load_pages(self):
        try:
            with open(f"{self._root_path}/pages.json", "r") as fin:
                self.pages = [DictaPage(**page_dict) for page_dict in json.load(fin)]
        except FileNotFoundError:
            print(f"No directory for {self.displayName}")

    def parse(self):
        self.__create_index()
        parallels = self.__create_version()
        return parallels

    def save(self):
        print(f"Saving: {self.displayNameEnglish}")
        save_index(self.index)
        save_text(self.version)

    def __create_index(self):
        self.index = create_simple_index(self.displayNameEnglish, self.displayName, ["Responsa"], ["Daf", "Paragraph"], address_types=['Integer', 'Integer'])

    def __create_version(self):
        parsed_pages = []
        all_parallels = []
        for page in tqdm(self.pages, desc=self.fileName):
            paragraphs, index, parallels = page.parse(self._root_path, self.index['title'], with_nikud=True)
            all_parallels += [list(parallels)]
            while len(parsed_pages) < index:
                parsed_pages += [[]]
                # all_parallels += [[]] TODO currently don't require parallels to have same pagination as base text
            parsed_pages += [paragraphs]
        self.version = {
            "title": self.displayNameEnglish,
            "chapter": parsed_pages,
            "versionTitle": "Dicta Library",
            "versionSource": "https://library.dicta.org.il",
            "language": "he",
            "versionNotes": VERSION_NOTES
        }
        return all_parallels


def init_argparse() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    parser.add_argument("book", help='name of book to parse')
    parser.add_argument('-l', '--library-path', dest='library_path')
    return parser


if __name__ == '__main__':
    parser = init_argparse()
    args = parser.parse_args()
    dicta = DictaLibraryManager(args.library_path)
    dicta.parse_and_save_book(args.book)
    book = dicta.get_book(args.book)
    dicta.create_and_save_parallels_commentary(book)

# https://dicta-library.cauldron.sefaria.org

"""
Steps to run on server

kubectl cp "~/sefaria/data/sources/Dicta Library/parse.py" $POD:/app/data
kubectl exec -it $POD -- bash
cd data
git clone https://github.com/Dicta-Israel-Center-for-Text-Analysis/dicta-library-data.git
python parse.py "דברי ריבות" -l "dicta-library-data"
"""
