// Fill in the middle prompts

import { CompletionOptions } from "../../index.js"
import { getLastNUriRelativePathParts, getShortestUniqueRelativeUriPaths } from "../../util/uri.js"
import { AutocompleteCodeSnippet, AutocompleteSnippet, AutocompleteSnippetType } from "../types.js"

type TemplateRenderer = (
  prefix: string,
  suffix: string,
  filepath: string,
  reponame: string,
  language: string,
  snippets: AutocompleteSnippet[],
  workspaceUris: string[],
) => string

export interface AutocompleteTemplate {
  compilePrefixSuffix?: (
    prefix: string,
    suffix: string,
    filepath: string,
    reponame: string,
    snippets: AutocompleteSnippet[],
    workspaceUris: string[],
  ) => [string, string]
  template: TemplateRenderer
  completionOptions?: Partial<CompletionOptions>
}

// https://huggingface.co/stabilityai/stable-code-3b
const stableCodeFimTemplate: AutocompleteTemplate = {
  template: (prefix: string, suffix: string): string => {
    return `<fim_prefix>${prefix}<fim_suffix>${suffix}<fim_middle>`
  },
  completionOptions: {
    stop: ["<fim_prefix>", "<fim_suffix>", "<fim_middle>", "<file_sep>", "<|endoftext|>", "</fim_middle>", "</code>"],
  },
}

// https://github.com/QwenLM/Qwen2.5-Coder?tab=readme-ov-file#3-file-level-code-completion-fill-in-the-middle
// This issue asks about the use of <|repo_name|> and <|file_sep|> together with <|fim_prefix|>, <|fim_suffix|> and <|fim_middle|>
// https://github.com/QwenLM/Qwen2.5-Coder/issues/343
const qwenCoderFimTemplate: AutocompleteTemplate = {
  template: (prefix: string, suffix: string): string => {
    return `<|fim_prefix|>${prefix}<|fim_suffix|>${suffix}<|fim_middle|>`
  },
  completionOptions: {
    stop: [
      "<|endoftext|>",
      "<|fim_prefix|>",
      "<|fim_middle|>",
      "<|fim_suffix|>",
      "<|fim_pad|>",
      "<|repo_name|>",
      "<|file_sep|>",
      "<|im_start|>",
      "<|im_end|>",
    ],
  },
}

const seedCoderFimTemplate: AutocompleteTemplate = {
  template: (prefix: string, suffix: string): string => {
    return `<[fim-prefix]>${prefix}<[fim-suffix]>${suffix}<[fim-middle]>`
  },
  completionOptions: {
    stop: [
      "<[end▁of▁sentence]>",
      "<[fim-prefix]>",
      "<[fim-middle]>",
      "<[fim-suffix]>",
      "<[PAD▁TOKEN]>",
      "<[SEP▁TOKEN]>",
      "<[begin▁of▁sentence]>",
    ],
  },
}

const codestralMultifileFimTemplate: AutocompleteTemplate = {
  compilePrefixSuffix: (prefix, suffix, filepath, _reponame, snippets, workspaceUris): [string, string] => {
    function getFileName(snippet: { uri: string; uniquePath: string }) {
      return snippet.uri.startsWith("file://") ? snippet.uniquePath : snippet.uri
    }

    if (snippets.length === 0) {
      if (suffix.trim().length === 0 && prefix.trim().length === 0) {
        return [`+++++ ${getLastNUriRelativePathParts(workspaceUris, filepath, 2)}\n${prefix}`, suffix]
      }
      return [prefix, suffix]
    }

    const relativePaths = getShortestUniqueRelativeUriPaths(
      [...snippets.map((snippet) => ("filepath" in snippet ? snippet.filepath : "file:///Untitled.txt")), filepath],
      workspaceUris,
    )

    const otherFiles = snippets
      .map((snippet, i) => {
        if (snippet.type === AutocompleteSnippetType.Diff) {
          return snippet.content
        }

        return `+++++ ${getFileName(relativePaths[i])} \n${snippet.content}`
      })
      .join("\n\n")

    return [`${otherFiles}\n\n+++++ ${getFileName(relativePaths[relativePaths.length - 1])}\n${prefix}`, suffix]
  },
  template: (prefix: string, suffix: string): string => {
    return `[SUFFIX]${suffix}[PREFIX]${prefix}`
  },
  completionOptions: {
    stop: ["[PREFIX]", "[SUFFIX]", "\n+++++ "],
  },
}

const mercuryMultifileFimTemplate: AutocompleteTemplate = {
  compilePrefixSuffix: (prefix, suffix, filepath, _reponame, snippets, workspaceUris): [string, string] => {
    function getFileName(snippet: { uri: string; uniquePath: string }) {
      return snippet.uri.startsWith("file://") ? snippet.uniquePath : snippet.uri
    }

    // Our current snippet format doesn't work well with mercury. We need to clean this up
    snippets = []

    if (snippets.length === 0) {
      if (suffix.trim().length === 0 && prefix.trim().length === 0) {
        return [
          `<|file_sep|>${getLastNUriRelativePathParts(workspaceUris, filepath, 2)}\n<|fim_prefix|>${prefix}`,
          suffix,
        ]
      }
      return [`<|fim_prefix|>${prefix}`, suffix]
    }

    const relativePaths = getShortestUniqueRelativeUriPaths(
      [...snippets.map((snippet) => ("filepath" in snippet ? snippet.filepath : "file:///Untitled.txt")), filepath],
      workspaceUris,
    )

    const otherFiles = snippets
      .map((snippet, i) => {
        if (snippet.type === AutocompleteSnippetType.Diff) {
          return snippet.content
        }

        return `<|file_sep|>${getFileName(relativePaths[i])} \n${snippet.content}`
      })
      .join("\n\n")

    return [
      `${otherFiles}${otherFiles ? "\n\n" : ""}<|file_sep|>${getFileName(relativePaths[relativePaths.length - 1])}\n<|fim_prefix|>${prefix}`,
      suffix,
    ]
  },
  template: (prefix: string, suffix: string): string => {
    return `${prefix}<|fim_suffix|>${suffix}<|fim_middle|>`
  },
}

const codegemmaFimTemplate: AutocompleteTemplate = {
  template: (prefix: string, suffix: string): string => {
    return `<|fim_prefix|>${prefix}<|fim_suffix|>${suffix}<|fim_middle|>`
  },
  completionOptions: {
    stop: ["<|fim_prefix|>", "<|fim_suffix|>", "<|fim_middle|>", "<|file_separator|>", "<end_of_turn>", "<eos>"],
  },
}

const codeLlamaFimTemplate: AutocompleteTemplate = {
  template: (prefix: string, suffix: string): string => {
    return `<PRE> ${prefix} <SUF>${suffix} <MID>`
  },
  completionOptions: { stop: ["<PRE>", "<SUF>", "<MID>", "<EOT>"] },
}

// https://huggingface.co/deepseek-ai/deepseek-coder-1.3b-base
const deepseekFimTemplate: AutocompleteTemplate = {
  template: (prefix: string, suffix: string): string => {
    return `<｜fim▁begin｜>${prefix}<｜fim▁hole｜>${suffix}<｜fim▁end｜>`
  },
  completionOptions: {
    stop: ["<｜fim▁begin｜>", "<｜fim▁hole｜>", "<｜fim▁end｜>", "//", "<｜end▁of▁sentence｜>"],
  },
}

// https://github.com/THUDM/CodeGeeX4/blob/main/guides/Infilling_guideline.md
const codegeexFimTemplate: AutocompleteTemplate = {
  template: (prefix, suffix, filepath, _reponame, language, allSnippets, workspaceUris): string => {
    const snippets = allSnippets.filter(
      (snippet) => snippet.type === AutocompleteSnippetType.Code,
    ) as AutocompleteCodeSnippet[]

    const relativePaths = getShortestUniqueRelativeUriPaths(
      [...snippets.map((snippet) => snippet.filepath), filepath],
      workspaceUris,
    )
    const baseTemplate = `###PATH:${
      relativePaths[relativePaths.length - 1]
    }\n###LANGUAGE:${language}\n###MODE:BLOCK\n<|code_suffix|>${suffix}<|code_prefix|>${prefix}<|code_middle|>`
    if (snippets.length === 0) {
      return `<|user|>\n${baseTemplate}<|assistant|>\n`
    }
    const references = `###REFERENCE:\n${snippets
      .map((snippet, i) => `###PATH:${relativePaths[i]}\n${snippet.content}\n`)
      .join("###REFERENCE:\n")}`
    const prompt = `<|user|>\n${references}\n${baseTemplate}<|assistant|>\n`
    return prompt
  },
  completionOptions: {
    stop: ["<|user|>", "<|code_suffix|>", "<|code_prefix|>", "<|code_middle|>", "<|assistant|>", "<|endoftext|>"],
  },
}

const holeFillerTemplate: AutocompleteTemplate = {
  template: (prefix: string, suffix: string) => {
    // From https://github.com/VictorTaelin/AI-scripts
    const SYSTEM_MSG = `You are a HOLE FILLER. You are provided with a file containing holes, formatted as '{{HOLE_NAME}}'. Your TASK is to complete with a string to replace this hole with, inside a <COMPLETION/> XML tag, including context-aware indentation, if needed.  All completions MUST be truthful, accurate, well-written and correct.

## EXAMPLE QUERY:

<QUERY>
function sum_evens(lim) {
  var sum = 0;
  for (var i = 0; i < lim; ++i) {
    {{FILL_HERE}}
  }
  return sum;
}
</QUERY>

TASK: Fill the {{FILL_HERE}} hole.

## CORRECT COMPLETION

<COMPLETION>if (i % 2 === 0) {
      sum += i;
    }</COMPLETION>

## EXAMPLE QUERY:

<QUERY>
def sum_list(lst):
  total = 0
  for x in lst:
  {{FILL_HERE}}
  return total

print sum_list([1, 2, 3])
</QUERY>

## CORRECT COMPLETION:

<COMPLETION>  total += x</COMPLETION>

## EXAMPLE QUERY:

<QUERY>
// data Tree a = Node (Tree a) (Tree a) | Leaf a

// sum :: Tree Int -> Int
// sum (Node lft rgt) = sum lft + sum rgt
// sum (Leaf val)     = val

// convert to TypeScript:
{{FILL_HERE}}
</QUERY>

## CORRECT COMPLETION:

<COMPLETION>type Tree<T>
  = {$:"Node", lft: Tree<T>, rgt: Tree<T>}
  | {$:"Leaf", val: T};

function sum(tree: Tree<number>): number {
  switch (tree.$) {
    case "Node":
      return sum(tree.lft) + sum(tree.rgt);
    case "Leaf":
      return tree.val;
  }
}</COMPLETION>

## EXAMPLE QUERY:

The 5th {{FILL_HERE}} is Jupiter.

## CORRECT COMPLETION:

<COMPLETION>planet from the Sun</COMPLETION>

## EXAMPLE QUERY:

function hypothenuse(a, b) {
  return Math.sqrt({{FILL_HERE}}b ** 2);
}

## CORRECT COMPLETION:

<COMPLETION>a ** 2 + </COMPLETION>`

    const fullPrompt =
      SYSTEM_MSG +
      `\n\n<QUERY>\n${prefix}{{FILL_HERE}}${suffix}\n</QUERY>\nTASK: Fill the {{FILL_HERE}} hole. Answer only with the CORRECT completion, and NOTHING ELSE. Do it now.\n<COMPLETION>`
    return fullPrompt
  },
  completionOptions: {
    stop: ["</COMPLETION>"],
  },
}

export function getTemplateForModel(model: string): AutocompleteTemplate {
  const lowerCaseModel = model.toLowerCase()

  // if (lowerCaseModel.includes("starcoder2")) {
  //   return starcoder2FimTemplate;
  // }
  if (lowerCaseModel.includes("mercury")) {
    return mercuryMultifileFimTemplate
  }

  if (lowerCaseModel.includes("qwen") && lowerCaseModel.includes("coder")) {
    return qwenCoderFimTemplate
  }

  if (lowerCaseModel.includes("seed") && lowerCaseModel.includes("coder")) {
    return seedCoderFimTemplate
  }

  if (
    lowerCaseModel.includes("starcoder") ||
    lowerCaseModel.includes("star-coder") ||
    lowerCaseModel.includes("starchat") ||
    lowerCaseModel.includes("octocoder") ||
    lowerCaseModel.includes("stable") ||
    lowerCaseModel.includes("codeqwen") ||
    lowerCaseModel.includes("qwen")
  ) {
    return stableCodeFimTemplate
  }

  if (lowerCaseModel.includes("codestral")) {
    return codestralMultifileFimTemplate
  }

  if (lowerCaseModel.includes("codegemma")) {
    return codegemmaFimTemplate
  }

  if (lowerCaseModel.includes("codellama")) {
    return codeLlamaFimTemplate
  }

  if (lowerCaseModel.includes("deepseek")) {
    return deepseekFimTemplate
  }

  if (lowerCaseModel.includes("codegeex")) {
    return codegeexFimTemplate
  }

  // if (
  // 	lowerCaseModel.includes("gpt") ||
  // 	lowerCaseModel.includes("davinci-002") ||
  // 	lowerCaseModel.includes("claude") ||
  // 	lowerCaseModel.includes("granite3") ||
  // 	lowerCaseModel.includes("granite-3")
  // ) {
  // Default Fallback Mode
  return holeFillerTemplate
  //}

  //return stableCodeFimTemplate
}
