import { afterEach, describe, expect, it, vi } from "vitest"

// ---------- Module mocks ----------

// Simple Handlebars mock that does naive placeholder substitution
vi.mock("handlebars", () => {
  return {
    default: {
      compile: (template: string) => (ctx: Record<string, string>) => {
        return template
          .replace(/{{prefix}}/g, ctx.prefix)
          .replace(/{{suffix}}/g, ctx.suffix)
          .replace(/{{filename}}/g, ctx.filename ?? "")
          .replace(/{{reponame}}/g, ctx.reponame ?? "")
          .replace(/{{language}}/g, ctx.language ?? "")
      },
    },
  }
})

// Token utilities – we map 1 char = 1 token for simplicity
vi.mock("../../../llm/countTokens", () => {
  const countTokens = (str: string) => str.length
  const pruneLinesFromTop = (str: string, allowed: number) => str.slice(Math.max(0, str.length - allowed))
  const pruneLinesFromBottom = (str: string, allowed: number) => str.slice(0, allowed)
  const getTokenCountingBufferSafety = () => 0

  return {
    countTokens,
    pruneLinesFromTop,
    pruneLinesFromBottom,
    getTokenCountingBufferSafety,
  }
})

// Snippet selection – configurable via constant return value
vi.mock("../filtering", () => ({
  getSnippets: () => [],
}))

// Snippet formatting
const FORMATTED_SNIPPETS = "[FORMATTED_SNIPPETS]"
vi.mock("../formatting", () => ({
  formatSnippets: () => FORMATTED_SNIPPETS,
}))

// Stop tokens helper – we expose a variable so each test can override it
let stopTokenReturn: string[] = ["<STOP>"]
vi.mock("../getStopTokens", () => ({
  getStopTokens: () => stopTokenReturn,
}))

// AutocompleteTemplate – provide overridable template + compiler + completionOptions
let templateOverride: any = (prefix: string, suffix: string) => `${prefix}|${suffix}`
let compileFnOverride: ((...args: any[]) => [string, string]) | undefined
let completionOptionsOverride: Record<string, any> | undefined
vi.mock("../AutocompleteTemplate", () => ({
  getTemplateForModel: () => ({
    template: templateOverride,
    compilePrefixSuffix: compileFnOverride,
    completionOptions: completionOptionsOverride ?? {},
  }),
}))

// ---------- Imports after mocks ----------
import { renderPrompt, renderPromptWithTokenLimit } from ".."
import { AutocompleteLanguageInfo } from "../../constants/AutocompleteLanguageInfo"
import { SnippetPayload } from "../../snippets"
import { HelperVars } from "../../util/HelperVars"

// ---------- Helper builders ----------

const tsLang: AutocompleteLanguageInfo = {
  name: "TypeScript",
  topLevelKeywords: [],
  singleLineComment: "//",
  endOfLine: [";"],
}

const emptySnippetPayload: SnippetPayload = {
  rootPathSnippets: [],
  importDefinitionSnippets: [],
  ideSnippets: [],
  recentlyEditedRangeSnippets: [],
  recentlyVisitedRangesSnippets: [],
  diffSnippets: [],
  clipboardSnippets: [],
  recentlyOpenedFileSnippets: [],
  staticSnippet: [],
}

function makeHelper() {
  return {
    input: {
      filepath: "file:///test.ts",
      pos: { line: 0, character: 0 },
      recentlyEditedRanges: [],
      recentlyVisitedRanges: [],
    },
    prunedPrefix: "PRUNED_PREFIX",
    prunedSuffix: "PRUNED_SUFFIX",
    lang: tsLang,
    modelName: "test-model",
    filepath: "file:///test.ts",
    workspaceUris: [],
    options: {
      maxPromptTokens: 2048,
      prefixPercentage: 0.5,
      maxSuffixPercentage: 0.5,
      experimental_includeClipboard: false,
      useRecentlyOpened: false,
      experimental_includeRecentlyVisitedRanges: false,
      experimental_includeRecentlyEditedRanges: false,
      experimental_includeDiff: false,
      onlyMyCode: false,
    },
  } as unknown as HelperVars
}

afterEach(() => {
  // reset overridable mocks
  templateOverride = (prefix: string, suffix: string) => `${prefix}|${suffix}`
  compileFnOverride = undefined
  completionOptionsOverride = undefined
  stopTokenReturn = ["<STOP>"]
  vi.restoreAllMocks()
})

// ---------- Test suite ----------

describe("renderPrompt prefix/suffix selection", () => {
  it("uses manuallyPassPrefix when provided", () => {
    const helper = makeHelper()
    helper.input.manuallyPassPrefix = "MANUAL"
    const { prefix, suffix } = renderPrompt({
      snippetPayload: emptySnippetPayload,
      workspaceDirs: ["file:///workspace"],
      helper,
    })

    expect(suffix).toBe("\n")
    expect(prefix.endsWith("MANUAL")).toBe(true)
  })

  it("falls back to prunedPrefix when no manual prefix", () => {
    const helper = makeHelper()

    const { prefix } = renderPrompt({
      snippetPayload: emptySnippetPayload,
      workspaceDirs: ["file:///workspace"],
      helper,
    })

    expect(prefix.includes("PRUNED_PREFIX")).toBe(true)
  })
})

describe("template rendering paths", () => {
  it("handles function template", () => {
    templateOverride = (p: string, s: string, _filepath: string, _reponame: string) => `FUNC:${p}|${s}`

    const helper = makeHelper()

    const { prompt } = renderPrompt({
      snippetPayload: emptySnippetPayload,
      workspaceDirs: ["file:///workspace"],
      helper,
    })

    expect(prompt.startsWith("FUNC:")).toBe(true)
    expect(prompt.includes("PRUNED_PREFIX")).toBe(true)
  })
})

describe("compilePrefixSuffix vs snippet formatting", () => {
  it("applies compilePrefixSuffix when provided", () => {
    compileFnOverride = (p: string, s: string) => [`COMP_${p}`, `COMP_${s}`]
    templateOverride = (prefix: string, suffix: string) => `${prefix}|${suffix}`
    const helper = makeHelper()

    const { prefix: compiledPrefix } = renderPrompt({
      snippetPayload: emptySnippetPayload,
      workspaceDirs: ["file:///workspace"],
      helper,
    })

    expect(compiledPrefix.startsWith("COMP_PRUNED_PREFIX")).toBe(true)
  })

  it("prepends formatted snippets when no compiler present", () => {
    const helper = makeHelper()

    const { prefix: compiledPrefix } = renderPrompt({
      snippetPayload: emptySnippetPayload,
      workspaceDirs: ["file:///workspace"],
      helper,
    })

    expect(compiledPrefix.startsWith(`${FORMATTED_SNIPPETS}\n`)).toBe(true)
  })
})

describe("renderPromptWithTokenLimit parity & pruning", () => {
  it("matches renderPrompt when llm is undefined", () => {
    const helper = makeHelper()

    const res1 = renderPrompt({
      snippetPayload: emptySnippetPayload,
      workspaceDirs: ["file:///workspace"],
      helper,
    })

    const res2 = renderPromptWithTokenLimit({
      snippetPayload: emptySnippetPayload,
      workspaceDirs: ["file:///workspace"],
      helper,
      llm: undefined,
    })

    expect(res2).toEqual(res1)
  })

  it("prunes prefix/suffix to respect small context length", () => {
    const longPrefix = "A".repeat(300)

    const helper = makeHelper()
    ;(helper as any).prunedPrefix = longPrefix

    const llmStub = {
      contextLength: 120,
      completionOptions: { maxTokens: 10 },
      model: "test-model",
    } as any

    const { prefix: compiledPrefix } = renderPromptWithTokenLimit({
      snippetPayload: emptySnippetPayload,
      workspaceDirs: ["file:///workspace"],
      helper,
      llm: llmStub,
    })

    expect(compiledPrefix.length).toBeLessThan(120)
  })
})

describe("stop-token merging", () => {
  it("returns stop tokens from getStopTokens", () => {
    stopTokenReturn = ["LANG_STOP", "TEMPLATE_STOP"]
    completionOptionsOverride = { stop: ["TEMPLATE_STOP"] }
    templateOverride = (prefix: string, suffix: string) => `${prefix}|${suffix}`

    const helper = makeHelper()

    const { completionOptions } = renderPrompt({
      snippetPayload: emptySnippetPayload,
      workspaceDirs: ["file:///workspace"],
      helper,
    })

    expect(completionOptions?.stop).toEqual(stopTokenReturn)
  })
})
