#!/usr/bin/env python3
#
# Count .deb downloads per distribution on packages.ag-projects.com.
#
# pool/ is shared by all distributions, so the distribution is taken from
# the version suffix autopackager adds (blink_7.1.4trixie_all.deb) and the
# repository from the URL prefix (/debian, /ubuntu, /raspbian).
#
# Apache logs are rotated away after ~2 weeks, so --update keeps per day
# counts of every package in a state file (run it daily from cron). Reports
# read the state file; --from-logs reports on the logs alone.
#
# A day is replaced only when the logs now show at least as many downloads
# for it as stored: the oldest day in the logs is usually cut by rotation
# and must not overwrite the full count stored earlier. Running --update
# again never double counts.
#
# Needs read access to /var/log/apache2 (root, or debian in the adm group).
#

import argparse
import glob
import gzip
import json
import os
import re
import sys
import tempfile

from collections import Counter, defaultdict
from datetime import datetime


DEFAULT_LOGS = '/var/log/apache2/access-packages.ag-projects.com.log*'
DEFAULT_STATE = os.path.expanduser('~debian/packages/download-stats.json')

DEBIAN = ['squeeze', 'wheezy', 'jessie', 'stretch', 'buster', 'bullseye', 'bookworm', 'trixie', 'forky', 'sid', 'experimental']
UBUNTU = ['bionic', 'eoan', 'focal', 'jammy', 'noble', 'resolute']
KNOWN = DEBIAN + UBUNTU

APT_AGENTS = ('Debian APT', 'Ubuntu APT', 'Raspbian APT')

# combined log format
log_re = re.compile(r'^(?P<ip>\S+) \S+ \S+ \[(?P<time>[^\]]+)\] "(?P<method>[A-Z]+) (?P<url>\S+)[^"]*" (?P<status>\d{3}) (?P<size>\S+)(?: "[^"]*" "(?P<agent>[^"]*)")?')
deb_re = re.compile(r'^/(?P<repo>debian|ubuntu|raspbian)/(?:public/)?pool/[^/]+/[^/]+/[^/]+/(?P<package>[^_/]+)_(?P<version>[^_/]+)_(?P<arch>[^_/.]+)\.u?deb$')
dist_re = re.compile(r'(%s)' % '|'.join(sorted(KNOWN, key=len, reverse=True)))

# day record: {'n': {key: count}, 'ips': {package|dist|kind: unique ips}}
# key = package|dist|version|arch|repo|kind, kind = apt or other
SEP = '|'


def distribution(version):
    match = dist_re.search(version.split(':')[-1])
    if match:
        return match.group(1)
    match = re.search(r'[~+.]?([a-z]{4,})\d*$', version)  # codename we don't know yet
    return match.group(1) if match else 'unknown'


def upstream_version(version, dist):
    return version.split(':')[-1].split(dist)[0].rstrip('~+.-') or version


def open_log(path):
    if path.endswith('.gz'):
        return gzip.open(path, 'rt', errors='replace')
    return open(path, errors='replace')


def scan_logs(files, include_partial=False):
    statuses = {'200', '206'} if include_partial else {'200'}
    counts = defaultdict(Counter)
    ips = defaultdict(lambda: defaultdict(set))
    for path in files:
        try:
            log = open_log(path)
        except OSError as e:
            print('cannot read %s: %s' % (path, e), file=sys.stderr)
            continue
        with log:
            for line in log:
                if '/pool/' not in line or 'deb ' not in line:
                    continue
                match = log_re.match(line)
                if not match or match.group('method') != 'GET' or match.group('status') not in statuses:
                    continue
                deb = deb_re.match(match.group('url').split('?')[0])
                if not deb:
                    continue
                day = datetime.strptime(match.group('time').split()[0], '%d/%b/%Y:%H:%M:%S').strftime('%Y-%m-%d')
                dist = distribution(deb.group('version'))
                version = upstream_version(deb.group('version'), dist)
                if deb.group('repo') == 'raspbian':
                    dist = 'raspbian/' + dist
                kind = 'apt' if (match.group('agent') or '').startswith(APT_AGENTS) else 'other'
                package = deb.group('package')
                counts[day][SEP.join((package, dist, version, deb.group('arch'), deb.group('repo'), kind))] += 1
                ips[day][SEP.join((package, dist, kind))].add(match.group('ip'))
                ips[day][SEP.join((package, dist, 'all'))].add(match.group('ip'))
    return {day: {'n': dict(counts[day]), 'ips': {k: len(v) for k, v in ips[day].items()}} for day in counts}


def load_state(path):
    try:
        with open(path) as f:
            return json.load(f)
    except FileNotFoundError:
        return {'days': {}}


def save_state(path, state):
    directory = os.path.dirname(os.path.abspath(path))
    fd, tmp = tempfile.mkstemp(dir=directory, prefix='.download-stats.')
    with os.fdopen(fd, 'w') as f:
        json.dump(state, f, sort_keys=True, separators=(',', ':'))
    os.chmod(tmp, 0o644)
    os.replace(tmp, path)


def merge(state, scanned):
    added = replaced = kept = 0
    for day, record in scanned.items():
        old = state['days'].get(day)
        if old is None:
            added += 1
        elif sum(record['n'].values()) >= sum(old['n'].values()):
            replaced += 1
        else:
            kept += 1
            continue
        state['days'][day] = record
    state['updated'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    return added, replaced, kept


def dist_order(key):
    dist = key.split('/')[-1]
    return (key.startswith('raspbian/'), KNOWN.index(dist) if dist in KNOWN else len(KNOWN), key)


def version_order(version):
    return [(0, int(x), '') if x.isdigit() else (1, 0, x) for x in re.split(r'[.~+-]', version)]


def print_table(title, counter, ipdays=None, order=dist_order):
    if not counter:
        return
    print(title)
    width = max(5, max(len(k) for k in counter))
    for key in sorted(counter, key=order):
        line = '  %-*s %8d' % (width, key, counter[key])
        if ipdays is not None:
            line += '  %8d ip-days' % ipdays[key]
        print(line)
    line = '  %-*s %8d' % (width, 'total', sum(counter.values()))
    if ipdays is not None:
        line += '  %8d ip-days' % sum(ipdays.values())
    print(line)
    print()


def report(days, args, source):
    package_re = re.compile(args.package)
    kinds = {'apt'} if args.apt_only else {'apt', 'other'}
    ip_kind = 'apt' if args.apt_only else 'all'

    per_dist = Counter()
    ipdays = Counter()
    per_repo = Counter()
    per_arch = Counter()
    per_version = defaultdict(Counter)
    per_month = defaultdict(Counter)
    seen = []

    for day in sorted(days):
        if args.since and day < args.since or args.until and day > args.until:
            continue
        record = days[day]
        found = False
        for key, n in record['n'].items():
            package, dist, version, arch, repo, kind = key.split(SEP)
            if kind not in kinds or not package_re.search(package):
                continue
            found = True
            per_dist[dist] += n
            per_repo[repo] += n
            per_arch[arch] += n
            per_version[version][dist] += n
            per_month[day[:7]][dist] += n
        for key, n in record['ips'].items():
            package, dist, kind = key.split(SEP)
            if kind == ip_kind and package_re.search(package):
                ipdays[dist] += n
        if found:
            seen.append(day)

    if not per_dist:
        print('no matching downloads (%s)' % source)
        return

    print('Downloads of %s from %s to %s (%s)\n' % (args.package, seen[0], seen[-1], source))
    print_table('Per distribution:', per_dist, ipdays)
    print_table('Per repository URL:', per_repo, order=str)
    if args.by_arch:
        print_table('Per architecture:', per_arch, order=str)
    if args.by_version:
        for version in sorted(per_version, key=version_order):
            print_table('Version %s:' % version, per_version[version])
    if args.by_month:
        for month in sorted(per_month):
            print_table('Month %s:' % month, per_month[month])


def main():
    parser = argparse.ArgumentParser(description='Count .deb downloads per distribution from the apache logs',
                                     epilog='ip-days: unique IPs per day, summed over the days (the logs are not kept, so no unique count over longer periods)')
    parser.add_argument('logs', nargs='*', help='log files to scan (default: %s)' % DEFAULT_LOGS)
    parser.add_argument('-u', '--update', action='store_true', help='scan the logs into the state file (run daily from cron)')
    parser.add_argument('-l', '--from-logs', action='store_true', help='report on the logs only, without the state file')
    parser.add_argument('-q', '--quiet', action='store_true', help='with --update: no report, only errors')
    parser.add_argument('--state', default=DEFAULT_STATE, help='state file (default: %(default)s)')
    parser.add_argument('-p', '--package', default='^blink$', help='package name regex (default: %(default)s)')
    parser.add_argument('-s', '--since', help='first day to count (YYYY-MM-DD)')
    parser.add_argument('--until', help='last day to count (YYYY-MM-DD)')
    parser.add_argument('-v', '--by-version', action='store_true', help='also break down by upstream version')
    parser.add_argument('-m', '--by-month', action='store_true', help='also break down by month')
    parser.add_argument('-a', '--by-arch', action='store_true', help='also break down by architecture')
    parser.add_argument('--apt-only', action='store_true', help='only count requests made by apt (skips browsers and bots)')
    parser.add_argument('--include-partial', action='store_true', help='with --from-logs: also count 206 partial content responses')
    args = parser.parse_args()

    for date in (args.since, args.until):
        if date:
            datetime.strptime(date, '%Y-%m-%d')

    if args.update and args.from_logs:
        parser.error('--update and --from-logs are mutually exclusive')
    if args.include_partial and not args.from_logs:
        parser.error('--include-partial only works with --from-logs (the state file holds 200 responses only)')

    if args.update or args.from_logs:
        files = args.logs or sorted(glob.glob(DEFAULT_LOGS))
        if not files:
            sys.exit('no log files found matching %s' % DEFAULT_LOGS)
        scanned = scan_logs(files, args.include_partial)
        if args.from_logs:
            report(scanned, args, '%d log files' % len(files))
            return
        state = load_state(args.state)
        added, replaced, kept = merge(state, scanned)
        save_state(args.state, state)
        if args.quiet:
            return
        print('%s: %d days added, %d updated, %d kept (logs had fewer)\n' % (args.state, added, replaced, kept))
    elif args.logs:
        parser.error('log files given without --update or --from-logs')

    if not os.path.exists(args.state):
        sys.exit('%s does not exist, run with --update first (or use --from-logs)' % args.state)
    state = load_state(args.state)
    report(state['days'], args, 'state file, updated %s' % state.get('updated', '?'))


if __name__ == '__main__':
    main()