#!/usr/bin/env python3
"""Crawl the home page of the top N Tranco domains and record Open Graph / Twitter card tags.

Usage: python3 -I og_crawl.py <tranco.csv> <N> <out.csv>
Standard library only. One GET per home page, one GET per og:image URL.
"""
import csv
import os
import struct
import sys
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from html.parser import HTMLParser
from urllib.parse import urljoin

UA = 'Mozilla/5.0 (compatible; RenderScreenshot-OG-Study/1.0; +https://renderscreenshot.com/blog/state-of-og-images)'
HTML_LIMIT = 3 * 1024 * 1024
IMG_LIMIT = 20 * 1024 * 1024
TIMEOUT = 15
WALL_LIMIT = 900  # seconds for the whole crawl
WANTED = {
    'og:image', 'og:image:url', 'og:image:secure_url', 'og:image:width', 'og:image:height',
    'og:image:type', 'og:image:alt', 'og:title', 'og:description', 'og:type', 'og:url', 'og:site_name',
    'twitter:card', 'twitter:image', 'twitter:image:src', 'twitter:site',
}


class MetaParser(HTMLParser):
    def __init__(self):
        super().__init__(convert_charrefs=True)
        self.meta = {}
        self.done = False

    def handle_starttag(self, tag, attrs):
        if self.done:
            return
        if tag == 'body':
            self.done = True
            return
        if tag != 'meta':
            return
        a = {k.lower(): (v or '') for k, v in attrs}
        key = (a.get('property') or a.get('name') or '').strip().lower()
        if key in WANTED and key not in self.meta:
            self.meta[key] = a.get('content', '').strip()


def fetch(url, limit):
    req = urllib.request.Request(url, headers={'User-Agent': UA, 'Accept': '*/*', 'Accept-Language': 'en'})
    deadline = time.monotonic() + 3 * TIMEOUT  # stop slow-drip responses
    with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
        chunks, size = [], 0
        while size <= limit:
            if time.monotonic() > deadline:
                raise TimeoutError('body deadline')
            chunk = r.read(65536)
            if not chunk:
                break
            chunks.append(chunk)
            size += len(chunk)
        return r.status, r.geturl(), r.headers.get('Content-Type', ''), b''.join(chunks)


def sniff(b):
    """Return (format, width, height) from image bytes, or (format, None, None)."""
    try:
        if b[:8] == b'\x89PNG\r\n\x1a\n':
            w, h = struct.unpack('>II', b[16:24])
            return 'png', w, h
        if b[:6] in (b'GIF87a', b'GIF89a'):
            w, h = struct.unpack('<HH', b[6:10])
            return 'gif', w, h
        if b[:4] == b'RIFF' and b[8:12] == b'WEBP':
            chunk = b[12:16]
            if chunk == b'VP8 ':
                w, h = struct.unpack('<HH', b[26:30])
                return 'webp', w & 0x3FFF, h & 0x3FFF
            if chunk == b'VP8L':
                bits = int.from_bytes(b[21:25], 'little')
                return 'webp', (bits & 0x3FFF) + 1, ((bits >> 14) & 0x3FFF) + 1
            if chunk == b'VP8X':
                w = int.from_bytes(b[24:27], 'little') + 1
                h = int.from_bytes(b[27:30], 'little') + 1
                return 'webp', w, h
            return 'webp', None, None
        if b[:2] == b'\xff\xd8':
            i = 2
            while i < len(b) - 9:
                if b[i] != 0xFF:
                    i += 1
                    continue
                marker = b[i + 1]
                if marker in (0xD8, 0x01) or 0xD0 <= marker <= 0xD7:
                    i += 2
                    continue
                seg = struct.unpack('>H', b[i + 2:i + 4])[0]
                if marker in (0xC0, 0xC1, 0xC2, 0xC3, 0xC5, 0xC6, 0xC7, 0xC9, 0xCA, 0xCB, 0xCD, 0xCE, 0xCF):
                    h, w = struct.unpack('>HH', b[i + 5:i + 9])
                    return 'jpeg', w, h
                i += 2 + seg
            return 'jpeg', None, None
        if b[4:8] == b'ftyp' and b[8:12] in (b'avif', b'avis'):
            j = b.find(b'ispe')
            if j > 0:
                w, h = struct.unpack('>II', b[j + 8:j + 16])
                return 'avif', w, h
            return 'avif', None, None
        head = b[:512].lstrip().lower()
        if head.startswith(b'<svg') or (head.startswith(b'<?xml') and b'<svg' in b[:2048].lower()):
            return 'svg', None, None
        if b[:4] == b'\x00\x00\x01\x00':
            return 'ico', None, None
        if head.startswith(b'<!doctype html') or head.startswith(b'<html'):
            return 'html', None, None
    except Exception:
        pass
    return 'unknown', None, None


def crawl(rank, domain):
    row = {'rank': rank, 'domain': domain, 'home_status': '', 'home_error': '', 'final_url': '', 'is_html': 0}
    body = None
    for scheme_host in (f'https://{domain}/', f'https://www.{domain}/'):
        try:
            status, final, ctype, body = fetch(scheme_host, HTML_LIMIT)
            row.update(home_status=status, final_url=final, home_error='')
            row['is_html'] = int('html' in ctype.lower())
            break
        except urllib.error.HTTPError as e:
            row.update(home_status=e.code, home_error='http')
        except Exception as e:
            row.update(home_error=type(e).__name__)
    if body is None or not row['is_html']:
        return row
    p = MetaParser()
    try:
        p.feed(body[:HTML_LIMIT].decode('utf-8', errors='replace'))
    except Exception:
        pass
    m = p.meta
    for k in sorted(WANTED):
        row[k] = m.get(k, '')[:500]
    img = m.get('og:image') or m.get('og:image:url') or m.get('og:image:secure_url')
    row['has_og_image'] = int(bool(img))
    row['has_twitter_card'] = int(bool(m.get('twitter:card')))
    row['img_url'] = ''
    if not img:
        return row
    absolute = urljoin(row['final_url'], img)
    row['img_relative'] = int(not img.lower().startswith(('http://', 'https://', '//')))
    row['img_http_only'] = int(absolute.lower().startswith('http://'))
    row['img_url'] = absolute[:500]
    try:
        status, final, ctype, data = fetch(absolute, IMG_LIMIT)
        fmt, w, h = sniff(data)
        row.update(img_status=status, img_ctype=ctype.split(';')[0].strip().lower(), img_bytes=len(data),
                   img_format=fmt, img_width=w or '', img_height=h or '', img_error='')
    except urllib.error.HTTPError as e:
        row.update(img_status=e.code, img_error='http')
    except Exception as e:
        row.update(img_error=type(e).__name__)
    return row


FIELDS = ['rank', 'domain', 'home_status', 'home_error', 'final_url', 'is_html', 'has_og_image', 'has_twitter_card',
          'img_url', 'img_relative', 'img_http_only', 'img_status', 'img_error', 'img_ctype', 'img_format',
          'img_bytes', 'img_width', 'img_height'] + sorted(WANTED)


def main():
    src, n, out = sys.argv[1], int(sys.argv[2]), sys.argv[3]
    domains = []
    with open(src) as f:
        for line in csv.reader(f):
            domains.append((int(line[0]), line[1]))
            if len(domains) >= n:
                break
    started = time.time()
    rows = []
    ex = ThreadPoolExecutor(max_workers=96)
    futs = {ex.submit(crawl, r, d): (r, d) for r, d in domains}
    try:
        # Hard stop: a hung DNS lookup is not covered by the socket timeout.
        for i, fu in enumerate(as_completed(futs, timeout=WALL_LIMIT), 1):
            rows.append(fu.result())
            if i % 500 == 0:
                print(f'{i}/{n} {time.time() - started:.0f}s', flush=True)
    except TimeoutError:
        for fu, (r, d) in futs.items():
            if not fu.done():
                rows.append({'rank': r, 'domain': d, 'home_error': 'WallClockLimit', 'is_html': 0})
    rows.sort(key=lambda r: r['rank'])
    with open(out, 'w', newline='') as f:
        w = csv.DictWriter(f, fieldnames=FIELDS, extrasaction='ignore')
        w.writeheader()
        w.writerows(rows)
    print(f'done {len(rows)} rows in {time.time() - started:.0f}s', flush=True)
    os._exit(0)  # do not wait for hung worker threads


if __name__ == '__main__':
    main()