"""Shared definitions for the reproduction kit. Standard library only.

Everything here restates, without changing it, what the study's own scripts do
(analyze.py, aggregate.py and facts.py, published one folder up, beside the
article):

  * a domain is the registrable domain (eTLD+1) under the Public Suffix List;
  * only requests that finished within the first 60 seconds of a load count;
  * data: and blob: URLs are not requests;
  * a request is THIRD PARTY when the owner of its registrable domain differs
    from the owner of the page's registrable domain (owners come from
    entities.json; a domain with no entry is its own owner);
  * a domain is ADVERTISING/TRACKING when its categories include ad, analytics,
    marketing or tag-manager.
"""
import csv, glob, gzip, json, os, urllib.parse

HERE = os.path.dirname(os.path.abspath(__file__))
WINDOW_MS = 60_000
AD_TRACKING = {"ad", "analytics", "marketing", "tag-manager",
               "advertising", "ad motivated tracking", "audience measurement",
               "third-party analytics marketing", "action pixels", "ad fraud", "tracking"}
TABLE = os.path.join(HERE, "requests-noquery.csv.gz")   # the published request table
TABLE_COLUMNS = ["page", "portal", "load", "page_url", "seconds", "url", "resource_type"]


def find(name, *extra):
    """Locate a study file in this folder or one folder up (where the article publishes entities.json and supplement.csv)."""
    for p in [os.path.join(HERE, name), os.path.join(HERE, "..", name), *extra]:
        if os.path.exists(p):
            return p
    raise SystemExit(f"cannot find {name}: download it (README.md lists every file and its address) "
                     f"and put it in {HERE} or in the folder above it")


def load_psl(path=None):
    path = path or find("public_suffix_list.dat")
    rules, exceptions, wild = set(), set(), set()
    for line in open(path, encoding="utf-8"):
        line = line.strip()
        if not line or line.startswith("//"):
            continue
        rule = line.split()[0].lower()
        if rule.startswith("!"):
            exceptions.add(rule[1:])
        elif rule.startswith("*."):
            wild.add(rule[2:])
        else:
            rules.add(rule)
    return rules, exceptions, wild


_PSL = None


def registrable(host):
    """eTLD+1 of a hostname. Identical to analyze.registrable."""
    global _PSL
    if _PSL is None:
        _PSL = load_psl()
    host = (host or "").lower().rstrip(".")
    if not host or host.replace(".", "").isdigit():
        return host
    rules, exceptions, wild = _PSL
    labels = host.split(".")
    for i in range(len(labels)):
        cand = ".".join(labels[i:])
        if cand in exceptions:
            return ".".join(labels[i:])
        parent = ".".join(labels[i + 1:])
        if cand in rules or (parent in wild):
            return ".".join(labels[i - 1:]) if i > 0 else host
    return ".".join(labels[-2:])


def hostname(url):
    return urllib.parse.urlsplit(url).hostname or ""


class Entities:
    def __init__(self, path=None):
        self.map = json.load(open(path or find("entities.json")))

    def owner(self, reg):
        return (self.map.get(reg) or {}).get("owner") or reg

    def categories(self, reg):
        return [c for c in (self.map.get(reg) or {}).get("categories", [])]

    def is_ad_tracking(self, reg):
        return bool({c.lower() for c in self.categories(reg)} & AD_TRACKING)

    def label(self, reg):
        """ad/tracking | other (classified, not ad/tracking) | unknown (no entry, or category unknown)."""
        if self.is_ad_tracking(reg):
            return "ad/tracking"
        return "other" if [c for c in self.categories(reg) if c != "unknown"] else "unknown"


def loads_from_raw(raw_dir):
    """One dict per load, straight from a folder of measure.cjs captures (the study's own are not published)."""
    for f in sorted(glob.glob(os.path.join(raw_dir, "*.json"))):
        d = json.load(open(f))
        yield {"page": d["label"], "portal": d["label"].split("-")[0], "load": int(d["run"]),
               "page_url": d["url"],
               "requests": [(int(r["t_ms"]), r["url"], r.get("type") or "")
                            for r in d["requests"] if not r["url"].startswith(("data:", "blob:"))]}


def loads_from_table(path=TABLE):
    """The same loads, from the published table."""
    loads = {}
    with gzip.open(path, "rt", encoding="utf-8", newline="") as fh:
        for row in csv.DictReader(fh):
            key = (row["page"], int(row["load"]))
            if key not in loads:
                loads[key] = {"page": row["page"], "portal": row["portal"], "load": int(row["load"]),
                              "page_url": row["page_url"], "requests": []}
            loads[key]["requests"].append((round(float(row["seconds"]) * 1000), row["url"], row["resource_type"]))
    for key in sorted(loads):
        yield loads[key]


def open_loads(source="auto"):
    """source: 'table' (requests-noquery.csv.gz in this folder, the default), or a path to
    another .csv.gz table or to a folder of measure.cjs captures."""
    if source == "table":
        if not os.path.exists(TABLE):
            raise SystemExit(f"cannot find {os.path.basename(TABLE)}: download it from "
                             f"https://manxlab.com/research/free-game-trackers/reproduce/{os.path.basename(TABLE)} "
                             f"and put it in {HERE}")
        return os.path.basename(TABLE), loads_from_table(TABLE)
    if not os.path.exists(source):
        raise SystemExit(f"cannot find {source}")
    if os.path.isdir(source):
        return source, loads_from_raw(source)
    return source, loads_from_table(source)