from lxml import etree
import pandas as pd
from itertools import combinations
from string import ascii_lowercase
import os
import json
from PIL import Image

Image.MAX_IMAGE_PIXELS = 933120000

from src.scrape_api import get_other_ids

schema_url = 'http://www.loc.gov/standards/alto/ns-v4# http://www.loc.gov/alto/v4/alto-4-4.xsd'
nsmap = {
    'xlink': 'http://www.w3.org/1999/xlink',
    'xsi': 'http://www.w3.org/2001/XMLSchema-instance'
}
domain = 'iiif-test.onb.ac.at'
domain_img = 'iiif.onb.ac.at'


def load_template(filepath: str):
    """Loads the manifest template (JSON) from FILEPATH.

    :param filepath: path to the JSON template.
    :returns: dict of template.
    """
    with open(filepath, 'r') as f:
        return json.load(f)


def load_alto(filepath: str):
    """Load the ALTO-XML from FILEPATH.

    :param filepath: path to ALTO-XML OCR.
    :returns: etree root object.
    """
    try:
        with open(filepath, 'rb') as f:
            root = etree.parse(f)
            return root
    except FileNotFoundError:
        return None


def gen_annotation(id: str, page_id: str, text_line):
    """Generates an Annotation from an ALTO-XML TextLine.

    :param id: id for the manifest (under MISC).
    :param page_id: id for the canvas page (format pageX).
    :param text_line: etree TextLine object.
    :returns: IIIF Annotation
    """
    x = text_line.xpath('./@HPOS')[0]
    y = text_line.xpath('./@VPOS')[0]
    w = text_line.xpath('./@WIDTH')[0]
    h = text_line.xpath('./@HEIGHT')[0]
    chars = ''
    for elem in text_line.xpath('./*'):
        if elem.tag == 'SP':
            chars += ' '
        else:
            content = elem.xpath('./@CONTENT')
            if type(content) is list:
                chars += content[0]
            else:
                chars += content
    annotation = {
        '@type': 'oa:Annotation',
        'motivation': 'sc:painting',
        'on':
        f'https://iiif-test.onb.ac.at/presentation/MISC/{id}/canvas/{page_id}#xywh={x},{y},{w},{h}',
        'resource': {
            '@type': 'cnt:ContentAsText',
            'chars': chars,
            'format': 'text/plain'
        },
        'within': {
            '@id':
            f'https://iiif-test.onb.ac.at/presentation/MISC/{id}/manifest',
            '@type': 'sc:Manifest'
        }
    }
    return annotation


def gen_annotation_list(alto, id: str, page_id: str, res_id: str):
    """Generates an AnnotationList from ALTO-XML OCR.

    :param alto: etree root object of ALTO-XML.
    :param id: id for the manifest (under MISC).
    :param page_id: id for the canvas page (page\d*).
    :param res_id: id for the AnnotationList resource (\d{8}).
    :returns: dict of AnnotationList.
    """
    text_lines = alto.xpath('//TextLine')
    resources = []
    for text_line in text_lines:
        resources.append(gen_annotation(id, page_id, text_line))

    annotation_list = {
        '@context': 'http://iiif.io/api/presentation/2/context.json',
        '@id':
        f'https://iiif-test.onb.ac.at/presentation/MISC/{id}/resource/{res_id}.json',
        '@type': 'sc:AnnotationList',
        'resources': resources
    }
    return annotation_list


def gen_canvas(alto, id: str, page_id: str, res_id: str, im_path: str,
               ext_id_str: str):
    """Generates one canvas to be included in the manifest's
    `canvases` list.

    :param id: id for the manifest (under MISC).
    :param page_id: id for the canvas page (page\d*).
    :param res_id: id for the AnnotationList resource (\d{8}).
    :param im_path: path to the image on local file system.
    :param ext_id_str: externalId of the digital object (DOD_XXXXX or
    DTL_XXXXX).
    :returns: dict of canvas.
    """
    width, height = Image.open(im_path).size
    im_num = os.path.splitext(os.path.basename(im_path))[0]
    ext_id = ext_id_str.split('_')
    if ext_id[0] == 'DTL':
        ext_id[0] = 'REPO'
    images = [{
        '@type': 'oa:Annotation',
        'motivation': 'sc:painting',
        'on': f'https://{domain}/presentation/MISC/{id}/canvas/{page_id}',
        'resource': {
            '@id':
            f'https://{domain_img}/images/{ext_id[0]}/{ext_id[1]}/{im_num}.jp2/full/full/0/default.jpg',
            '@type': 'dctypes:Image',
            'format': 'image/jpeg',
            'height': height,
            'service': {
                '@context': 'http://iiif.io/api/image/2/context.json',
                '@id':
                f'https://{domain_img}/images/{ext_id[0]}/{ext_id[1]}/{im_num}.jp2',
                'profile': 'http://iiif.io/api/image/2/level2.json'
            },
            'width': width
        }
    }]
    other_content = [{
        '@id':
        f'https://{domain}/presentation/MISC/{id}/resource/{res_id}.json',
        '@type': 'sc:AnnotationList'
    }]
    if alto is not None:
        annotation_list = gen_annotation_list(alto, id, page_id, res_id)
        os.makedirs(f'/home/zenon/annotations/{pair_tree(id)}', exist_ok=True)
        with open(f'/home/zenon/annotations/{pair_tree(id)}/{res_id}.json',
                  'w',
                  encoding='utf-8') as f:
            json.dump(annotation_list, f, indent=True, ensure_ascii=False)
        canvas = {
            '@id': f'https://{domain}/presentation/MISC/{id}/canvas/{page_id}',
            '@type': 'sc:Canvas',
            'height': height,
            'images': images,
            'label': f'{page_id}',
            'otherContent': other_content,
            'width': width
        }
    else:
        canvas = {
            '@id': f'https://{domain}/presentation/MISC/{id}/canvas/{page_id}',
            '@type': 'sc:Canvas',
            'height': height,
            'images': images,
            'label': f'{page_id}',
            'width': width
        }
    return canvas


def populate_template(altos, template: dict, metadata: tuple, id: str,
                      ext_id_str: str):
    """Takes an ALTO-XML and additional metadata to populate the IIIF
    template.

    :param alto: etree root object of ALTO-XML.
    :param template: dict of IIIF manifest template.
    :param metadata: pandas row with metadata (itertuples).
    :returns: populated IIIF manifest.
    """
    ext_id = ext_id_str.split('_')
    # Add metadata to template['metadata']
    template['@id'] = f'https://{domain}/presentation/MISC/{id}/manifest'
    template['label'] = metadata.Titolo
    metadata_list = []
    if metadata.Revuo != '' and not pd.isnull(metadata.Revuo):
        metadata_list.append({'label': 'Newspaper', 'value': metadata.Revuo})
    if metadata.Dato != '' and not pd.isnull(metadata.Dato):
        metadata_list.append({'label': 'Date', 'value': metadata.Dato})
    if metadata.Titolo != '' and not pd.isnull(metadata.Titolo):
        metadata_list.append({'label': 'Title', 'value': metadata.Titolo})
    if metadata.Loko != '' and not pd.isnull(metadata.Loko):
        metadata_list.append({'label': 'Location', 'value': metadata.Loko})
    if metadata.Serĉvortoj != '' and not pd.isnull(metadata.Serĉvortoj):
        metadata_list.append({
            'label': 'Keywords',
            'value': metadata.Serĉvortoj
        })
    if metadata.Lingvo != '' and not pd.isnull(metadata.Lingvo):
        metadata_list.append({'label': 'Language', 'value': metadata.Lingvo})
    if metadata._11 != '' and not pd.isnull(metadata._11):
        metadata_list.append({'label': 'AC', 'value': metadata._11})
    template['metadata'] = metadata_list

    # Add canvases to template['sequences']['canvases']
    canvases = []
    if not pd.isnull(metadata.Images):
        for idx, im_path in enumerate(metadata.Images.split(',')):
            page_id = f'page{idx}'
            res_id = str(idx).zfill(8)
            canvases.append(
                gen_canvas(altos[idx], id, page_id, res_id, im_path,
                           ext_id_str))
        template['sequences'][0]['canvases'] = canvases
    else:
        return None
    return template


def get_dfs(csv_path):
    csv_files = [f for f in os.listdir(csv_path) if f.endswith('.csv')]
    # csv_files = [os.path.join(csv_path, elem) for elem in csv_files]
    csv_files.sort()
    id_to_df = {}
    for file in csv_files:
        file_path = os.path.join(csv_path, file)
        df = pd.read_csv(file_path, delimiter=';', encoding='utf-8')
        astor_id = os.path.splitext(file)[0]
        id_to_df[astor_id] = df
    return id_to_df


def pair_tree(string):
    return os.path.join(*[string[i:i + 2] for i in range(0, len(string), 2)])


if __name__ == '__main__':
    id_to_other_ids = get_other_ids()
    id_to_ext_id = id_to_other_ids.copy()
    for key, value in id_to_ext_id.items():
        id_to_ext_id[key] = value[-1]
    template = load_template(filepath='manifest_template.json')
    dfs = get_dfs('../artifacts/mapping_cleaned_links')
    id_it = combinations(list(ascii_lowercase), 5)
    for astor_id, df in dfs.items():
        for row in df.itertuples():
            if row.Images == '' or pd.isnull(row.Images):
                continue
            images = row.Images.split(',')
            article_ids = [
                os.path.splitext(os.path.basename(im))[0] for im in images
            ]
            altos = [
                load_alto(
                    os.path.join('..', 'artifacts', 'predictions', astor_id,
                                 article_id, article_id + '.alto.xml'))
                for article_id in article_ids
            ]
            ext_id_str = id_to_ext_id[astor_id]
            manifest_id = 'esperanto' + ''.join(next(id_it))
            manifest = populate_template(altos, template.copy(), row,
                                         manifest_id, ext_id_str)
            if manifest is None:
                continue
            os.makedirs(f'/home/zenon/manifests/{pair_tree(manifest_id)}',
                        exist_ok=True)
            with open(
                    f'/home/zenon/manifests/{pair_tree(manifest_id)}/{manifest_id}.json',
                    'w',
                    encoding='utf-8') as f:
                json.dump(manifest, f, indent=True, ensure_ascii=False)
