# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01_cdx.ipynb (unless otherwise specified).


from __future__ import annotations # For Python <3.9


__all__ = ['header_and_rows_to_dict', 'mimetypes_to_regex', 'query_wayback_cdx', 'IA_CDX_URL', 'CaptureIndexRecord',
           'fetch_internet_archive_content', 'make_session', 'WaybackResult', 'WaybackMachine', 'get_cc_indexes',
           'jsonl_loads', 'query_cc_cdx_num_pages', 'query_cc_cdx_page', 'fetch_cc', 'CC_DATA_URL', 'CommonCrawl']

# Cell
# Typing
#nbdev_comment from __future__ import annotations # For Python <3.9
from typing import Any, Optional, Union
from collections.abc import Iterable

import requests
from requests.sessions import Session
import json

# Cell
def header_and_rows_to_dict(rows: Iterable[list[Any]]) -> list[dict[Any, Any]]:
    header = None
    data = []
    for row in rows:
        if header is None:
            header = row
        else:
            assert len(row) == len(header), "Row should be same length as header"
            data.append(dict(zip(header, row)))
    return data

# Cell
IA_CDX_URL = 'http://web.archive.org/cdx/search/cdx'

# This could be more precise
CaptureIndexRecord = dict

def mimetypes_to_regex(mime: list[str], prefix='mimetype:') -> str:
    return prefix + '|'.join('(' + s.replace('*', '.*') + ')' for s in mime)

def query_wayback_cdx(url: str, start: Optional[str], end: Optional[str],
                      status_ok: bool = True,
                      mime: Optional[Union[str, Iterable[str]]] = None,
                      limit: Optional[int] = None, offset: Optional[int] = None,
                      session: Optional[Session] = None) -> list[CaptureIndexRecord]:
    """Get references to Wayback Machine Captures for url.

    Queries the Internet Archive Capture Index (CDX) for url.

    Arguments:
      * start: Minimum date in format YYYYmmddHHMMSS (or any substring) inclusive
      * end: Maximum date in format YYYYmmddHHMMSS (or any substring) inclusive
      * status_ok: Only return those with a HTTP status 200
      * mime: Filter on mimetypes, '*' is a wildcard (e.g. 'image/*')
      * limit: Only return first limit records
      * offset: Skip the first offset records, combine with limit
      * session: Session to use when making requests
    Filters results between start and end inclusive, in format YYYYmmddHHMMSS or any substring
    (e.g. start="202001", end="202001" will get all captures in January 2020)
    """
    if session is None:
        session = requests

    params = {'url': url,
              'output': 'json',
              'from': start,
              'to': end,
              'limit': limit,
              'offset': offset}

    filter = []
    if status_ok:
        filter.append('statuscode:200')
    if mime:
        # Turn our list of mimetypes into a regex
        if isinstance(mime, str):
            mime = [mime]
        filter.append(mimetypes_to_regex(mime))
    params['filter'] = filter

    params = {k:v for k,v in params.items() if v}
    response = session.get(IA_CDX_URL, params=params)
    response.raise_for_status()
    return header_and_rows_to_dict(response.json())

# Cell
def fetch_internet_archive_content(timestamp: str, url: str,
                                   session: Optional[Session] = None) -> bytes:
    if session is None:
        session = requests

    url = f'http://web.archive.org/web/{timestamp}/{url}'
    response = session.get(url)
    # Sometimes Internet Archive deletes records
    if response.status_code == 404:
        logging.warning(f'Missing {url}')
        return None
    response.raise_for_status()
    return response.content

# Cell
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry

def make_session(pool_maxsize):
    retry_strategy =  Retry(total=5, backoff_factor=1, status_forcelist=set([504, 500]))
    adapter = HTTPAdapter(max_retries=retry_strategy, pool_maxsize=pool_maxsize, pool_block=True)
    session = requests.Session()
    session.mount('http://', adapter)
    session.mount('https://', adapter)
    return session

# Cell
from joblib import delayed, Memory, Parallel

class WaybackResult:
    def __init__(self, params, session=None):
        self.params = params
        self.session = session
        self._query = self.memory.cache(query_wayback_cdx, ignore=['session'])

    def fetch(self, force=False):
        fetch = self._fetch.call if force else self._fetch
        result = fetch(record['timestamp'], record['original'], session=self.session)
        return result[0] if force else result


class WaybackMachine:
    def __init__(self, location: Optional[str],
                 session: Union[bool, Session] = True, threads: int=8, verbose=0) -> None:
        self.threads = threads
        self.memory = Memory(location, verbose=verbose)

        if session is True:
            session = make_session(pool_maxsize=threads)
        if not session:
            session = None
        self.session = session

        self._query = self.memory.cache(query_wayback_cdx, ignore=['session'])
        self._fetch = self.memory.cache(fetch_internet_archive_content, ignore=['session'])

    def query(self, url: str, start: Optional[str], end: Optional[str],
                      status_ok: bool = True, mime: Optional[Union[str, Iterable[str]]] = None,
                      limit: Optional[int] = None, offset: Optional[int] = None, force=False) -> List[CaptureIndexRecord]:
        query = self._query.call if force else self._query
        result = query(url, start, end, status_ok, mime, limit, offset, session=self.session)
        return result[0] if force else result

    def fetch_one(self, record, force=False) -> bytes:
        fetch = self._fetch.call if force else self._fetch
        result = fetch(record['timestamp'], record['original'], session=self.session)
        return result[0] if force else result

    def fetch(self, records, force=False, threads=None) -> bytes:
        fetch = self._fetch.call if force else self._fetch
        result = Parallel(threads or self.threads, prefer='threads')(delayed(self.fetch_one)(record, force) for record in records)
        return result

# Cell
from functools import lru_cache
@lru_cache(maxsize=None)
def get_cc_indexes() -> List[Dict[str, str]]:
    response = requests.get("https://index.commoncrawl.org/collinfo.json")
    response.raise_for_status()
    return response.json()

# Cell
import json

def jsonl_loads(jsonl):
    return [json.loads(line) for line in jsonl.splitlines()]

# Cell
def query_cc_cdx_num_pages(api: str,  url: str,
                           session: Optional[Session] = None) -> int:
    if session is None:
        session = requests

    response = session.get(api, params={'url': url, 'output': 'json', 'showNumPages': True})
    response.raise_for_status()
    data = response.json()
    return data["pages"]

def query_cc_cdx_page(
                 api: str, url: str, page: int,
                 start: Optional[str] = None, end: Optional[str] = None,
                 status_ok: bool = True, mime: Optional[Union[str, Iterable[str]]] = None,
                 limit: Optional[int] = None, offset: Optional[int] = None,
                 session: Optional[Session] = None) -> List[CaptureIndexRecord]:
    """Get references to Common Crawl Captures for url.

    Queries the Common Crawl Capture Index (CDX) for url.

    Filters:
      * api: API endpoint to use (e.g. 'https://index.commoncrawl.org/CC-MAIN-2021-43-index')
      * start: Minimum date in format YYYYmmddHHMMSS (or any substring) inclusive
      * end: Maximum date in format YYYYmmddHHMMSS (or any substring) inclusive
      * status_ok: Only return those with a HTTP status 200
      * mime: Filter on mimetypes, '*' is a wildcard (e.g. 'image/*')
      * limit: Only return first limit records
      * offset: Skip the first offset records, combine with limit
      * session: Session to use when making requests
    Filters results between start and end inclusive, in format YYYYmmddHHMMSS or any substring
    (e.g. start="202001", end="202001" will get all captures in January 2020)
    """
    if session is None:
        session = requests

    params = {'url': url,
              'page': page,
              'output': 'json',
              'page': page,
              'from': start,
              'to': end,
              'limit': limit,
              'offset': offset}

    filter = []
    if status_ok:
        # N.B. Different to IA
        filter.append('=status:200')
    if mime:
        if isinstance(mime, str):
            mime = [mime]
        # N.B. Different to IA
        filter.append(mimetypes_to_regex(mime, prefix='~mime:'))
    params['filter'] = filter

    params = {k:v for k,v in params.items() if v}
    response = session.get(api, params=params)
    response.raise_for_status()
    return jsonl_loads(response.content)

# Cell
from warcio import ArchiveIterator
from io import BytesIO

CC_DATA_URL = "https://commoncrawl.s3.amazonaws.com/"
def fetch_cc(filename: str, offset: int, length: int, session: Optional[Session] = None) -> bytes:
    if session is None:
        session = requests
    data_url = CC_DATA_URL + filename
    start_byte = int(offset)
    end_byte = start_byte + int(length)
    headers = {"Range": f"bytes={start_byte}-{end_byte}"}
    r = session.get(data_url, headers=headers)
    r.raise_for_status()

    # Decord WARC
    response_content = r.content
    archive = ArchiveIterator(BytesIO(response_content))
    record = next(archive)
    content = record.content_stream().read()

    # Archive should have just 1 record
    assert not any(True for _ in archive), "Expected 1 result in archive"

    return content

# Cell
from joblib import delayed, Memory, Parallel
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry

class CommonCrawl:
    def __init__(self, location: Optional[str],
                 session: Union[bool, Session] = True, threads: int=64, verbose=0) -> None:
        self.threads = threads
        self.memory = Memory(location, verbose=verbose)

        if session is True:
            session = make_session(pool_maxsize=threads)
        if not session:
            session = None
        self.session = session

        self._query_pages = self.memory.cache(query_cc_cdx_num_pages, ignore=['session'])
        self._query = self.memory.cache(query_cc_cdx_page, ignore=['session'])
        self._fetch = self.memory.cache(fetch_cc, ignore=['session'])

    @property
    def cdx_apis(self) -> Dict[str, str]:
        return {x['id']: x['cdx-api'] for x in get_cc_indexes()}

    # Note: don't use the default number of threads since we're hitting a single CC API
    # Todo several CC APIs
    def query(self, api: str, url: str,
              start: Optional[str] = None, end: Optional[str] = None,
              status_ok: bool = True, mime: Optional[Union[str, Iterable[str]]] = None,
              limit: Optional[int] = None, offset: Optional[int] = None, force=False,
              threads: int=4) -> List[CaptureIndexRecord]:
        # lookup by key
        api = self.cdx_apis.get(api, api)
        query = self._query.call if force else self._query

        # Wrap this up?
        query_pages = self._query_pages.call if force else self._query_pages
        num_pages = query_pages(api, url, session=self.session)
        num_pages = num_pages[0] if force else num_pages

        pages_results = Parallel(threads, prefer='threads')(delayed(query)(api, url, page, start, end, status_ok, mime, limit, offset, session=self.session) for page in range(num_pages))
        if force:
            pages_results = [r[0] for r in pages_results]
        return [result for page in pages_results for result in page]

    def fetch_one(self, record, force=False) -> bytes:
        fetch = self._fetch.call if force else self._fetch
        result = fetch(record['filename'], record['offset'], record['length'], session=self.session)
        return result[0] if force else result

    def fetch(self, records, force=False, threads=None) -> List[bytes]:
        fetch = self._fetch.call if force else self._fetch
        result = Parallel(threads or self.threads, prefer='threads')(delayed(self.fetch_one)(record, force) for record in records)
        return result