#!/usr/bin/env python3
"""Download every visible catalogue title/year and verify count and SHA-256."""
import argparse
import hashlib
import json
from datetime import datetime, timezone
from pathlib import Path
import time
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

URL = 'https://psilocybin-research.com/catalogue.php?format=txt'


def verify(raw, headers):
    if headers.get('Content-Type', '').split(';')[0].strip() != 'text/plain':
        raise ValueError('Expected UTF-8 plain text; no output saved')
    count = int(headers['X-Publication-Tracker-Export-Count'])
    total = int(headers['X-Publication-Tracker-Total'])
    digest = hashlib.sha256(raw).hexdigest()
    if headers.get('X-Publication-Tracker-Complete') != 'true' or count != total:
        raise ValueError('Server did not confirm a complete catalogue; no output saved')
    if digest != headers.get('X-Publication-Tracker-SHA256'):
        raise ValueError('Response checksum mismatch; no output saved')
    text = raw.decode('utf-8')
    if len(text.splitlines()) != count:
        raise ValueError('TXT line count does not match the record count; no output saved')
    return {
        'source_url': URL,
        'retrieved_at_utc': datetime.now(timezone.utc).isoformat(),
        'generated_at_utc': headers['X-Publication-Tracker-Generated-At'],
        'record_count': count,
        'total_records': total,
        'complete': True,
        'sha256': digest,
        'interpretation': 'One line per catalogue record, including preprints and trial registrations; n.d. means an unrecorded year. Records are not necessarily independent studies.',
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, default=Path('psilocybin_year_title.txt'))
    args = parser.parse_args()
    manifest = args.output.with_suffix(args.output.suffix + '.manifest.json')
    if args.output.exists() or manifest.exists():
        parser.error('Output or manifest already exists; choose a new filename')
    request = Request(URL, headers={'User-Agent': 'PsilocybinCatalogueClient/1.0', 'Accept': 'text/plain'})
    response = None
    for attempt in range(3):
        try:
            response = urlopen(request, timeout=60)
            break
        except HTTPError as error:
            if error.code not in (429, 503) or attempt == 2:
                raise
            delay = error.headers.get('Retry-After', '10')
            time.sleep(min(60, max(1, int(delay) if delay.isdigit() else 10)))
    with response:
        raw = response.read(32 * 1024 * 1024 + 1)
        if len(raw) > 32 * 1024 * 1024:
            raise ValueError('Response exceeded the download bound; no output saved')
        metadata = verify(raw, response.headers)
    created = []
    try:
        with args.output.open('xb') as target:
            created.append(args.output)
            with manifest.open('xb') as meta:
                created.append(manifest)
                target.write(raw)
                meta.write((json.dumps(metadata, indent=2) + '\n').encode('utf-8'))
    except BaseException:
        for path in created:
            path.unlink(missing_ok=True)
        raise
    print(f"Saved and verified {metadata['record_count']} records: {args.output}")
    print(f'Keep the verification metadata: {manifest}')


if __name__ == '__main__':
    try:
        main()
    except (URLError, ValueError, KeyError, OSError) as error:
        raise SystemExit(f'Download not completed: {error}. If network access is blocked, download the TXT in a browser and attach it to your AI session.')