#!/usr/bin/env python3
"""Re-run the share-icloud-photos-android check.

Two parts:

1. public-web: for every public iCloud Shared Album named in dataset.json, ask
   Apple's sharedstreams service (the one the album's web page calls, no account)
   for the photo list, read the size of the largest copy it offers for every
   photo, read the file name of every copy, and download a sample of copies to
   read their real format and pixel size from the bytes.
2. primary-docs: fetch every page named in dataset.json and report, per option
   and property, whether each supporting quote still appears on that page.

    python3 check.py [path/to/dataset.json] [--json out.json]

Without a path it reads dataset.json next to this file, or downloads it from
https://www.viallo.app/blog/share-icloud-photos-android/dataset.json.
"""
import collections
import html
import json
import os
import re
import struct
import sys
import urllib.error
import urllib.request

DOC_UA = ('Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 '
          '(KHTML, like Gecko) Version/18.0 Safari/605.1.15')
# The albums are read the way an Android phone's browser would ask for them.
PHONE_UA = ('Mozilla/5.0 (Linux; Android 14; Pixel 8) AppleWebKit/537.36 '
            '(KHTML, like Gecko) Chrome/130.0 Mobile Safari/537.36')
ACCEPT_IMAGE = 'image/avif,image/webp,image/heic,image/heif,image/apng,*/*'
REMOTE = 'https://www.viallo.app/blog/share-icloud-photos-android/dataset.json'
SAMPLE_PER_ALBUM = 5


def fetch(url, ua=DOC_UA, accept=None):
    headers = {'User-Agent': ua, 'Accept-Language': 'en-US,en;q=0.9'}
    if accept:
        headers['Accept'] = accept
    with urllib.request.urlopen(urllib.request.Request(url, headers=headers), timeout=60) as res:
        return res.read()


def normalise(text):
    text = text.replace('’', "'").replace('‘', "'").replace('“', '"').replace('”', '"')
    return re.sub(r'\s+', ' ', text).strip()


def page_texts(raw):
    raw = re.sub(r'(?is)<(script|style|noscript)[^>]*>.*?</\1>', ' ', raw)
    # A quote can cross an inline link, so compare against both joinings of the tags.
    return [normalise(html.unescape(re.sub(r'(?s)<[^>]+>', sep, raw))) for sep in ('', ' ')]


def stream_post(host, token, path, body):
    req = urllib.request.Request(
        f'https://{host}/{token}/sharedstreams/{path}',
        data=json.dumps(body).encode(),
        headers={'User-Agent': PHONE_UA, 'Content-Type': 'text/plain'},
        method='POST',
    )
    try:
        with urllib.request.urlopen(req, timeout=60) as res:
            return res.status, json.loads(res.read())
    except urllib.error.HTTPError as exc:
        body = exc.read()
        try:
            return exc.code, json.loads(body)
        except ValueError:
            return exc.code, body[:200]


def jpeg_size(data):
    i = 2
    while i + 9 < len(data):
        if data[i] != 0xFF:
            return None
        marker = data[i + 1]
        length = struct.unpack('>H', data[i + 2:i + 4])[0]
        if marker in (0xC0, 0xC1, 0xC2):
            height, width = struct.unpack('>HH', data[i + 5:i + 9])
            return width, height
        i += 2 + length
    return None


def sniff(data):
    if data[:3] == b'\xff\xd8\xff':
        return 'JPEG'
    if data[4:8] == b'ftyp':
        return data[8:12].decode('ascii', 'replace').strip().upper()
    if data[:8] == b'\x89PNG\r\n\x1a\n':
        return 'PNG'
    if data[:4] == b'RIFF' and data[8:12] == b'WEBP':
        return 'WEBP'
    return 'unknown'


def largest(photo):
    return max(photo['derivatives'].values(), key=lambda d: int(d['width']) * int(d['height']))


def measure_album(token):
    # The first host answers 330 with the partition that holds the album.
    host = 'p23-sharedstreams.icloud.com'
    status, stream = stream_post(host, token, 'webstream', {'streamCtag': None})
    if status == 330:
        host = stream['X-Apple-MMe-Host']
        status, stream = stream_post(host, token, 'webstream', {'streamCtag': None})
    if status != 200:
        return {'error': f'webstream HTTP {status}'}
    photos = stream['photos']
    images = [p for p in photos if p.get('mediaAssetType', 'image') == 'image']
    long_edges, short_edges = collections.Counter(), collections.Counter()
    for p in images:
        d = largest(p)
        w, h = int(d['width']), int(d['height'])
        long_edges[max(w, h)] += 1
        short_edges[min(w, h)] += 1
    extensions = collections.Counter()
    urls = {}
    for start in range(0, len(images), 25):
        batch = images[start:start + 25]
        status, assets = stream_post(host, token, 'webasseturls', {'photoGuids': [p['photoGuid'] for p in batch]})
        if status != 200:
            return {'error': f'webasseturls HTTP {status}'}
        for p in batch:
            item = assets['items'].get(largest(p)['checksum'])
            if not item:
                extensions['(no url)'] += 1
                continue
            loc = assets['locations'][item['url_location']]
            url = f"{loc['scheme']}://{loc['hosts'][0]}{item['url_path']}"
            urls[p['photoGuid']] = url
            name = url.split('?')[0].rsplit('/', 1)[-1]
            extensions[name.rsplit('.', 1)[-1].lower() if '.' in name else '(none)'] += 1
    sample = []
    for p in images[:SAMPLE_PER_ALBUM]:
        data = fetch(urls[p['photoGuid']], ua=PHONE_UA, accept=ACCEPT_IMAGE)
        fmt = sniff(data)
        size = jpeg_size(data) if fmt == 'JPEG' else None
        sample.append({'format': fmt, 'bytes': len(data), 'width': size[0] if size else None,
                       'height': size[1] if size else None})
    return {
        'streamName': stream.get('streamName'),
        'images': len(images),
        'videos': len(photos) - len(images),
        'largestCopyLongEdge': {'min': min(long_edges), 'max': max(long_edges)},
        'largestCopyShortEdgeCounts': {str(k): v for k, v in sorted(short_edges.items())},
        'imagesWithLongEdgeAbove2049': sum(v for k, v in long_edges.items() if k > 2049),
        'copyFileExtensions': dict(extensions),
        'sample': sample,
    }


def load_dataset(path):
    if path:
        return json.load(open(path))
    local = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'dataset.json')
    if os.path.exists(local):
        return json.load(open(local))
    return json.loads(fetch(REMOTE))


def main():
    recheck = '--recheck' in sys.argv
    args = [a for a in sys.argv[1:] if a != '--recheck']
    out = None
    if '--json' in args:
        out = args[args.index('--json') + 1]
        del args[args.index('--json'):args.index('--json') + 2]
    data = load_dataset(args[0] if args else None)
    print(f"dataset checkedAt {data['checkedAt']}, method {data['method']}")

    print('\npublic-web: public iCloud Shared Albums')
    measured, changed = {}, 0
    # Re-measuring 720 photos takes minutes, so the weekly recheck reads the pages only.
    for album in ([] if recheck else data['albums']):
        try:
            m = measure_album(album['token'])
        except Exception as exc:
            m = {'error': str(exc)}
        measured[album['token']] = m
        if 'error' in m:
            changed += 1
            print(f"  ERROR   {album['label']}: {m['error']}")
            continue
        recorded = album.get('measured', {})
        same = all(recorded.get(k) == m[k] for k in m)
        changed += 0 if same else 1
        fmts = collections.Counter(s['format'] for s in m['sample'])
        dims = collections.Counter(f"{s['width']}x{s['height']}" for s in m['sample'])
        print(f"  {'SAME' if same else 'CHANGED':<7} {album['label']}: {m['images']} photos, {m['videos']} videos, "
              f"largest copy long edge {m['largestCopyLongEdge']['min']}-{m['largestCopyLongEdge']['max']} px, "
              f"short edge {m['largestCopyShortEdgeCounts']}, long edge above 2049 px on "
              f"{m['imagesWithLongEdgeAbove2049']}, copy names {m['copyFileExtensions']}, "
              f"sample {dict(fmts)} {dict(dims)}")
    if out:
        json.dump(measured, open(out, 'w'), indent=2)

    print('\nprimary-docs: supporting quotes')
    pages = {}
    for url in data['inputs']:
        if 'sharedstreams' in url or 'github.com' in url:
            continue
        try:
            pages[url] = page_texts(fetch(url).decode('utf-8', errors='replace'))
        except Exception as exc:
            pages[url] = exc
    found = missing = unreachable = 0
    for option, props in data['results'].items():
        for prop, cell in props.items():
            if not cell['evidence']:
                print(f"  -       {option} / {prop}: {cell['value']} (no quote)")
                continue
            for ev in cell['evidence']:
                if 'url' not in ev:
                    continue
                page = pages.get(ev['url'])
                if isinstance(page, Exception) or page is None:
                    unreachable += 1
                    status = 'UNREACHABLE'
                elif any(normalise(ev['quote']) in t for t in page):
                    found += 1
                    status = 'FOUND'
                else:
                    missing += 1
                    status = 'MISSING'
                print(f"  {status:<7} {option} / {prop}: {cell['value']}\n          \"{ev['quote']}\" {ev['url']}")
    print(f'\n{0 if recheck else len(data["albums"]) - changed} albums as recorded, {changed} changed or unreachable; '
          f'{found} quotes found, {missing} missing, {unreachable} on unreachable pages')
    return 1 if missing or unreachable or changed else 0


if __name__ == '__main__':
    sys.exit(main())