codeflash-internal/js/cf-api/github/create-pr-from-diffcontents.ts

import type { FileDiffContent } from "@codeflash-ai/code-suggester/build/src/types.js"
import type { Octokit } from "octokit"
import { addLabelToPullRequest } from "./github-utils.js"
import {
  buildBenchmarkInfo,
  buildDependentPrTitle,
  buildPrTitle,
  buildResultFooter,
  generateOptimizationReviewTemplate,
  originalPRComment,
  buildResultHeader,
  buildResultDetails,
  buildResultTestReport,
} from "./pr-changes-utils.js"
import type { RestEndpointMethodTypes } from "@octokit/rest"
import { AnyOctokit, PullRequestCreationResponse } from "../types.js"
import * as Sentry from "@sentry/node"

type GitCreateTreeParamsTree =
  RestEndpointMethodTypes["git"]["createTree"]["parameters"]["tree"][number]

interface BenchmarkDetail {
  benchmark_name: string
  test_function: string
  original_timing: number
  expected_new_timing: number
  speedup_percent: number
}

export interface PrCommentFields {
  best_runtime: string
  original_runtime: string
  function_name: string
  file_path: string
  speedup_x: string
  speedup_pct: string
  loop_count: string
  optimization_explanation: string
  report_table: any
  benchmark_details?: BenchmarkDetail[]
  original_async_throughput?: string
  best_async_throughput?: string
}

// Dependencies interface for easier testing
export interface CreatePrFromDiffContentsDependencies {
  addLabelToPullRequest: typeof addLabelToPullRequest
  buildBenchmarkInfo: typeof buildBenchmarkInfo
  buildDependentPrTitle: typeof buildDependentPrTitle
  buildPrTitle: typeof buildPrTitle
  buildResultFooter: typeof buildResultFooter
  buildResultHeader: typeof buildResultHeader
  buildResultDetails: typeof buildResultDetails
  buildResultTestReport: typeof buildResultTestReport
  originalPRComment: typeof originalPRComment
}

// Default dependencies
let dependencies: CreatePrFromDiffContentsDependencies = {
  addLabelToPullRequest,
  buildBenchmarkInfo,
  buildDependentPrTitle,
  buildPrTitle,
  buildResultFooter,
  buildResultHeader,
  buildResultDetails,
  buildResultTestReport,
  originalPRComment,
}

// For testing - allow dependency injection
export function setCreatePrFromDiffContentsDependencies(
  deps: Partial<CreatePrFromDiffContentsDependencies>,
) {
  dependencies = { ...dependencies, ...deps }
}

export function resetCreatePrFromDiffContentsDependencies() {
  dependencies = {
    addLabelToPullRequest,
    buildBenchmarkInfo,
    buildDependentPrTitle,
    buildPrTitle,
    buildResultFooter,
    buildResultHeader,
    buildResultDetails,
    buildResultTestReport,
    originalPRComment,
  }
}

export async function createNewBranchFromDiffContents(
  installationOctokit: Octokit,
  owner: string,
  repo: string,
  newBranchName: string,
  baseBranch: string,
  diffContentsMap: Map<string, FileDiffContent>,
  commitMessage: string,
): Promise<boolean> {
  try {
    // Create a new branch off of the specified base branch
    await installationOctokit.rest.git.createRef({
      owner,
      repo,
      ref: `refs/heads/${newBranchName}`,
      sha: (
        await installationOctokit.rest.git.getRef({
          owner,
          repo,
          ref: `heads/${baseBranch}`,
        })
      ).data.object.sha,
    })

    const treeItems: GitCreateTreeParamsTree[] = []
    for (const [filePath, fileDiffContent] of diffContentsMap) {
      try {
        await installationOctokit.rest.repos.getContent({
          owner,
          repo,
          path: filePath,
          ref: baseBranch,
        })
      } catch (err: any) {
        if (err.status === 404) {
          const error = new Error(
            `File does not exist in base branch: ${filePath} in ${owner}/${repo}`,
          )
          Sentry.captureException(error)
          throw error
        }
        throw err
      }

      // Create blobs for the changed files with newContent
      const blobData = await installationOctokit.rest.git.createBlob({
        owner,
        repo,
        content: fileDiffContent.newContent,
        encoding: "utf-8",
      })

      treeItems.push({
        path: filePath,
        mode: "100644", // blob
        type: "blob",
        sha: blobData.data.sha,
      })
    }

    // Get the current commit SHA of the new branch
    const refData = await installationOctokit.rest.git.getRef({
      owner,
      repo,
      ref: `heads/${newBranchName}`,
    })
    const currentCommitSha = refData.data.object.sha

    // Get the tree SHA of the current commit
    const commitData = await installationOctokit.rest.git.getCommit({
      owner,
      repo,
      commit_sha: currentCommitSha,
    })
    const currentTreeSha = commitData.data.tree.sha

    // Create a new tree with the changes
    const newTreeData = await installationOctokit.rest.git.createTree({
      owner,
      repo,
      base_tree: currentTreeSha,
      tree: treeItems,
    })

    // Create a new commit with the new tree
    const newCommitData = await installationOctokit.rest.git.createCommit({
      owner,
      repo,
      message: commitMessage,
      tree: newTreeData.data.sha,
      parents: [currentCommitSha],
    })

    // Update the new branch reference to point to the new commit
    const result = await installationOctokit.rest.git.updateRef({
      owner,
      repo,
      ref: `heads/${newBranchName}`,
      sha: newCommitData.data.sha,
    })

    return result.status === 200
  } catch (error) {
    console.error("Error creating branch from diff contents:", error)
    Sentry.captureException(`Failed to create branch: ${error.message}`, {
      extra: { owner, repo, newBranchName, baseBranch, commitMessage, diffContentsMap },
    })
    return false
  }
}

/**
 * Create a branch with a single new file using Git Tree API
 * This supports creating new files (unlike createNewBranchFromDiffContents which requires files to exist)
 */
export async function createBranchWithNewFile(
  installationOctokit: Octokit,
  owner: string,
  repo: string,
  newBranchName: string,
  baseBranch: string,
  filePath: string,
  fileContent: string,
  commitMessage: string,
): Promise<boolean> {
  try {
    // Step 1: Get base branch SHA
    const baseRef = await installationOctokit.rest.git.getRef({
      owner,
      repo,
      ref: `heads/${baseBranch}`,
    })
    const baseSha = baseRef.data.object.sha

    // Step 2: Create new branch
    await installationOctokit.rest.git.createRef({
      owner,
      repo,
      ref: `refs/heads/${newBranchName}`,
      sha: baseSha,
    })

    // Step 3: Create blob for the file content
    const blobData = await installationOctokit.rest.git.createBlob({
      owner,
      repo,
      content: fileContent,
      encoding: "utf-8",
    })

    // Step 4: Get current commit and tree from the branch
    const refData = await installationOctokit.rest.git.getRef({
      owner,
      repo,
      ref: `heads/${newBranchName}`,
    })
    const currentCommitSha = refData.data.object.sha

    const commitData = await installationOctokit.rest.git.getCommit({
      owner,
      repo,
      commit_sha: currentCommitSha,
    })
    const currentTreeSha = commitData.data.tree.sha

    // Step 5: Create new tree with the file
    // Using base_tree automatically preserves all existing files - we just add/update our file
    const newTreeData = await installationOctokit.rest.git.createTree({
      owner,
      repo,
      base_tree: currentTreeSha,
      tree: [
        {
          path: filePath,
          mode: "100644", // Regular file
          type: "blob",
          sha: blobData.data.sha,
        },
      ],
    })

    // Step 6: Create new commit
    const newCommitData = await installationOctokit.rest.git.createCommit({
      owner,
      repo,
      message: commitMessage,
      tree: newTreeData.data.sha,
      parents: [currentCommitSha],
    })

    // Step 7: Update branch to point to new commit
    const result = await installationOctokit.rest.git.updateRef({
      owner,
      repo,
      ref: `heads/${newBranchName}`,
      sha: newCommitData.data.sha,
    })

    return result.status === 200
  } catch (error) {
    console.error("Error creating branch with new file:", error)
    Sentry.captureException(error)
    return false
  }
}

export async function createNewPullRequest(
  installationOctokit: AnyOctokit,
  owner: string,
  repo: string,
  title: string,
  body: string,
  newBranchName: string,
  baseBranch: string,
): Promise<PullRequestCreationResponse> {
  return await installationOctokit.rest.pulls.create({
    owner,
    repo,
    title,
    head: newBranchName,
    base: baseBranch,
    body,
  })
}

export async function assignReviewer(
  installationOctokit: Octokit,
  owner: string,
  repo: string,
  pullNumber: number,
  reviewer: string,
): Promise<void> {
  await installationOctokit.rest.pulls.requestReviewers({
    owner,
    repo,
    pull_number: pullNumber,
    reviewers: [reviewer],
  })
}

export async function createStandalonePullRequest(
  installationOctokit: AnyOctokit,
  owner: string,
  repo: string,
  newBranchName: string,
  baseBranch: string,
  prCommentFields: PrCommentFields,
  existingTests: string,
  generatedTests: string,
  coverage_message: string,
  replayTests: string = "",
  concolicTests: string = "",
): Promise<PullRequestCreationResponse> {
  // Open a new standalone Codeflash PR (that doesn't reference an original PR, likely just to main)
  const prCommentHeader = dependencies.buildResultHeader(prCommentFields)
  // Build benchmark info if available
  const benchmarkInfo =
    prCommentFields.benchmark_details && prCommentFields.benchmark_details.length > 0
      ? dependencies.buildBenchmarkInfo(prCommentFields)
      : ""
  // Open a new PR from the new branch onto the original PR's head branch
  const prCommentBody = dependencies.buildResultDetails(prCommentFields)
  const prCommentTestReport = dependencies.buildResultTestReport(
    prCommentFields,
    existingTests,
    generatedTests,
    coverage_message,
    replayTests,
    concolicTests,
  )
  const prCommentFooter = dependencies.buildResultFooter(newBranchName)
  const title: string = dependencies.buildPrTitle(
    prCommentFields.function_name,
    prCommentFields.speedup_pct,
    prCommentFields.speedup_x,
  )
  const body: string = benchmarkInfo
    ? `${prCommentHeader}\n${benchmarkInfo}\n${prCommentBody}\n${prCommentTestReport}\n${prCommentFooter}`
    : `${prCommentHeader}\n${prCommentBody}\n${prCommentTestReport}\n${prCommentFooter}`

  return await createNewPullRequest(
    installationOctokit,
    owner,
    repo,
    title,
    body,
    newBranchName,
    baseBranch,
  )
}

export async function createDependentPullRequest(
  installationOctokit: AnyOctokit,
  owner: string,
  repo: string,
  origPrNumber: number,
  newBranchName: string,
  baseBranch: string,
  prCommentFields: PrCommentFields,
  existingTests: string,
  generatedTests: string,
  coverage_message: string,
  replayTests: string = "",
  concolicTests: string = "",
  optimizationReview: string = "",
): Promise<PullRequestCreationResponse> {
  const { title, body } = createDependentPRTitleAndBody(
    prCommentFields,
    existingTests,
    generatedTests,
    coverage_message,
    origPrNumber,
    newBranchName,
    baseBranch,
    replayTests,
    concolicTests,
    optimizationReview,
  )

  const newPrData = await createNewPullRequest(
    installationOctokit,
    owner,
    repo,
    title,
    body,
    newBranchName,
    baseBranch,
  )

  await dependencies.addLabelToPullRequest(installationOctokit, owner, repo, newPrData.data.number)

  await dependencies.addLabelToPullRequest(
    installationOctokit,
    owner,
    repo,
    newPrData.data.number,
    `🎯 Quality: ${optimizationReview.charAt(0).toUpperCase() + optimizationReview.slice(1).toLowerCase()}`,
    "FFC043",
    "Optimization Quality according to codeflash",
  )

  // Make a comment on the original PR with a link to the new PR
  const commentBody: string = dependencies.originalPRComment(
    prCommentFields,
    newPrData.data.number,
    baseBranch,
    optimizationReview,
  )
  await installationOctokit.rest.issues.createComment({
    owner,
    repo,
    issue_number: origPrNumber,
    body: commentBody,
  })

  return newPrData
}

function createDependentPRTitleAndBody(
  prCommentFields: PrCommentFields,
  existingTests: string,
  generatedTests: string,
  coverage_message: string,
  origPrNumber: number,
  newBranchName: string,
  baseBranch: string,
  replayTests: string = "",
  concolicTests: string = "",
  optimizationReview: string = "",
): { title: string; body: string } {
  const prCommentHeader = dependencies.buildResultHeader(prCommentFields)
  // Build benchmark info if available
  const benchmarkInfo =
    prCommentFields.benchmark_details && prCommentFields.benchmark_details.length > 0
      ? dependencies.buildBenchmarkInfo(prCommentFields)
      : ""
  // Open a new PR from the new branch onto the original PR's head branch
  const prCommentBody = dependencies.buildResultDetails(prCommentFields)
  const prCommentTestReport = dependencies.buildResultTestReport(
    prCommentFields,
    existingTests,
    generatedTests,
    coverage_message,
    replayTests,
    concolicTests,
  )
  const prCommentFooter = dependencies.buildResultFooter(newBranchName)
  const title = dependencies.buildDependentPrTitle(
    prCommentFields.function_name,
    prCommentFields.speedup_pct,
    prCommentFields.speedup_x,
    origPrNumber,
    baseBranch,
  )
  const introSection =
    `## ⚡️ This pull request contains optimizations for PR #${origPrNumber}
If you approve this dependent PR, these changes will be merged into the original PR branch \`${baseBranch}\`.
>This PR will be automatically closed if the original PR is merged.\n` + `----\n`
  let optReviewBadge = generateOptimizationReviewTemplate(optimizationReview)
  optReviewBadge &&= ` ${optReviewBadge}\n`
  // Conditionally build the body based on whether benchmark info exists
  const body = benchmarkInfo
    ? `${introSection}${prCommentHeader}\n${benchmarkInfo}\n${prCommentBody}\n${prCommentTestReport}\n${prCommentFooter}${optReviewBadge}`
    : `${introSection}${prCommentHeader}\n${prCommentBody}\n${prCommentTestReport}\n${prCommentFooter}${optReviewBadge}`

  return { title, body }
}