#!/usr/bin/env python3
# Expects traefik json access logs either by passing a log file as the first argument, or via stdin

import json
import re
import sys
import argparse
from datetime import datetime
from collections import Counter
from typing import List, Tuple, Optional
from dataclasses import dataclass

from tabulate import tabulate
import requests_cache
from netaddr import IPSet, IPAddress


@dataclass
class LogEntry:
    ClientHost: str
    DownstreamStatus: int
    RequestHost: str
    RequestMethod: str
    RequestPath: str
    RequestProtocol: str
    UserAgent: str
    time: str
    client_info: "Optional[ClientInfo]"


@dataclass
class UserAgent:
    pacman_version: Tuple[int, int, int]
    windows_version: Tuple[int, int]
    build_number: int
    cygwin_arch: str
    host_arch: str
    libalpm_version: Tuple[int, int, int]


@dataclass
class ClientInfo:
    pacman_version: str
    windows_edition: str
    user_agent: UserAgent
    ci: str


def is_valid_user_agent(ua: str) -> bool:
    return ua.startswith("pacman") and "_NT-" in ua


def parse_user_agent(ua: str) -> UserAgent:
    assert is_valid_user_agent(ua)
    m = re.match(r"pacman/([^\s]+) \(([^-]+)-([^-]+)-?([0-9]+|)-?(.+|) ([^)]+)\) libalpm/(.+)", ua)
    assert m is not None
    pacman_version, _, windows_version, build_number, extra, arch, alpm_version = m.groups()

    pacman_version = tuple(map(int, pacman_version.split("-")[0].split(".")))
    windows_version = tuple(map(int, windows_version.split(".")))
    build_number = int(build_number or "-1")
    cygwin_arch = arch
    if extra in ("WOW64", "WOW") and cygwin_arch == "i686":
        host_arch = "x86_64"
    elif extra == "ARM64":
        host_arch = "aarch64"
    elif extra:
        host_arch = "???"
    else:
        host_arch = cygwin_arch
    alpm_version = tuple(map(int, alpm_version.split(".")))

    return UserAgent(
        pacman_version,
        windows_version,
        build_number,
        cygwin_arch,
        host_arch,
        alpm_version
    )


def test_parse_user_agent():
    ua = "pacman/6.0.1 (MSYS_NT-10.0-19042 x86_64) libalpm/13.0.1"
    assert parse_user_agent(ua).pacman_version == (6, 0, 1)
    ua = "pacman/6.0.1 (MSYS_NT-10.0-19042-WOW64 i686) libalpm/13.0.1"
    assert parse_user_agent(ua).cygwin_arch == "i686"
    assert parse_user_agent(ua).host_arch == "x86_64"
    assert parse_user_agent(ua).build_number == 19042
    ua = "pacman/5.2.1 (MSYS_NT-6.1-7601-WOW64 i686) libalpm/12.0.1"
    assert parse_user_agent(ua).cygwin_arch == "i686"
    assert parse_user_agent(ua).host_arch == "x86_64"
    assert parse_user_agent(ua).build_number == 7601
    ua = "pacman/5.1.0 (MINGW64_NT-10.0 x86_64) libalpm/11.0.0"
    assert parse_user_agent(ua).build_number == -1
    assert parse_user_agent(ua).cygwin_arch == "x86_64"
    assert parse_user_agent(ua).host_arch == "x86_64"
    ua = "pacman/5.0.1 (UCRT64_NT-10.0-WOW i686) libalpm/10.0.1"
    assert parse_user_agent(ua).build_number == -1
    assert parse_user_agent(ua).host_arch == "x86_64"
    assert parse_user_agent(ua).cygwin_arch == "i686"
    ua = "pacman/4.2.1-313-g5535-dirty (MINGW64_NT-6.1 x86_64) libalpm/9.0.1"
    assert parse_user_agent(ua).pacman_version == (4, 2, 1)
    ua = "pacman/6.0.1 (MSYS_NT-6.0-6002 x86_64) libalpm/13.0.1"
    assert parse_user_agent(ua).windows_version == (6, 0)
    ua = "pacman/6.1.0 (MSYS_NT-10.0-22631-ARM64 x86_64) libalpm/14.0.0"
    assert parse_user_agent(ua).host_arch == "aarch64"
    assert parse_user_agent(ua).cygwin_arch == "x86_64"
    ua = "pacman/6.1.0 (MSYS_NT-10.0-22631-ARM64 i686) libalpm/14.0.0"
    assert parse_user_agent(ua).host_arch == "aarch64"
    assert parse_user_agent(ua).cygwin_arch == "i686"


def test_get_windows_edition():
    ua = "pacman/6.0.1 (MSYS_NT-6.0-6002 x86_64) libalpm/13.0.1"
    assert get_windows_edition(parse_user_agent(ua)) == "Vista"
    ua = "pacman/4.2.1-463-gbb493-dirty (MINGW32_NT-5.0 i686) libalpm/10.0.0"
    assert get_windows_edition(parse_user_agent(ua)) == "2000"


def get_windows_edition(ua: UserAgent):
    if ua.windows_version >= (10, 0):
        if ua.build_number >= 22000:
            return "11"
        else:
            return "10"
    elif ua.windows_version == (5, 0):
        return "2000"
    elif ua.windows_version == (5, 1):
        return "XP"
    elif ua.windows_version == (5, 2):
        return "XP"
    elif ua.windows_version == (6, 0):
        return "Vista"
    elif ua.windows_version == (6, 1):
        return "7"
    elif ua.windows_version == (6, 2):
        return "8"
    elif ua.windows_version == (6, 3):
        return "8.1"
    else:
        assert 0, ua


def get_ci_networks():
    session = requests_cache.CachedSession(
        'msys2-logstats-ci', expire_after=360, use_cache_dir=True)
    r = session.get("https://api.github.com/meta")
    r.raise_for_status()
    gha = IPSet(r.json()["actions"])

    r = session.get("https://www.appveyor.com/ips.json")
    r.raise_for_status()
    appveyor = IPSet(r.json())

    r = session.get("https://www.gstatic.com/ipranges/cloud.json")
    r.raise_for_status()
    prefixes = set()
    for prefix in r.json()["prefixes"]:
        prefixes.add(prefix.get("ipv4Prefix", prefix.get("ipv6Prefix")))
    gcp = IPSet(prefixes)

    r = session.get("https://ip-ranges.amazonaws.com/ip-ranges.json")
    r.raise_for_status()
    prefixes = set()
    for prefix in r.json()["prefixes"]:
        prefixes.add(prefix["ip_prefix"])
    for prefix in r.json()["ipv6_prefixes"]:
        prefixes.add(prefix["ipv6_prefix"])
    aws = IPSet(prefixes)

    r = session.get(
        "https://www.microsoft.com/en-us/download/details.aspx?id=56519",
        headers={
            'User-Agent':
                'Mozilla/5.0 (X11; Linux x86_64; rv:109.0) Gecko/20100101 Firefox/112.0'
        }
    )
    r.raise_for_status()
    match = re.search("(https://download.microsoft.com/.*?/ServiceTags_Public_.*?\\.json)", r.text)
    r = session.get(match.groups(0)[0])
    r.raise_for_status()
    prefixes = set()
    for value in r.json()["values"]:
        for nw in value["properties"]["addressPrefixes"]:
            prefixes.add(nw)
    azure = IPSet(prefixes)

    return {"GHA": gha, "APPV": appveyor, "GCP": gcp, "AWS": aws, "AZ": azure}


def get_repo_for_path(path: str) -> str:
    repo = path.rsplit("/", 1)[0].lstrip("/")
    if repo == "mingw/i686":
        repo = "mingw/mingw32"
    elif repo == "mingw/x86_64":
        repo = "mingw/mingw64"

    return repo


def get_type_for_path(path: str) -> str:
    if path.endswith(".db") or ".db." in path:
        return "db"
    elif path.endswith(".files") or ".files." in path:
        return "db"
    else:
        return "pkg"


def print_repos(entries, show_ci):
    for request_type in ["pkg", "db"]:
        type_requests = [e for e in entries if get_type_for_path(e.RequestPath) == request_type]
        table = []
        for (repo, type_, ci), count in Counter([
                (get_repo_for_path(e.RequestPath), get_type_for_path(e.RequestPath),
                 e.client_info.ci) for e in type_requests]).most_common():
            pcnt = count / len(type_requests) * 100
            line = [repo, type_, ci, f"{pcnt:.2f}%", f"{count}"]
            if not show_ci:
                line.pop(2)
            table.append(line)
        headers = ["Repo", "Type", "CI", "% Requests", "Requests"]
        if not show_ci:
            headers.pop(2)
        print()
        print(tabulate(table, headers, stralign="right", numalign="right"))


def print_windows_major(clients, entries, show_ci):
    per_request = {}
    for (edition, ci), count in Counter([(e.client_info.windows_edition, e.client_info.ci) for e in entries]).most_common():
        per_request[(edition, ci)] = count
    table = []
    for (edition, ci), count_clients in Counter([(u.windows_edition, u.ci) for u in clients]).most_common():
        pcnt_clients = count_clients / len(clients) * 100
        count_req = per_request[(edition, ci)]
        pcnt_req = count_req / len(entries) * 100
        line = [edition, ci, f"{pcnt_clients:.2f}%", f"{count_clients}", f"{pcnt_req:.2f}%", f"{count_req}"]
        if not show_ci:
            line.pop(1)
        table.append(line)
    headers = ["Windows", "CI", "% Clients", "Clients", "% Requests", "Requests"]
    if not show_ci:
        headers.pop(1)
    print()
    print(tabulate(table, headers, stralign="right", numalign="right"))


def print_ci_systems(clients, entries):
    per_request = {}
    for ci, count in Counter([e.client_info.ci for e in entries]).most_common():
        per_request[ci] = count
    table = []
    for ci, count_clients in Counter([u.ci for u in clients]).most_common():
        pcnt_clients = count_clients / len(clients) * 100
        count_req = per_request[ci]
        pcnt_req = count_req / len(entries) * 100
        line = [ci, f"{pcnt_clients:.2f}%", f"{count_clients}", f"{pcnt_req:.2f}%", f"{count_req}"]
        table.append(line)
    headers = ["CI", "% Clients", "Clients", "% Requests", "Requests"]
    print()
    print(tabulate(table, headers, stralign="right", numalign="right"))


def print_windows_version_details(clients, show_ci):
    table = []
    for (windows_version, build_number, ci), count in Counter(
            [(u.user_agent.windows_version, u.user_agent.build_number, u.ci) for u in clients]).most_common():
        pcnt = count / len(clients) * 100
        line = [".".join(map(str, windows_version)), build_number, ci, f"{pcnt:.2f}%", f"{count}"]
        if not show_ci:
            line.pop(2)
        table.append(line)
    headers = ["Win Ver", "Build Number", "CI", "% Clients", "Clients"]
    if not show_ci:
        headers.pop(2)
    print()
    print(tabulate(table, headers, stralign="right", numalign="right"))


def print_pacman(clients, show_ci):
    table = []
    for (version, ci), count in Counter([(u.pacman_version, u.ci) for u in clients]).most_common():
        pcnt = count / len(clients) * 100
        line = [version, ci, f"{pcnt:.2f}%", f"{count}"]
        if not show_ci:
            line.pop(1)
        table.append(line)
    headers = ["Pacman Ver", "CI", "% Clients", "Clients"]
    if not show_ci:
        headers.pop(1)
    print()
    print(tabulate(table, headers, stralign="right", numalign="right"))


def print_system_arch(clients, show_ci):
    table = []
    for (cygwin_arch, host_arch, ci), count in Counter([(u.user_agent.cygwin_arch, u.user_agent.host_arch, u.ci) for u in clients]).most_common():
        pcnt = count / len(clients) * 100
        line = [cygwin_arch, host_arch, ci, f"{pcnt:.2f}%", f"{count}"]
        if not show_ci:
            line.pop(2)
        table.append(line)
    headers = ["Cygwin Arch", "Host Arch", "CI", "% Clients", "Clients"]
    if not show_ci:
        headers.pop(2)
    print()
    print(tabulate(table, headers, stralign="right", numalign="right"))


def datetime_fromisoformat(value: str) -> datetime:
    # For Python <3.11
    if value.endswith("Z"):
        value = value[:-1] + "+00:00"
    return datetime.fromisoformat(value)


def main(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('infile', nargs='?', type=argparse.FileType('r', encoding="utf-8"), default=sys.stdin)
    parser.add_argument('--show-ci', action='store_true', help='show CI/cloud providers')
    parser.add_argument('--skip-ci', action='store_true', help='skip CI/cloud IP ranges')
    parser.add_argument('--only-ci', action='store_true', help='only CI/cloud IP ranges')
    parser.add_argument('--show-summary', action='store_true', help='show only a CI/cloud summary')
    args = parser.parse_args(argv[1:])

    assert not (args.skip_ci and args.only_ci)

    detect_ci = False

    if args.show_summary:
        assert not args.skip_ci
        assert not args.only_ci
        args.show_ci = True
        detect_ci = True

    if args.skip_ci or args.only_ci:
        detect_ci = True

    if detect_ci:
        ci_networks = get_ci_networks()

    entries: List[LogEntry] = []
    decoder = json.JSONDecoder()
    with args.infile as h:
        for line in h:
            if not line.startswith("{"):
                continue
            decoded = decoder.decode(line)
            if "RequestHost" not in decoded or "request_User-Agent" not in decoded:
                continue
            entry = LogEntry(
                decoded["ClientHost"],
                decoded["DownstreamStatus"],
                decoded["RequestHost"],
                decoded["RequestMethod"],
                decoded["RequestPath"],
                decoded["RequestProtocol"],
                decoded["request_User-Agent"],
                decoded["time"],
                None,
            )
            if entry.RequestHost not in ["repo.msys2.org", "mirror.msys2.org"]:
                continue
            if entry.RequestMethod != "GET":
                continue
            if not is_valid_user_agent(entry.UserAgent):
                continue
            entries.append(entry)

    def user_key(entry):
        # only take the system part of the UA, so we don't count
        # a client twice if it upgrades during a job for example
        system_id = "".join(entry.UserAgent.split()[1:3])
        return system_id + entry.ClientHost

    first = entries[0].time
    last = entries[0].time
    grouped = {}
    for entry in entries:
        if entry.time < first:
            first = entry.time
        if entry.time > last:
            last = entry.time
        key = user_key(entry)
        grouped.setdefault(key, []).append(entry)

    ip_to_ci = {}
    if detect_ci:
        def get_ip_to_ci(ip_addr: str) -> str:
            ip = IPAddress(ip_addr)
            for name, ipset in ci_networks.items():
                if ip in ipset:
                    return name
            return ""

        for group_entries in grouped.values():
            ip = group_entries[0].ClientHost
            ip_to_ci[ip] = get_ip_to_ci(ip)

    clients: List[ClientInfo] = []
    for group_entries in grouped.values():
        user_agent = parse_user_agent(group_entries[0].UserAgent)
        ci = ip_to_ci.get(group_entries[0].ClientHost, "")
        client_info = ClientInfo(
            ".".join(map(str, user_agent.pacman_version)),
            get_windows_edition(user_agent),
            user_agent,
            ci
        )
        clients.append(client_info)
        for e in group_entries:
            e.client_info = client_info

    if args.skip_ci:
        entries = [e for e in entries if not ip_to_ci.get(e.ClientHost, "")]
        clients = [c for c in clients if not c.ci]

    if args.only_ci:
        entries = [e for e in entries if ip_to_ci.get(e.ClientHost, "")]
        clients = [c for c in clients if c.ci]

    # Log info
    diff = datetime_fromisoformat(last) - datetime_fromisoformat(first)
    duration = (diff).total_seconds()
    requests_per_second = len(entries) / duration
    print(tabulate([
        ["Duration", f"from {first} to {last} ({diff})"],
        ["Requests", f"{len(entries)} ({requests_per_second:.2f}/s)"],
        ["Clients", f"{len(clients)} (clients are grouped by IP+WinVer+Arch)"],
        ["Included", "CI only" if args.only_ci else "non-CI only" if args.skip_ci else "all"],
    ]))

    # Repos
    if not args.show_summary:
        print_repos(entries, args.show_ci)

    # CI Systems
    if args.show_ci:
        print_ci_systems(clients, entries)

    # Windows versions
    if not args.show_summary:
        print_windows_major(clients, entries, args.show_ci)

    # Windows versions detailed
    if not args.show_summary:
        print_windows_version_details(clients, args.show_ci)

    # Pacman
    if not args.show_summary:
        print_pacman(clients, args.show_ci)

    # CPU Arch
    if not args.show_summary:
        print_system_arch(clients, args.show_ci)


if __name__ == "__main__":
    main(sys.argv)
