import { type FileDiffContent } from "code-suggester/build/src/types"
import { type Octokit } from "octokit"
import fs from "fs"
import { addLabelToPullRequest, buildDependentPrTitle, buildPrTitle } from "./github-utils"
import { logger } from "@azure/identity"

type ReportTable = Record<
  string,
  {
    passed: number
    failed: number
  }
>

const PR_HEADER_TEMPLATE = fs.readFileSync("github/optimization_results_header.md", "utf8")
const PR_DETAILS_TEMPLATE = fs.readFileSync("github/optimization_results_details.md", "utf8")
const PR_TEST_REPORT_TEMPLATE = fs.readFileSync(
  "github/optimization_results_test_report.md",
  "utf8",
)

type PrCommentFields = Record<string, any>

export function buildResultHeader(fields: PrCommentFields): string {
  return PR_HEADER_TEMPLATE.replace(/\{best_runtime}/g, fields.best_runtime)
    .replace(/\{original_runtime}/g, fields.original_runtime)
    .replace(/\{function_name}/g, fields.function_name)
    .replace(/\{file_path}/g, fields.file_path)
    .replace(/\{speedup_x}/g, fields.speedup_x)
    .replace(/\{speedup_pct}/g, fields.speedup_pct)
    .replace(/\{num_runs}/g, fields.loop_count)
}

export function buildResultDetails(fields: PrCommentFields): string {
  return (
    PR_DETAILS_TEMPLATE.replace(/\{optimization_explanation}/g, fields.optimization_explanation) +
    "\n"
  )
}

export function buildResultTestReport(
  fields: PrCommentFields,
  existingTests: string,
  generatedTests: string,
  coverage_message: string,
): string {
  const reportTableDict: ReportTable = fields.report_table

  let reportTableMd: string = ""
  // Initialize the markdown table
  reportTableMd += '| Test                        | Status            |\n';
  reportTableMd += '| --------------------------- | ----------------- |\n';

  // Loop through each test type and construct the table rows
  for (const [testType, { passed = 0, failed = 0 }] of Object.entries(reportTableDict)) {
    if (testType.includes('🎨')) {
      continue;
    }
    let status = '';
    let detailsNote = '';

    // Determine the status based on passed counts
   if (passed > 0) {
      status = `✅ **${passed} Passed**`;
    } else {
      status = '🔘 **None Found**';
    }

    // If there are details, note that they are available below
    if ((passed > 0) && (testType.includes('Existing') || testType.includes('Generated'))) {
      detailsNote = 'See below';
    } else {
      detailsNote = '';
    }

    // Add the row to the markdown table
    reportTableMd += `| ${testType} | ${status} |\n`;
  }
  reportTableMd += `|📊 Tests Coverage       | ${coverage_message} |\n`;


  // Add detailed test outputs below the table
  for (const [testType, { passed = 0, failed = 0 }] of Object.entries(reportTableDict)) {
    // Only include details for Existing and Generated tests with results
    if ((passed > 0) && (testType.includes('Existing') || testType.includes('Generated'))) {
      // Add a heading for the test type details

      reportTableMd += `<details>\n`;
      reportTableMd += `<summary>${testType} Details</summary>\n\n`;

      // Include the relevant test code
      if (testType.includes('Existing')) {
        reportTableMd += '```python\n';
        reportTableMd += existingTests.trim();
        reportTableMd += '\n```\n';
      } else if (testType.includes('Generated')) {
        reportTableMd += '```python\n';
        reportTableMd += generatedTests.trim();
        reportTableMd += '\n```\n';
      } else {
        reportTableMd += '_No additional details available._\n';
      }

      reportTableMd += `\n</details>\n\n`;
    }
  }
  // Add the final markdown content (e.g., the feedback section)
  const finalMarkdown = `
${reportTableMd}

:loudspeaker: Feedback on this optimization? [![Discord](https://img.shields.io/badge/Discord-Join%20Our%20Community-7289DA)](https://codeflash.ai/discord)
`;

  return PR_TEST_REPORT_TEMPLATE.replace(/\{report_table}/g, finalMarkdown)
}

export async function createNewBranchFromDiffContents(
  installationOctokit: Octokit,
  owner: string,
  repo: string,
  newBranchName: string,
  baseBranch: string,
  diffContentsMap: Map<string, FileDiffContent>,
  commitMessage: string,
): Promise<boolean> {
  // 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,
  })

  // Apply changes from diffContentsMap to the new branch and commit
  const treeItems: Octokit.GitCreateTreeParamsTree[] = []
  for (const [filePath, fileDiffContent] of diffContentsMap) {
    // 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
}

async function createNewPullRequest(
  installationOctokit: Octokit,
  owner: string,
  repo: string,
  title: string,
  body: string,
  newBranchName: string,
  baseBranch: string,
): Promise<Octokit.Response<Octokit.PullsCreateResponse>> {
  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: Octokit,
  owner: string,
  repo: string,
  newBranchName: string,
  baseBranch: string,
  prCommentFields: PrCommentFields,
  existingTests: string,
  generatedTests: string,
  coverage_message: string,
): Promise<Octokit.Response<Octokit.PullsCreateResponse>> {
  // Open a new standalone Codeflash PR (that doesn't reference an original PR, likely just to main)
  const prCommentHeader = buildResultHeader(prCommentFields)
  // Open a new PR from the new branch onto the original PR's head branch
  const prCommentBody = buildResultDetails(prCommentFields)
  const prCommentTestReport = buildResultTestReport(prCommentFields, existingTests, generatedTests, coverage_message)

  const title: string = buildPrTitle(prCommentFields.function_name, prCommentFields.speedup_pct)
  const body: string = `${prCommentHeader}\n` + `${prCommentBody}\n` + `${prCommentTestReport}`

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

export async function createDependentPullRequest(
  installationOctokit: Octokit,
  owner: string,
  repo: string,
  origPrNumber: number,
  newBranchName: string,
  baseBranch: string,
  prCommentFields: PrCommentFields,
  existingTests: string,
  generatedTests: string,
): Promise<Octokit.Response<Octokit.PullsCreateResponse>> {
  const prCommentHeader = buildResultHeader(prCommentFields)
  // Open a new PR from the new branch onto the original PR's head branch
  const prCommentBody = buildResultDetails(prCommentFields)
  const prCommentTestReport = buildResultTestReport(prCommentFields, existingTests, generatedTests, coverage_message)

  const title = buildDependentPrTitle(
    prCommentFields.function_name,
    prCommentFields.speedup_pct,
    origPrNumber,
    baseBranch,
  )
  const body =
    `## ⚡️ 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` +
    `${prCommentHeader}\n` +
    `${prCommentBody}\n` +
    `${prCommentTestReport}`
  const newPrData = await createNewPullRequest(
    installationOctokit,
    owner,
    repo,
    title,
    body,
    newBranchName,
    baseBranch,
  )
  await addLabelToPullRequest(installationOctokit, owner, repo, newPrData.data.number)

  // Make a comment on the original PR with a link to the new PR
  const result = await installationOctokit.rest.issues.createComment({
    owner,
    repo,
    issue_number: origPrNumber,
    body: `## ⚡️ Codeflash found optimizations for this PR
${prCommentHeader}
#### I created a new dependent PR with the suggested changes. Please review: 

- ### #${newPrData.data.number}

If you approve, it will be merged into this PR (branch \`${baseBranch}\`).`,
  })

  return newPrData
}
