const GENERATED = [
  /^mlflow\/protos\/.*\.pyi$/,
  /_pb2(_grpc)?\.py$/,
  /^uv\.lock$/,
  /package-lock\.json$/,
  /yarn\.lock$/,
  /^docs\/api_reference\/source\/rest-api\.rst$/,
];

const THRESHOLDS = {
  XS: 9,
  S: 49,
  M: 199,
  L: 499,
  XL: Infinity,
};

async function isGenerated(owner, repo, sha, filename) {
  if (GENERATED.some((p) => p.test(filename))) return true;
  if (filename.endsWith(".java")) {
    try {
      const resp = await fetch(
        `https://raw.githubusercontent.com/${owner}/${repo}/${sha}/${filename}`,
        { headers: { Range: "bytes=0-255", Authorization: `token ${process.env.GH_TOKEN}` } }
      );
      if (resp.ok) {
        const text = await resp.text();
        if (text.includes("Generated by the protocol buffer compiler")) return true;
      }
    } catch (e) {
      console.warn(filename, e.message);
    }
  }
  return false;
}

function getSize(total) {
  return Object.entries(THRESHOLDS).find(([, max]) => total <= max)[0];
}

function extractStackedPrBaseSha(body, headRef) {
  if (!body || !body.includes("Stacked PR")) return null;
  const marker = `[**${headRef}**]`;
  for (const line of body.split("\n")) {
    if (line.includes(marker)) {
      const match = line.match(/\/files\/([a-f0-9]{7,40})\.\.([a-f0-9]{7,40})/);
      if (match) return match[1];
    }
  }
  return null;
}

async function getFiles(github, owner, repo, pr) {
  const baseSha = extractStackedPrBaseSha(pr.body, pr.head.ref);
  if (baseSha) {
    console.log(`Stacked PR detected, using incremental diff: ${baseSha}..${pr.head.sha}`);
    try {
      const resp = await github.rest.repos.compareCommits({
        owner,
        repo,
        base: baseSha,
        head: pr.head.sha,
      });
      return resp.data.files;
    } catch (e) {
      console.warn(`Failed to fetch incremental diff, falling back to full diff: ${e.message}`);
    }
  }
  return github.paginate(github.rest.pulls.listFiles, {
    owner,
    repo,
    pull_number: pr.number,
    per_page: 100,
  });
}

module.exports = async ({ github, context }) => {
  const { owner, repo } = context.repo;
  const pr = context.payload.pull_request;
  const headSha = pr.head.sha;
  const files = await getFiles(github, owner, repo, pr);

  let total = 0;
  const maxThreshold = Math.max(...Object.values(THRESHOLDS).filter(isFinite));
  for (const f of files) {
    if (!(await isGenerated(owner, repo, headSha, f.filename))) {
      total += f.additions + f.deletions;
    }
    if (total > maxThreshold) break;
  }

  const sizeLabel = `size/${getSize(total)}`;
  console.log(`Size: ${total} lines -> ${sizeLabel}`);

  // Manage labels
  const currentLabels = (
    await github.paginate(github.rest.issues.listLabelsOnIssue, {
      owner,
      repo,
      issue_number: pr.number,
    })
  ).map((l) => l.name);

  // Remove stale size labels
  for (const label of currentLabels) {
    if (label.startsWith("size/") && label !== sizeLabel) {
      console.log(`Removing stale label: ${label}`);
      await github.rest.issues
        .removeLabel({ owner, repo, issue_number: pr.number, name: label })
        .catch((e) => console.warn(`Failed to remove label ${label}: ${e.message}`));
    }
  }

  // Add the correct label if not already present
  if (!currentLabels.includes(sizeLabel)) {
    try {
      await github.rest.issues.getLabel({ owner, repo, name: sizeLabel });
    } catch (e) {
      if (e.status !== 404) throw e;
      console.log(`Creating label: ${sizeLabel}`);
      await github.rest.issues.createLabel({
        owner,
        repo,
        name: sizeLabel,
        color: "ededed",
        description: `Pull request size: ${sizeLabel.replace("size/", "")}`,
      });
    }
    await github.rest.issues.addLabels({
      owner,
      repo,
      issue_number: pr.number,
      labels: [sizeLabel],
    });
  }
};