import { distance } from "fastest-levenshtein"

import { ChatMessage } from "../index.js"
import { renderChatMessage } from "../util/messageContent.js"

export type LineStream = AsyncGenerator<string>

type MatchLineResult = {
  /**
   * -1 if it's a new line, otherwise the index of the first match
   * in the old lines.
   */
  matchIndex: number
  isPerfectMatch: boolean
  newLine: string
}

function linesMatchPerfectly(lineA: string, lineB: string): boolean {
  return lineA === lineB && lineA !== ""
}

const END_BRACKETS = ["}", "});", "})"]

function linesMatch(lineA: string, lineB: string, linesBetween = 0): boolean {
  // Require a perfect (without padding) match for these lines
  // Otherwise they are edit distance 1 from empty lines and other single char lines (e.g. each other)
  if (["}", "*", "});", "})"].includes(lineA.trim())) {
    return lineA.trim() === lineB.trim()
  }

  const d = distance(lineA, lineB)

  return (
    // Should be more unlikely for lines to fuzzy match if they are further away
    (d / Math.max(lineA.length, lineB.length) <= Math.max(0, 0.48 - linesBetween * 0.06) ||
      lineA.trim() === lineB.trim()) &&
    lineA.trim() !== ""
  )
}

/**
 * Used to find a match for a new line in an array of old lines.
 *
 * Return the index of the first match and whether it is a perfect match
 * Also return a version of the line with correct indentation if needs fixing
 */
export function matchLine(newLine: string, oldLines: string[], permissiveAboutIndentation = false): MatchLineResult {
  // Only match empty lines if it's the next one:
  if (newLine.trim() === "" && oldLines[0]?.trim() === "") {
    return {
      matchIndex: 0,
      isPerfectMatch: true,
      newLine: newLine.trim(),
    }
  }

  const isEndBracket = END_BRACKETS.includes(newLine.trim())

  for (let i = 0; i < oldLines.length; i++) {
    // trims trailing whitespaces from the lines before comparison
    //this ensures trailing spaces don't affect matching.
    const oldLineTrimmed = oldLines[i].trimEnd()
    const newLineTrimmed = newLine.trimEnd()

    // Don't match end bracket lines if too far away
    if (i > 4 && isEndBracket) {
      return { matchIndex: -1, isPerfectMatch: false, newLine }
    }

    if (linesMatchPerfectly(newLineTrimmed, oldLineTrimmed)) {
      return { matchIndex: i, isPerfectMatch: true, newLine }
    }
    if (linesMatch(newLineTrimmed, oldLineTrimmed, i)) {
      // This is a way to fix indentation, but only for sufficiently long lines to avoid matching whitespace or short lines
      if (
        newLineTrimmed.trimStart() === oldLineTrimmed.trimStart() &&
        (permissiveAboutIndentation || newLine.trim().length > 8)
      ) {
        return {
          matchIndex: i,
          isPerfectMatch: true,
          newLine: oldLines[i],
        }
      }
      return { matchIndex: i, isPerfectMatch: false, newLine }
    }
  }

  return { matchIndex: -1, isPerfectMatch: false, newLine }
}

/**
 * Convert a stream of arbitrary chunks to a stream of lines
 */
export async function* streamLines(
  streamCompletion: AsyncGenerator<string | ChatMessage>,
  log: boolean = false,
): LineStream {
  const allLines = []

  let buffer = ""

  try {
    for await (const update of streamCompletion) {
      const chunk = typeof update === "string" ? update : renderChatMessage(update)
      buffer += chunk
      const lines = buffer.split("\n")
      buffer = lines.pop() ?? ""
      for (const line of lines) {
        yield line
        allLines.push(line)
      }

      // if (buffer === "" && chunk.endsWith("\n")) {
      //   yield "";
      //   allLines.push("");
      // }
    }
    if (buffer.length > 0) {
      yield buffer
      allLines.push(buffer)
    }
  } finally {
    if (log) {
      console.log("Streamed lines: ", allLines.join("\n"))
    }
  }
}

export async function* generateLines<T>(lines: T[]): AsyncGenerator<T> {
  for (const line of lines) {
    yield line
    // await new Promise((resolve, reject) => setTimeout(() => resolve(null), 50));
  }
}
