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


from __future__ import annotations


__all__ = ['get_warc_metadata', 'QueryWARC']

# Cell
#nbdev_comment from __future__ import annotations
from typing import Union
from pathlib import Path
from collections.abc import Iterable

from warcio.recordloader import ArcWarcRecord
import warcio

# Cell

def get_warc_metadata(record: ArcWarcRecord) -> dict[str, str]:
    return {
        'url': record.rec_headers.get_header('WARC-Target-URI'),
        # Normalise timestamp?
        'timestamp': record.rec_headers.get_header('WARC-Date'),
        'mime': record.http_headers.get_header('Content-Type').split(';')[0],
        'status':  record.http_headers.get_statuscode(),
    }

class QueryWARC:
    def __init__(self, path: Union[str, Path]) -> None:
        self.path = path

    def query(self) -> list[dict[str, str]]:
        with open(self.path, 'rb') as f:
            archive = warcio.ArchiveIterator(f)
            for record in archive:
                if record.rec_type != 'response':
                    continue
                metadata = get_warc_metadata(record)
                metadata['offset'] = str(archive.get_record_offset())
                metadata['filename'] = self.path
                yield metadata

    def fetch_one(self, record: dict[str, str]) -> bytes:
        with open(record['filename'], 'rb') as f:
            f.seek(int(record['offset']))
            record = next(warcio.ArchiveIterator(f))
            return record.content_stream().read()

    def fetch(self, records: Iterable[dict[str, str]]) -> list[bytes]:
        # Inefficient to keep opening/closing file
        return [self.fetch_one(record) for record in records]