"""This module provides functions for downloading images with the
signature ``V70/SXXX`` from the `ONB API <https://astor.onb.ac.at/>`_.

"""

import requests
import json
import argparse
import os
import logging
import hashlib
from bs4 import BeautifulSoup
from urllib.parse import urlparse, parse_qs
from pathlib import Path
from dotenv import load_dotenv

load_dotenv('../.env')
ASTOR_API_KEY = os.getenv('ASTOR_API_KEY')

API_URL = 'https://astor.onb.ac.at/discovery/internal/search?query=signature:V70&from=1&rows=200&apikey=' + ASTOR_API_KEY
COLLECTION_URL = 'https://astor.onb.ac.at/access/dip/{}?apikey=' + ASTOR_API_KEY


def get_other_ids():
    """Collect all otherIds for every folder with signature ``V70/SXXX``.

    :returns: Dict of IDs to otherIds List
    :rtype: Dict
    """
    collections = json.loads(requests.get(API_URL).content)
    return {
        doc['id']:
        [otherid for otherid in doc['otherIds'] + [doc['externalId']]]
        for doc in collections['documents']
    }


def get_ids(signatures=False):
    """Collect IDs of all collections with signature ``V70/SXXX``.

    :returns: List of IDs (len: 154)
    :rtype: List[str]

    """
    collections = json.loads(requests.get(API_URL).content)
    total_number = sum(
        [doc['totalImages'] for doc in collections['documents']])
    if signatures:
        return total_number, {
            doc['signature']: doc['id']
            for doc in collections['documents']
        }
    return total_number, [doc['id'] for doc in collections['documents']]


def download_imgs_from_collection(collection_id, output_path, count, limit,
                                  quality) -> int:
    """Takes a COLLECTION_ID and downloads all images associated with
    that ID into OUTPUT_PATH. The JSON for the collection is available
    at f'https://astor.onb.ac.at/access/dip/131ABBA7?apikey={ASTOR_API_KEY}' where
    ``131ABBA7`` is the COLLECTION_ID to download.

    :param collection_id: ID to query and download images from
    :param output_path: Path to save images to
    :param limit: Limit the number of downloaded images per collection
    :param quality: Either lossy or lossless
    :returns: number of images already downloaded (count)
    :rtype: int

    """
    session = requests.Session()

    # Download the JSON response from COLLECTION_URL
    doc = json.loads(session.get(COLLECTION_URL.format(collection_id)).content)
    # Parse the doc to get all files in a list (id_00000000.jp2)
    files = doc['files']

    for _, value in doc['files'].items():
        sha1_server = value["sha1Hash"]
        if count >= limit:
            break

        # Set path of jp2 image
        newpath = os.path.join(output_path, value['aipId'])
        # Make dirs necessary for saving
        os.makedirs(newpath, exist_ok=True)

        count += 1
        # Change fileext from .jp2 to .jpg if lossy desired
        if quality == 'lossy':
            filename = os.path.splitext(
                os.path.join(newpath, value['filename']))[0] + '.jpg'
        else:
            filename = os.path.join(newpath, value['filename'])

        if os.path.exists(filename):
            logging.info(f'{filename} {count}/{limit} already exists.')
            if quality == 'lossless':
                # Check if downloaded file has same hash as specified by the API
                with open(filename, 'r+b') as f:
                    sha1_downloaded = hashlib.sha1(f.read()).hexdigest()
                    if sha1_server != sha1_downloaded:
                        logging.error(
                            f'Local image hash {sha1_downloaded} does not match server hash {sha1_server}'
                        )
                        logging.error(
                            f'Collection ID: {collection_id}, filename: {filename}'
                        )
                    else:
                        logging.info(f'Hashes match!')
            continue

        # imageUrl contains lossy jpegs, fileUrl lossless jp2s
        if quality == 'lossy':
            resp = session.get(value['imageUrl'])
        else:
            resp = session.get(value['fileUrl'])

        if resp.status_code != 200:
            logging.warn(f'Image download failed with code {resp.status_code}')

        # Download image and check local hash against server hash if
        # lossless
        with open(filename, 'w+b') as f:
            f.write(resp.content)
            logging.info(f'Downloaded image {count}/{limit}: {filename}')
            if quality == 'lossless':
                # Only do hash check if lossless
                sha1_downloaded = hashlib.sha1(f.read()).hexdigest()
                if sha1_server != sha1_downloaded:
                    logging.error(
                        f'Local image hash {sha1_downloaded} does not match server hash {sha1_server}'
                    )
                    logging.error(
                        f'Collection ID: {collection_id}, filename: {filename}'
                    )
    return count


def dir_path(string) -> str:
    """Check if STRING is a valid dir path.

    :param string: Dir path as string.
    :returns: Either valid string or exception.
    :rtype: str
    :raises: NotADirectoryError
    """
    if os.path.isdir(string) or not os.path.exists(string):
        return string
    else:
        raise NotADirectoryError(string)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        prog='scrape',
        description='Download Esperanto newspaper images from Astor API.')

    parser.add_argument('-o',
                        '--output',
                        type=dir_path,
                        default=Path('.', 'images'),
                        help='Path of directory to save images to')
    parser.add_argument('-l',
                        '--limit',
                        default=30000,
                        type=int,
                        help='How many images to download')
    parser.add_argument(
        '-q',
        '--quality',
        choices=['lossy', 'lossless'],
        default='lossy',
        help='Whether to download lossy jpegs or (presumably) lossless jp2s')

    parser.add_argument(
        '--log',
        choices=['debug', 'info', 'warning', 'error', 'critical'],
        help='log level',
        default='warning')

    args = parser.parse_args()

    # Set logging level and message format
    numeric_level = getattr(logging, args.log.upper(), None)
    logging.basicConfig(format='%(levelname)s::%(asctime)s::%(message)s',
                        datefmt='%Y-%m-%dT%H:%M:%S',
                        level=numeric_level)

    # Get IDs of all collections (one collection = one signature)
    total_number, ids = get_ids()

    count = 0
    for id in ids:
        count = download_imgs_from_collection(
            id, args.output, count,
            args.limit if args.limit <= total_number else total_number,
            args.quality)
