Save PR description with media content and PR patch

ImportInterestingPulls.java

///usr/bin/env jbang --jdk-providers=mise "$0" "$@" ; exit $?
//JAVA 27
//PREVIEW

/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.nio.file.DirectoryStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardCopyOption;
import java.text.Normalizer;
import java.util.ArrayList;
import java.util.Base64;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Locale;
import java.util.concurrent.StructuredTaskScope;
import java.util.concurrent.StructuredTaskScope.Joiner;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class ImportInterestingPulls {
  private static final Object GIT_LOCK = new Object();
  private static final Pattern PR_URL = Pattern.compile(
      "https://github\\.com/([^/]+)/([^/]+)/pull/(\\d+)(?:[/#?].*)?", Pattern.CASE_INSENSITIVE);
  private static final Pattern GITHUB_REMOTE = Pattern.compile(
      "github\\.com[:/]([^/]+)/([^/]+)$", Pattern.CASE_INSENSITIVE);
  private static final Pattern URL = Pattern.compile("https://[^\\s<>)\\\"']+");
  private static final Pattern ATTACHMENT_URL = Pattern.compile(
      "https://(?:github\\.com/(?:user-attachments/assets|[^/\\s]+/[^/\\s]+/assets)/[^\\s<>)\\\"']+|(?:private-)?user-images\\.githubusercontent\\.com/[^\\s<>)\\\"']+)");
  private static final Pattern MEDIA_ELEMENT = Pattern.compile(
      "<(?:img|video|source)\\b[^>]*>", Pattern.CASE_INSENSITIVE);
  private static final Pattern SRC_ATTRIBUTE = Pattern.compile(
      "\\bsrc=\\\"([^\\\"]+)\\\"", Pattern.CASE_INSENSITIVE);

  public static void main(String... args) throws Exception {
    if (args.length == 1 && args[0].equals("--self-test")) {
      selfTest();
      return;
    }
    if (args.length == 0) {
      throw new IllegalArgumentException(
          "Usage: ImportInterestingPulls.java <PR number or GitHub PR URL>...");
    }

    Path repositoryRoot = Path.of(runText("rtk", "git", "rev-parse", "--show-toplevel").trim());
    String repository = githubRepository(runText("rtk", "git", "remote", "get-url", "origin").trim());
    List<PullSpec> pullRequests = List.of(args).stream()
        .map(argument -> parsePull(argument, repository))
        .distinct()
        .toList();
    Path outputRoot = repositoryRoot.resolve("interesting-pulls");
    Files.createDirectories(outputRoot);

    try (HttpClient http = HttpClient.newBuilder()
        .followRedirects(HttpClient.Redirect.ALWAYS)
        .build();
        var scope = StructuredTaskScope.open(Joiner.<String>allSuccessfulOrThrow())) {
      pullRequests.forEach(pullRequest ->
          scope.fork(() -> importPullRequest(http, outputRoot, pullRequest)));
      scope.join().forEach(System.out::println);
    }
  }

  private static String importPullRequest(HttpClient http, Path outputRoot, PullSpec spec) throws Exception {
    PullRequest pullRequest = readPullRequest(spec);
    Path directory = findOrCreateDirectory(outputRoot, pullRequest);

    fetchBranch(pullRequest);
    String localizedBody = downloadMedia(http, directory, pullRequest.body(), pullRequest.bodyHtml());

    String markdown = "# " + pullRequest.title().stripTrailing() + "\n\n"
        + "Branch: `" + pullRequest.branch() + "`\n\n"
        + localizedBody;
    Files.writeString(directory.resolve("README.md"), markdown, StandardCharsets.UTF_8);

    byte[] diff = run(
        "rtk", "gh", "api",
        "-H", "Accept: application/vnd.github.diff",
        "repos/" + pullRequest.repository() + "/pulls/" + pullRequest.number());
    Files.write(directory.resolve("pr-" + pullRequest.number() + ".diff"), diff);

    return "PR " + pullRequest.number() + ": " + directory.getFileName() + " -> " + pullRequest.branch();
  }

  private static PullRequest readPullRequest(PullSpec spec) throws Exception {
    String output = runText(
        "rtk", "gh", "api",
        "-H", "Accept: application/vnd.github.full+json",
        "repos/" + spec.repository() + "/pulls/" + spec.number(),
        "--jq", "[.title, (.body // \"\"), (.body_html // \"\"), .head.ref, .head.sha] | .[] | @base64");
    String[] lines = output.split("\\R", -1);
    if (lines.length < 5) {
      throw new IOException("Unexpected metadata response for PR " + spec.number());
    }

    return new PullRequest(
        spec.repository(),
        spec.number(),
        decode(lines[0]),
        decode(lines[1]),
        decode(lines[2]),
        decode(lines[3]),
        decode(lines[4]));
  }

  private static Path findOrCreateDirectory(Path outputRoot, PullRequest pullRequest) throws IOException {
    List<Path> existing = new ArrayList<>();
    try (DirectoryStream<Path> stream = Files.newDirectoryStream(
        outputRoot, "pr-" + pullRequest.number() + "-*")) {
      stream.forEach(existing::add);
    }
    if (existing.size() > 1) {
      throw new IOException("Multiple folders already exist for PR " + pullRequest.number() + ": " + existing);
    }

    Path directory = existing.isEmpty()
        ? outputRoot.resolve("pr-" + pullRequest.number() + "-" + slug(pullRequest.title()))
        : existing.getFirst();
    Files.createDirectories(directory);
    return directory;
  }

  private static void fetchBranch(PullRequest pullRequest) throws Exception {
    // ponytail: Git mutates shared refs, so serialize this small section unless large batches need more throughput.
    synchronized (GIT_LOCK) {
      run("rtk", "git", "check-ref-format", "--branch", pullRequest.branch());
      run("rtk", "git", "fetch", "origin", "refs/pull/" + pullRequest.number() + "/head");

      String fetchedSha = runText("rtk", "git", "rev-parse", "FETCH_HEAD").trim();
      if (!fetchedSha.equals(pullRequest.sha())) {
        throw new IOException("Fetched PR " + pullRequest.number() + " at " + fetchedSha
            + ", but GitHub reported " + pullRequest.sha());
      }

      String ref = "refs/heads/" + pullRequest.branch();
      if (exitCode("rtk", "git", "show-ref", "--verify", "--quiet", ref) == 0) {
        String localSha = runText("rtk", "git", "rev-parse", ref).trim();
        if (!localSha.equals(pullRequest.sha())) {
          throw new IOException("Local branch " + pullRequest.branch() + " points to " + localSha
              + "; refusing to overwrite it with " + pullRequest.sha());
        }
        return;
      }

      run("rtk", "git", "branch", pullRequest.branch(), pullRequest.sha());
    }
  }

  private static PullSpec parsePull(String argument, String repository) {
    if (argument.matches("[1-9]\\d*")) {
      return new PullSpec(repository, Integer.parseInt(argument));
    }

    Matcher matcher = PR_URL.matcher(argument);
    if (!matcher.matches()) {
      throw new IllegalArgumentException("Not a PR number or GitHub PR URL: " + argument);
    }

    String urlRepository = matcher.group(1) + "/" + matcher.group(2);
    if (!urlRepository.equalsIgnoreCase(repository)) {
      throw new IllegalArgumentException(
          "PR URL repository " + urlRepository + " does not match origin " + repository);
    }
    return new PullSpec(repository, Integer.parseInt(matcher.group(3)));
  }

  private static String githubRepository(String remoteUrl) {
    Matcher matcher = GITHUB_REMOTE.matcher(remoteUrl);
    if (!matcher.find()) {
      throw new IllegalArgumentException("Origin is not a GitHub repository: " + remoteUrl);
    }
    return matcher.group(1) + "/" + matcher.group(2).replaceFirst("\\.git$", "");
  }

  private static String downloadMedia(HttpClient http, Path directory, String body, String bodyHtml)
      throws Exception {
    List<RenderedMedia> renderedMedia = renderedMedia(bodyHtml);
    List<String> renderedUrls = urls(URL, bodyHtml.replace("&amp;", "&")).stream().toList();
    String localizedBody = body;
    int index = 1;

    for (String originalUrl : mediaUrls(body, renderedMedia)) {
      String key = assetKey(originalUrl);
      String downloadUrl = renderedMedia.stream()
          .filter(media -> media.html().contains(originalUrl))
          .map(RenderedMedia::url)
          .findFirst()
          .orElseGet(() -> renderedUrls.stream()
          .filter(url -> url.contains(key))
          .findFirst()
          .orElse(originalUrl));
      String extension = extension(downloadUrl);
      Path temporaryFile = Files.createTempFile(directory, ".media-", ".download");

      try {
        HttpResponse<Path> response = http.send(
            HttpRequest.newBuilder(URI.create(downloadUrl)).GET().build(),
            HttpResponse.BodyHandlers.ofFile(temporaryFile));
        if (response.statusCode() / 100 != 2) {
          throw new IOException("Downloading " + originalUrl + " returned HTTP " + response.statusCode());
        }

        String contentType = response.headers().firstValue("content-type").orElse("");
        if (contentType.startsWith("text/html")) {
          throw new IOException("Downloading " + originalUrl + " returned HTML instead of media");
        }
        if (extension.equals("bin")) {
          extension = extensionFor(contentType);
        }

        String filename = "%02d-%s.%s".formatted(index++, safeName(key), extension);
        Files.move(temporaryFile, directory.resolve(filename), StandardCopyOption.REPLACE_EXISTING);
        localizedBody = replaceOutsideFences(localizedBody, originalUrl, "./" + filename);
      } finally {
        Files.deleteIfExists(temporaryFile);
      }
    }

    return localizedBody;
  }

  private static LinkedHashSet<String> mediaUrls(String body, List<RenderedMedia> renderedMedia) {
    LinkedHashSet<String> mediaUrls = new LinkedHashSet<>();
    for (String url : urlsOutsideFences(body)) {
      if (ATTACHMENT_URL.matcher(url).matches()
          || renderedMedia.stream().anyMatch(media -> media.html().contains(url))) {
        mediaUrls.add(url);
      }
    }
    return mediaUrls;
  }

  private static List<RenderedMedia> renderedMedia(String bodyHtml) {
    List<RenderedMedia> media = new ArrayList<>();
    Matcher elementMatcher = MEDIA_ELEMENT.matcher(bodyHtml);
    while (elementMatcher.find()) {
      String html = elementMatcher.group().replace("&amp;", "&");
      Matcher sourceMatcher = SRC_ATTRIBUTE.matcher(html);
      if (sourceMatcher.find()) {
        media.add(new RenderedMedia(html, sourceMatcher.group(1)));
      }
    }
    return media;
  }

  private static LinkedHashSet<String> urlsOutsideFences(String body) {
    LinkedHashSet<String> urls = new LinkedHashSet<>();
    boolean inFence = false;
    for (String line : body.split("(?<=\\n)", -1)) {
      String trimmed = line.stripLeading();
      if (trimmed.startsWith("```") || trimmed.startsWith("~~~")) {
        inFence = !inFence;
      } else if (!inFence) {
        urls.addAll(urls(URL, line));
      }
    }
    return urls;
  }

  private static String replaceOutsideFences(String body, String target, String replacement) {
    StringBuilder result = new StringBuilder(body.length());
    boolean inFence = false;
    for (String line : body.split("(?<=\\n)", -1)) {
      String trimmed = line.stripLeading();
      if (trimmed.startsWith("```") || trimmed.startsWith("~~~")) {
        inFence = !inFence;
      }
      if (inFence) {
        result.append(line);
      } else if (line.strip().equals(target)) {
        String filename = Path.of(replacement).getFileName().toString();
        result.append(line.replace(target, "[" + filename + "](" + replacement + ")"));
      } else {
        result.append(line.replace(target, replacement));
      }
    }
    return result.toString();
  }

  private static LinkedHashSet<String> urls(Pattern pattern, String text) {
    LinkedHashSet<String> urls = new LinkedHashSet<>();
    Matcher matcher = pattern.matcher(text);
    while (matcher.find()) {
      urls.add(matcher.group());
    }
    return urls;
  }

  private static String assetKey(String url) {
    String path = URI.create(url).getPath();
    return path.substring(path.lastIndexOf('/') + 1).replaceFirst("\\.[A-Za-z0-9]{2,5}$", "");
  }

  private static String extension(String url) {
    String path = URI.create(url.replace("&amp;", "&")).getPath();
    Matcher matcher = Pattern.compile("\\.([A-Za-z0-9]{2,5})$").matcher(path);
    return matcher.find() ? matcher.group(1).toLowerCase(Locale.ROOT) : "bin";
  }

  private static String extensionFor(String contentType) {
    return switch (contentType.split(";", 2)[0]) {
      case "image/png" -> "png";
      case "image/jpeg" -> "jpg";
      case "image/gif" -> "gif";
      case "image/webp" -> "webp";
      case "image/svg+xml" -> "svg";
      case "video/mp4" -> "mp4";
      case "video/quicktime" -> "mov";
      case "video/webm" -> "webm";
      default -> "bin";
    };
  }

  private static String safeName(String value) {
    String safe = value.replaceAll("[^A-Za-z0-9._-]+", "-");
    return safe.isBlank() ? "media" : safe;
  }

  private static String slug(String title) {
    String slug = Normalizer.normalize(title, Normalizer.Form.NFKD)
        .replaceAll("\\p{M}+", "")
        .toLowerCase(Locale.ROOT)
        .replaceAll("[^a-z0-9]+", "-")
        .replaceAll("(^-+|-+$)", "");
    if (slug.length() <= 60) {
      return slug.isEmpty() ? "untitled" : slug;
    }

    slug = slug.substring(0, 60);
    int lastDash = slug.lastIndexOf('-');
    return lastDash >= 40 ? slug.substring(0, lastDash) : slug.replaceFirst("-+$", "");
  }

  private static String decode(String value) {
    return new String(Base64.getDecoder().decode(value), StandardCharsets.UTF_8);
  }

  private static String runText(String... command) throws Exception {
    return new String(run(command), StandardCharsets.UTF_8);
  }

  private static byte[] run(String... command) throws Exception {
    Process process = new ProcessBuilder(command).redirectError(ProcessBuilder.Redirect.INHERIT).start();
    byte[] output = process.getInputStream().readAllBytes();
    int exitCode = process.waitFor();
    if (exitCode != 0) {
      throw new IOException(String.join(" ", command) + " exited with " + exitCode);
    }
    return output;
  }

  private static int exitCode(String... command) throws Exception {
    return new ProcessBuilder(command)
        .redirectOutput(ProcessBuilder.Redirect.DISCARD)
        .redirectError(ProcessBuilder.Redirect.DISCARD)
        .start()
        .waitFor();
  }

  private static void selfTest() {
    check(githubRepository("git@github.com:DataDog/datadog-intellij-plugin.git")
        .equals("DataDog/datadog-intellij-plugin"));
    check(githubRepository("https://github.com/DataDog/datadog-intellij-plugin.git")
        .equals("DataDog/datadog-intellij-plugin"));
    check(parsePull("2307", "DataDog/datadog-intellij-plugin")
        .equals(new PullSpec("DataDog/datadog-intellij-plugin", 2307)));
    check(parsePull("https://github.com/DataDog/datadog-intellij-plugin/pull/2307/files",
        "DataDog/datadog-intellij-plugin")
        .equals(new PullSpec("DataDog/datadog-intellij-plugin", 2307)));
    check(slug("IDE-3813 Experiment to configure Docker files ")
        .equals("ide-3813-experiment-to-configure-docker-files"));
    String media = "![demo](https://github.com/user-attachments/assets/01234567-89ab-cdef-0123-456789abcdef)";
    check(urls(ATTACHMENT_URL, media).size() == 1);
    check(assetKey(urls(ATTACHMENT_URL, media).getFirst())
        .equals("01234567-89ab-cdef-0123-456789abcdef"));
    String fenced = "```markdown\n" + media + "\n```\n" + media;
    check(replaceOutsideFences(fenced, "https://github.com", "./local").equals(
        "```markdown\n" + media + "\n```\n![demo](./local/user-attachments/assets/01234567-89ab-cdef-0123-456789abcdef)"));
    check(replaceOutsideFences("https://example.com/demo.mov\n", "https://example.com/demo.mov", "./demo.mov")
        .equals("[demo.mov](./demo.mov)\n"));
    System.out.println("Self-test passed");
  }

  private static void check(boolean condition) {
    if (!condition) {
      throw new AssertionError();
    }
  }

  private record PullRequest(
      String repository,
      int number,
      String title,
      String body,
      String bodyHtml,
      String branch,
      String sha) {}

  private record PullSpec(String repository, int number) {}

  private record RenderedMedia(String html, String url) {}
}