#!/usr/bin/env python3
"""Prune the Go build cache to a size limit.

Walks the cache directory, sorts all files by (atime, mtime, size, path)
descending, accumulates them until the size limit is reached, then deletes
the remainder.
"""

import argparse
import os
import re
import subprocess
import sys
import time


def parse_limit(s: str) -> int:
    suffixes = {"K": 1024, "M": 1024**2, "G": 1024**3}
    s = s.strip()
    if s[-1].upper() in suffixes:
        return int(s[:-1]) * suffixes[s[-1].upper()]
    return int(s)


def format_ago(ts: float) -> str:
    secs = time.time() - ts
    if secs < 60:
        return "{:.0f}s ago".format(secs)
    mins = secs / 60
    if mins < 60:
        return "{:.0f}m ago".format(mins)
    hours = mins / 60
    if hours < 24:
        return "{:.0f}h ago".format(hours)
    days = hours / 24
    if days < 14:
        return "{:.0f}d ago".format(days)
    weeks = days / 7
    if weeks < 52:
        return "{:.0f}w ago".format(weeks)
    return "{:.1f}y ago".format(days / 365)


def format_size(n: float) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024:
            fmt = "{:.0f} {}" if n % 1 == 0 else "{:.1f} {}"
            return fmt.format(n, unit)
        n /= 1024
    return "{:.1f} TB".format(n)


def gocache_default() -> str:
    result = subprocess.run(
        ["go", "env", "GOCACHE"],
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        print("error: go env GOCACHE failed: " + result.stderr.strip(), file=sys.stderr)
        sys.exit(1)
    return result.stdout.strip()


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "gocache",
        nargs="?",
        metavar="GOCACHE",
        help="Go build cache path (default: go env GOCACHE)",
    )
    parser.add_argument(
        "--limit",
        default="20G",
        type=parse_limit,
        metavar="N",
        help="Maximum cache size. Accepts K/M/G suffixes. Default: 20G",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print stats without deleting anything",
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Print each file with its atime, mtime, size, and keep/delete decision",
    )
    args = parser.parse_args()

    gocache = args.gocache if args.gocache is not None else gocache_default()

    if not os.path.isdir(gocache):
        print("error: not a directory: " + gocache, file=sys.stderr)
        return 1

    # Walk files of interest, calculating sort fields.
    # skip=True for files not directly under a 2-hex subdir; floats to top when sorted.
    hex2 = re.compile(r'^[0-9a-f]{2}$')
    files = []
    for dirpath, dirnames, filenames in os.walk(gocache):
        first = os.path.relpath(dirpath, gocache).split(os.sep)[0]
        skip = first == '.' or not hex2.match(first)
        for filename in filenames:
            path = os.path.join(dirpath, filename)
            try:
                st = os.lstat(path)
                files.append((skip, st.st_atime, st.st_mtime, st.st_size, path))
            except OSError:
                pass

    # Sort: skip files first (True > False), then by (atime-mtime, -size) descending.
    # atime-mtime measures reuse: a large positive delta means the file was written long ago
    # but accessed recently — a genuine cross-run cache hit. Files written and never reused
    # have delta≈0 and sort last. This is robust to noatime/relatime mounts where all files
    # in a single CI run appear accessed at roughly the same time. -size prefers smaller files
    # as tie-breaker, since action files are tiny and critical; evicting them orphans outputs.
    files.sort(key=lambda f: (f[0], f[1] - f[2], -f[3]), reverse=True)

    if args.debug:
        print("{:<12} {:<12} {:<10} {:<11} {}".format("atime", "mtime", "size", "decision", "path"))

    total = 0
    kept = 0
    deleted = 0
    deleted_size = 0
    errors = 0
    retained_first = retained_last = None
    deleted_first = deleted_last = None
    for skip, atime, mtime, size, path in files:
        if skip:
            decision = "skip"
            delete = False
        else:
            keep = total + size <= args.limit
            if keep:
                total += size
                kept += 1
                t = (atime, mtime, size)
                if retained_first is None:
                    retained_first = t
                retained_last = t
            delete = not keep
            decision = "retain" if keep else "delete"

        if args.debug:
            print("{:<12} {:<12} {:<10} {:<11} {}".format(
                format_ago(atime), format_ago(mtime), format_size(size), decision, path))

        if delete:
            deleted += 1
            deleted_size += size
            t = (atime, mtime, size)
            if deleted_first is None:
                deleted_first = t
            deleted_last = t
            if not args.dry_run:
                try:
                    os.remove(path)
                except OSError as e:
                    print("failed to remove {}: {}".format(path, e), file=sys.stderr)
                    errors += 1

    def fmt_tuple(t):
        return "({}, {}, {})".format(format_ago(t[0]), format_ago(t[1]), format_size(t[2]))

    action_retain = "would retain" if args.dry_run else "retained"
    action_delete = "would delete" if args.dry_run else "deleted"

    if retained_first is not None:
        print("{} {} files ({:.2f} GB) [{}, {}]".format(
            action_retain, kept, total / 1024**3, fmt_tuple(retained_first), fmt_tuple(retained_last)))
    else:
        print("{} {} files ({:.2f} GB)".format(action_retain, kept, total / 1024**3))

    if deleted_first is not None:
        print("{} {} files ({:.2f} GB) [{}, {}]".format(
            action_delete, deleted, deleted_size / 1024**3, fmt_tuple(deleted_first), fmt_tuple(deleted_last)))
    elif deleted:
        print("{} {} files ({:.2f} GB)".format(action_delete, deleted, deleted_size / 1024**3))
    if errors:
        print("errors: {}".format(errors), file=sys.stderr)
        return 1
    return 0


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