export const BRACKETS: { [key: string]: string } = {
  "(": ")",
  "{": "}",
  "[": "]",
}
export const BRACKETS_REVERSE: { [key: string]: string } = {
  ")": "(",
  "}": "{",
  "]": "[",
}
/**
 * We follow the policy of only completing bracket pairs that we started
 * But sometimes we started the pair in a previous autocomplete suggestion
 */
export class BracketMatchingService {
  private openingBracketsFromLastCompletion: string[] = []
  private lastCompletionFile: string | undefined = undefined

  handleAcceptedCompletion(completion: string, filepath: string) {
    this.openingBracketsFromLastCompletion = []
    const stack: string[] = []

    for (let i = 0; i < completion.length; i++) {
      const char = completion[i]
      if (Object.keys(BRACKETS).includes(char)) {
        // It's an opening bracket
        stack.push(char)
      } else if (Object.values(BRACKETS).includes(char)) {
        // It's a closing bracket
        if (stack.length === 0 || BRACKETS[stack.pop()!] !== char) {
          break
        }
      }
    }

    // Any remaining opening brackets in the stack are uncompleted
    this.openingBracketsFromLastCompletion = stack
    this.lastCompletionFile = filepath
  }

  async *stopOnUnmatchedClosingBracket(
    stream: AsyncGenerator<string>,
    prefix: string,
    suffix: string,
    filepath: string,
    multiline: boolean, // Whether this is a multiline completion or not
  ): AsyncGenerator<string> {
    let stack: string[] = []
    if (multiline) {
      // Add opening brackets from the previous response
      if (this.lastCompletionFile === filepath) {
        stack = [...this.openingBracketsFromLastCompletion]
      } else {
        this.lastCompletionFile = undefined
      }
    } else {
      // If single line completion, then allow completing bracket pairs that are
      // started on the current line but not finished on the current line
      if (!multiline) {
        const currentLine = (prefix.split("\n").pop() ?? "") + (suffix.split("\n")[0] ?? "")
        for (let i = 0; i < currentLine.length; i++) {
          const char = currentLine[i]
          if (Object.keys(BRACKETS).includes(char)) {
            // It's an opening bracket
            stack.push(char)
          } else if (Object.values(BRACKETS).includes(char)) {
            // It's a closing bracket
            if (stack.length === 0 || BRACKETS[stack.pop()!] !== char) {
              break
            }
          }
        }
      }
    }

    // Add corresponding open brackets from suffix to stack
    // because we overwrite them and the diff is displayed, and this allows something to be edited after that
    for (let i = 0; i < suffix.length; i++) {
      if (suffix[i] === " ") {
        continue
      }
      const openBracket = BRACKETS_REVERSE[suffix[i]]
      if (!openBracket) {
        break
      }
      stack.unshift(openBracket)
    }

    let seenNonWhitespaceOrClosingBracket = false
    for await (let chunk of stream) {
      // Allow closing brackets before any non-whitespace characters
      if (!seenNonWhitespaceOrClosingBracket) {
        const firstNonWhitespaceOrClosingBracketIndex = chunk.search(/[^\s)}\]]/)
        if (firstNonWhitespaceOrClosingBracketIndex !== -1) {
          yield chunk.slice(0, firstNonWhitespaceOrClosingBracketIndex)
          chunk = chunk.slice(firstNonWhitespaceOrClosingBracketIndex)
          seenNonWhitespaceOrClosingBracket = true
        } else {
          yield chunk
          continue
        }
      }

      for (let i = 0; i < chunk.length; i++) {
        const char = chunk[i]
        if (Object.values(BRACKETS).includes(char)) {
          // It's a closing bracket
          if (stack.length === 0 || BRACKETS[stack.pop()!] !== char) {
            // If the stack is empty or the top of the stack doesn't match the current closing bracket
            yield chunk.slice(0, i)
            return // Stop the generator if the closing bracket doesn't have a matching opening bracket in the stream
          }
        } else if (Object.keys(BRACKETS).includes(char)) {
          // It's an opening bracket
          stack.push(char)
        }
      }
      yield chunk
    }
  }
}
