#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-only
"""Measure the effort to recreate each final Linux release using SLOCCount.

Example (also fetch release tags for a shallow checkout):
    scripts/sloccount-releases.py --fetch-tags -o linux-effort.csv

Requires Python 3.9+, Git, and SLOCCount. Sources and SLOCCount's cache live
in Linux tmpfs or an automatically created, case-sensitive macOS RAM disk.
Only the CSV and explicitly fetched Git objects are written to disk.

CSV columns: date,version,years. Dates are UTC commit dates, with tag dates
as a fallback. The undated v2.6.11 tree tag uses its recorded release date,
2005-03-02. Rows are flushed in descending date order, regardless of worker
completion order. Years means
COCOMO person-years using sloccount --effort 4.646 1.12:
    person-years = 4.646 * (KSLOC ** 1.12) / 12
See https://dwheeler.com/sloccount/sloccount.html for the estimation model.

The default total memory budget is at most 16 GiB, half currently available
RAM, and one quarter of physical RAM. Each worker reserves 512 MiB for
processes in addition to its estimated source, cache, and filesystem space.
--jobs is an upper bound: the memory budget can reduce concurrency further.
RAM-backed files can still be swapped by the OS under memory pressure.
"""

import argparse
from collections import deque
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from datetime import datetime, timezone
import csv
import math
import os
from pathlib import Path, PurePosixPath
import platform
import plistlib
import re
import shutil
import signal
import subprocess
import sys
import tarfile
import tempfile
import threading
import uuid


MIB = 1024 ** 2
GIB = 1024 ** 3
# A full 96,039-file Linux tree with SLOCCount 2.26 used about 28 MiB
# peak process-tree RSS and 23 MiB of cache in a macOS profile. Allow
# substantially more for other versions, Git, extraction, and Python.
PROCESS_MEMORY = 512 * MIB
COORDINATOR_MEMORY = 256 * MIB
RELEASE = re.compile(r"v[0-9]+(?:\.[0-9]+)+\Z")
# This pre-Git release points to a tree and has no tagger header either.
# Match the original tag object so unrelated undated tags still fail.
# Release date: https://www.kernel.org/pub/linux/kernel/v2.6/
RELEASE_DATE_FALLBACKS = {
    ("v2.6.11", "5dc01c595e6c6ec9ccda4f6f69c131c0dd945f8c"):
        "2005-03-02T00:00:00+00:00",
}
EFFORT = re.compile(
    r"^.*Effort[^\n]*Person-Years\s*\(Person-Months\)\s*="
    r"\s*([0-9][0-9,]*\.[0-9]+)\s*\(", re.MULTILINE)


class Commands:
    """Track child process groups so cancellation also stops SLOCCount helpers."""

    def __init__(self):
        self.lock = threading.Lock()
        self.children = set()
        self.stopped = threading.Event()

    @contextmanager
    def process(self, argv, **kwargs):
        with self.lock:
            if self.stopped.is_set():
                raise RuntimeError("Run cancelled")
            child = subprocess.Popen(argv, start_new_session=True, **kwargs)
            self.children.add(child)
        try:
            yield child
        finally:
            if child.poll() is None:
                self.terminate(child)
            child.wait()
            for stream in (child.stdin, child.stdout, child.stderr):
                if stream is not None:
                    stream.close()
            with self.lock:
                self.children.discard(child)

    @staticmethod
    def terminate(child):
        try:
            os.killpg(child.pid, signal.SIGTERM)
        except ProcessLookupError:
            return
        try:
            child.wait(timeout=5)
        except subprocess.TimeoutExpired:
            try:
                os.killpg(child.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass

    def cancel(self):
        with self.lock:
            self.stopped.set()
            children = list(self.children)
        for child in children:
            self.terminate(child)

    def run(self, argv, **kwargs):
        with self.process(argv, stdout=subprocess.PIPE,
                          stderr=subprocess.PIPE, **kwargs) as child:
            output, errors = child.communicate()
            if child.returncode:
                detail = (errors or output)[-4000:].decode(errors="replace").strip()
                raise RuntimeError("{} failed: {}".format(argv[0], detail))
            return output


def memory_status(commands):
    """Use conservative available-memory figures, including Linux cgroup caps."""
    if platform.system() == "Darwin":
        total = int(commands.run(["sysctl", "-n", "hw.memsize"]))
        stats = commands.run(["vm_stat"]).decode()
        page_size = int(re.search(r"page size of (\d+) bytes", stats)[1])
        pages = sum(int(re.search(r"^" + name + r":\s+(\d+)",
                                  stats, re.MULTILINE)[1])
                    for name in ("Pages free", "Pages speculative"))
        return total, pages * page_size
    if platform.system() != "Linux":
        raise RuntimeError("RAM-backed workspaces require Linux or macOS")
    entries = dict(re.findall(r"^(\w+):\s+(\d+)",
                             Path("/proc/meminfo").read_text(), re.MULTILINE))
    total = int(entries["MemTotal"]) * 1024
    available = int(entries.get("MemAvailable", entries["MemFree"])) * 1024
    for limit_path, usage_path in (
            ("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory.current"),
            ("/sys/fs/cgroup/memory/memory.limit_in_bytes",
             "/sys/fs/cgroup/memory/memory.usage_in_bytes")):
        try:
            limit = int(Path(limit_path).read_text())
            usage = int(Path(usage_path).read_text())
        except (OSError, ValueError):
            continue
        total = min(total, limit)
        available = min(available, max(0, limit - usage))
    return total, available


@contextmanager
def ram_workspace(commands, capacity, tmpfs):
    if platform.system() == "Linux":
        parent = Path(tmpfs).resolve()
        kind = commands.run(["stat", "-f", "-c", "%T", str(parent)]).strip()
        if kind not in (b"tmpfs", b"ramfs"):
            raise RuntimeError("{} is not tmpfs/ramfs".format(parent))
        if shutil.disk_usage(parent).free < capacity:
            raise RuntimeError("Not enough tmpfs space; enlarge it or lower "
                               "--memory-gib/--jobs (no disk fallback is used)")
        with tempfile.TemporaryDirectory(prefix="linux-sloc-", dir=parent) as tmp:
            yield Path(tmp)
        return

    # Format only the new anonymous RAM device returned by hdiutil.
    device = commands.run([
        "hdiutil", "attach", "-nomount", "ram://{}".format(capacity // 512)
    ]).decode().strip()
    if not re.fullmatch(r"/dev/disk[0-9]+", device):
        raise RuntimeError("Unexpected RAM device response: {!r}".format(device))
    try:
        label = "linux-sloc-" + uuid.uuid4().hex[:12]
        commands.run(["diskutil", "eraseVolume", "HFSX", label, device])
        info = plistlib.loads(commands.run(["diskutil", "info", "-plist", device]))
        yield Path(info["MountPoint"])
    finally:
        # Cleanup must still work after the command runner has been cancelled.
        result = subprocess.run(["hdiutil", "detach", device],
                                stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if result.returncode:
            print("Could not detach RAM disk {}: {}".format(
                device, result.stderr.decode(errors="replace").strip()), file=sys.stderr)


class WorkspaceBudget:
    def __init__(self, capacity, commands, keep_free):
        self.capacity = capacity
        self.free = capacity
        self.commands = commands
        self.keep_free = keep_free
        self.condition = threading.Condition()

    @contextmanager
    def reserve(self, size):
        if size > self.capacity:
            raise RuntimeError("A release needs {:.2f} GiB of workspace; only "
                               "{:.2f} GiB is budgeted".format(size / GIB,
                                                             self.capacity / GIB))
        with self.condition:
            while self.free < size:
                if self.commands.stopped.is_set():
                    raise RuntimeError("Run cancelled")
                self.condition.wait(timeout=1)
            if memory_status(self.commands)[1] < self.keep_free:
                raise RuntimeError("Available RAM fell below the safety reserve; "
                                   "close other applications and retry")
            self.free -= size
        try:
            yield
        finally:
            with self.condition:
                self.free += size
                self.condition.notify_all()


def release_tags(commands, git):
    # Dereferenced commit dates handle annotated tags; direct commit dates
    # handle lightweight tags. Tree tags need a tag date or known fallback.
    fmt = ("%(refname:strip=2)%09%(objectname)%09%(*committerdate:iso-strict)"
           "%09%(committerdate:iso-strict)%09%(taggerdate:iso-strict)")
    releases = []
    for line in commands.run(git + ["for-each-ref", "--format=" + fmt,
                                    "refs/tags"]).decode().splitlines():
        version, oid, *dates = line.split("\t")
        if RELEASE.fullmatch(version):
            date = next((value for value in dates if value), None)
            if date is None:
                date = RELEASE_DATE_FALLBACKS.get((version, oid))
            if date is None:
                raise RuntimeError("Cannot determine release date of " + version)
            when = datetime.fromisoformat(date).astimezone(timezone.utc)
            releases.append((when, version, oid))
    return sorted(releases, reverse=True)


def fetch_tags(commands, git, remote):
    refs = commands.run(git + ["ls-remote", "--tags", "--refs", remote])
    names = [line.split()[1].decode() for line in refs.splitlines()
             if RELEASE.fullmatch(line.split()[1].decode().removeprefix("refs/tags/"))]
    if not names:
        raise RuntimeError("No final release tags found on " + remote)
    fetch = git + ["fetch", "--no-tags"]
    if commands.run(git + ["rev-parse", "--is-shallow-repository"]).strip() == b"true":
        fetch.append("--depth=1")
    print("Fetching {} final release tags from {}...".format(len(names), remote),
          file=sys.stderr, flush=True)
    commands.run(fetch + [remote] + [name + ":" + name for name in names])


def workspace_size(commands, git, oid):
    """Budget file allocation, metadata, SLOCCount lists, and a 20% margin."""
    files = size = 0
    page = os.sysconf("SC_PAGE_SIZE") if platform.system() == "Linux" else 4096
    listing = commands.run(git + ["ls-tree", "-rlz", oid])
    for record in listing.split(b"\0"):
        if not record:
            continue
        mode, kind, _, length = record.split(b"\t", 1)[0].split()
        if kind == b"blob" and mode != b"120000":
            size += ((int(length) + page - 1) // page) * page
            files += 1
    # 1 KiB metadata + 2 KiB of SLOCCount cache per tracked file, plus
    # fixed space for filesystem overhead, logs, and small source trees.
    return math.ceil((size + files * 3072 + 128 * MIB) * 1.2 / MIB) * MIB


def extract_source(commands, git, oid, destination, env):
    with commands.process(git + ["archive", "--format=tar", oid],
                          stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                          env=env) as child:
        with tarfile.open(fileobj=child.stdout, mode="r|") as archive:
            for member in archive:
                path = PurePosixPath(member.name)
                if path.is_absolute() or ".." in path.parts:
                    raise RuntimeError("Unsafe archive path: " + member.name)
                target = destination.joinpath(*path.parts)
                if member.isdir():
                    target.mkdir(parents=True, exist_ok=True)
                elif member.isfile():
                    target.parent.mkdir(parents=True, exist_ok=True)
                    with archive.extractfile(member) as source, target.open("xb") as output:
                        shutil.copyfileobj(source, output)
                    target.chmod(member.mode & 0o777)
                # SLOCCount does not follow symlinks by default. Do not
                # materialize them, or retain a full tar member index in RAM.
                archive.members.clear()
        errors = child.stderr.read()
        if child.wait():
            raise RuntimeError("git archive failed: " + errors.decode(errors="replace"))


def measure(release, commands, git, root, budget):
    when, version, oid = release
    size = workspace_size(commands, git, oid)
    with budget.reserve(size), tempfile.TemporaryDirectory(prefix=version + "-",
                                                          dir=root) as tmp:
        workspace = Path(tmp)
        source, data = workspace / "source", workspace / "data"
        source.mkdir()
        data.mkdir()
        env = dict(os.environ, LC_ALL="C", LANG="C", TMPDIR=str(workspace),
                   TMP=str(workspace), TEMP=str(workspace))
        print("Counting {} ({})...".format(version, when.date()),
              file=sys.stderr, flush=True)
        extract_source(commands, git, oid, source, env)
        output = commands.run(["sloccount", "--effort", "4.646", "1.12",
                               "--datadir", str(data), str(source)],
                              cwd=workspace, env=env).decode(errors="replace")
        match = EFFORT.search(output)
        if not match:
            raise RuntimeError("No SLOCCount person-years estimate for {}:\n{}".format(
                version, output[-4000:]))
        return when.date().isoformat(), version, match[1].replace(",", "")


def positive_int(value):
    number = int(value)
    if number < 1:
        raise argparse.ArgumentTypeError("must be positive")
    return number


def main():
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("-o", "--output", default="linux-effort.csv",
                        help="CSV file to overwrite, or - for stdout (default: %(default)s)")
    parser.add_argument("--repo", type=Path, default=Path(__file__).resolve().parents[1],
                        help="Git repository (default: this script's Linux tree)")
    parser.add_argument("--fetch-tags", action="store_true",
                        help="fetch final release tags and objects before counting")
    parser.add_argument("--remote", default="origin", help="remote for --fetch-tags")
    parser.add_argument("-j", "--jobs", type=positive_int, default=os.cpu_count() or 1,
                        help="maximum worker threads (default: CPU count, capped by RAM)")
    parser.add_argument("--memory-gib", type=float, default=16,
                        help="total RAM budget ceiling in GiB (default: %(default)s)")
    parser.add_argument("--tmpfs", default="/dev/shm",
                        help="Linux tmpfs directory (default: %(default)s)")
    parser.add_argument("--limit", type=positive_int,
                        help="only measure this many newest releases")
    args = parser.parse_args()
    if not math.isfinite(args.memory_gib) or args.memory_gib <= 0:
        parser.error("--memory-gib must be a positive finite number")
    for executable in ("git", "sloccount"):
        if not shutil.which(executable):
            parser.error("Required executable not found: " + executable)
    commands = Commands()
    git = ["git", "-C", str(args.repo.resolve()), "-c", "core.packedGitLimit=128m",
           "-c", "pack.threads=1", "-c", "core.deltaBaseCacheLimit=64m"]
    try:
        if args.fetch_tags:
            fetch_tags(commands, git, args.remote)
        releases = release_tags(commands, git)
        if not releases:
            raise RuntimeError("No final release tags found. Re-run with --fetch-tags "
                               "to fetch them from this repository's remote.")
        releases = releases[:args.limit]
        total, available = memory_status(commands)
        ceiling = min(int(args.memory_gib * GIB), total // 4, available // 2)
        workspace = workspace_size(commands, git, releases[0][2])
        jobs = min(args.jobs, len(releases),
                   (ceiling - COORDINATOR_MEMORY) // (workspace + PROCESS_MEMORY))
        if jobs < 1:
            raise RuntimeError("Insufficient safe RAM budget: {:.2f} GiB available to "
                               "this run; one worker needs {:.2f} GiB".format(
                                   ceiling / GIB,
                                   (workspace + PROCESS_MEMORY + COORDINATOR_MEMORY) / GIB))
        capacity = min(jobs * workspace,
                       (ceiling - COORDINATOR_MEMORY - jobs * PROCESS_MEMORY) // MIB * MIB)
        print("{} releases; {} workers; {:.2f} GiB total budget "
              "({:.2f} GiB workspace, {:.2f} GiB process reserve per worker).".format(
                  len(releases), jobs, ceiling / GIB, capacity / GIB,
                  PROCESS_MEMORY / GIB), file=sys.stderr, flush=True)
        with ram_workspace(commands, capacity, args.tmpfs) as root:
            budget = WorkspaceBudget(capacity, commands, min(total // 4, available // 2))
            output = (os.fdopen(os.dup(sys.stdout.fileno()), "w", buffering=1, newline="")
                      if args.output == "-" else open(args.output, "w", buffering=1, newline=""))
            with output:
                writer = csv.writer(output, lineterminator="\n")
                writer.writerow(("date", "version", "years"))
                output.flush()
                executor = ThreadPoolExecutor(max_workers=jobs)
                pending = deque()
                remaining = iter(releases)
                try:
                    for release in (next(remaining, None) for _ in range(jobs)):
                        if release is not None:
                            pending.append(executor.submit(measure, release, commands,
                                                           git, root, budget))
                    while pending:
                        writer.writerow(pending.popleft().result())
                        output.flush()
                        release = next(remaining, None)
                        if release is not None:
                            pending.append(executor.submit(measure, release, commands,
                                                           git, root, budget))
                except BaseException:
                    commands.cancel()
                    raise
                finally:
                    executor.shutdown(wait=True, cancel_futures=True)
        return 0
    except (OSError, ValueError, RuntimeError, tarfile.TarError) as error:
        print("error: " + str(error), file=sys.stderr)
        return 1
    except KeyboardInterrupt:
        commands.cancel()
        print("Interrupted; completed CSV rows have been preserved.", file=sys.stderr)
        return 130


if __name__ == "__main__":
    def interrupt(signum, frame):
        raise KeyboardInterrupt

    signal.signal(signal.SIGTERM, interrupt)
    sys.exit(main())