/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.facebook.ktfmt.format

import com.google.common.collect.ImmutableList
import com.google.googlejavaformat.DocBuilder
import com.google.googlejavaformat.Op

/**
 * A dummy [Op] that prevents comments from being moved ahead of it, into parent [Level]s.
 *
 * If a comment is the first thing in a [Level], [OpBuilder] moves it outside of that level during
 * [Doc] building. It does so recursively, until the comment is not the first element of the level.
 * This behaviour can be very confusing, where comments seem to absorb or ignore expected
 * indentation.
 */
object FenceCommentsOp : Op {
  val AS_LIST = ImmutableList.of<Op>(FenceCommentsOp)

  override fun add(builder: DocBuilder) {
    // Do nothing. This Op simply needs to be in the OpsBuilder.
  }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.facebook.ktfmt.format

import com.facebook.ktfmt.debughelpers.printOps
import com.facebook.ktfmt.format.FormattingOptions.Style.DROPBOX
import com.facebook.ktfmt.format.FormattingOptions.Style.GOOGLE
import com.facebook.ktfmt.format.RedundantElementRemover.dropRedundantElements
import com.facebook.ktfmt.format.WhitespaceTombstones.indexOfWhitespaceTombstone
import com.facebook.ktfmt.kdoc.Escaping
import com.facebook.ktfmt.kdoc.KDocCommentsHelper
import com.google.common.collect.ImmutableList
import com.google.common.collect.Range
import com.google.googlejavaformat.Doc
import com.google.googlejavaformat.DocBuilder
import com.google.googlejavaformat.Newlines
import com.google.googlejavaformat.OpsBuilder
import com.google.googlejavaformat.java.FormatterException
import com.google.googlejavaformat.java.JavaOutput
import org.jetbrains.kotlin.com.intellij.openapi.util.text.StringUtil
import org.jetbrains.kotlin.com.intellij.openapi.util.text.StringUtilRt
import org.jetbrains.kotlin.com.intellij.psi.PsiComment
import org.jetbrains.kotlin.com.intellij.psi.PsiElement
import org.jetbrains.kotlin.com.intellij.psi.PsiElementVisitor
import org.jetbrains.kotlin.com.intellij.psi.PsiWhiteSpace
import org.jetbrains.kotlin.psi.KtImportDirective
import org.jetbrains.kotlin.psi.psiUtil.endOffset
import org.jetbrains.kotlin.psi.psiUtil.startOffset

object Formatter {

  @JvmField
  val GOOGLE_FORMAT = FormattingOptions(style = GOOGLE, blockIndent = 2, continuationIndent = 2)

  /** A format that attempts to reflect https://kotlinlang.org/docs/coding-conventions.html. */
  @JvmField
  val KOTLINLANG_FORMAT = FormattingOptions(style = GOOGLE, blockIndent = 4, continuationIndent = 4)

  @JvmField
  val DROPBOX_FORMAT = FormattingOptions(style = DROPBOX, blockIndent = 4, continuationIndent = 4)

  private val MINIMUM_KOTLIN_VERSION = KotlinVersion(1, 4)

  /**
   * format formats the Kotlin code given in 'code' and returns it as a string. This method is
   * accessed through Reflection.
   */
  @JvmStatic
  @Throws(FormatterException::class, ParseError::class)
  fun format(code: String): String = format(FormattingOptions(), code)

  /**
   * format formats the Kotlin code given in 'code' with 'removeUnusedImports' and returns it as a
   * string. This method is accessed through Reflection.
   */
  @JvmStatic
  @Throws(FormatterException::class, ParseError::class)
  fun format(code: String, removeUnusedImports: Boolean): String =
      format(FormattingOptions(removeUnusedImports = removeUnusedImports), code)

  /**
   * format formats the Kotlin code given in 'code' with the 'maxWidth' and returns it as a string.
   */
  @JvmStatic
  @Throws(FormatterException::class, ParseError::class)
  fun format(options: FormattingOptions, code: String): String {
    val (shebang, kotlinCode) =
        if (code.startsWith("#!")) {
          code.split("\n".toRegex(), limit = 2)
        } else {
          listOf("", code)
        }
    checkEscapeSequences(kotlinCode)

    val lfCode = StringUtilRt.convertLineSeparators(kotlinCode)
    val sortedImports = sortedAndDistinctImports(lfCode)
    val noRedundantElements = dropRedundantElements(sortedImports, options)
    val prettyCode =
        prettyPrint(noRedundantElements, options, Newlines.guessLineSeparator(kotlinCode)!!)
    return if (shebang.isNotEmpty()) shebang + "\n" + prettyCode else prettyCode
  }

  /** prettyPrint reflows 'code' using google-java-format's engine. */
  private fun prettyPrint(code: String, options: FormattingOptions, lineSeparator: String): String {
    val file = Parser.parse(code)
    val kotlinInput = KotlinInput(code, file)
    val javaOutput =
        JavaOutput(lineSeparator, kotlinInput, KDocCommentsHelper(lineSeparator, options.maxWidth))
    val builder = OpsBuilder(kotlinInput, javaOutput)
    file.accept(createAstVisitor(options, builder))
    builder.sync(kotlinInput.text.length)
    builder.drain()
    val ops = builder.build()
    if (options.debuggingPrintOpsAfterFormatting) {
      printOps(ops)
    }
    val doc = DocBuilder().withOps(ops).build()
    doc.computeBreaks(javaOutput.commentsHelper, options.maxWidth, Doc.State(+0, 0))
    doc.write(javaOutput)
    javaOutput.flush()

    val tokenRangeSet =
        kotlinInput.characterRangesToTokenRanges(ImmutableList.of(Range.closedOpen(0, code.length)))
    return WhitespaceTombstones.replaceTombstoneWithTrailingWhitespace(
        JavaOutput.applyReplacements(code, javaOutput.getFormatReplacements(tokenRangeSet)))
  }

  private fun createAstVisitor(options: FormattingOptions, builder: OpsBuilder): PsiElementVisitor {
    if (KotlinVersion.CURRENT < MINIMUM_KOTLIN_VERSION) {
      throw RuntimeException("Unsupported runtime Kotlin version: " + KotlinVersion.CURRENT)
    }
    return KotlinInputAstVisitor(options, builder)
  }

  private fun checkEscapeSequences(code: String) {
    var index = code.indexOfWhitespaceTombstone()
    if (index == -1) {
      index = Escaping.indexOfCommentEscapeSequences(code)
    }
    if (index != -1) {
      throw ParseError(
          "ktfmt does not support code which contains one of {\\u0003, \\u0004, \\u0005} character" +
              "; escape it",
          StringUtil.offsetToLineColumn(code, index))
    }
  }

  private fun sortedAndDistinctImports(code: String): String {
    val file = Parser.parse(code)

    val importList = file.importList ?: return code
    if (importList.imports.isEmpty()) {
      return code
    }

    val commentList = mutableListOf<PsiElement>()
    // Find non-import elements; comments are moved, in order, to the top of the import list. Other
    // non-import elements throw a ParseError.
    var element = importList.firstChild
    while (element != null) {
      if (element is PsiComment) {
        commentList.add(element)
      } else if (element !is KtImportDirective && element !is PsiWhiteSpace) {
        throw ParseError(
            "Imports not contiguous: " + element.text,
            StringUtil.offsetToLineColumn(code, element.startOffset))
      }
      element = element.nextSibling
    }
    fun canonicalText(importDirective: KtImportDirective) =
        importDirective.importedFqName?.asString() +
            " " +
            importDirective.alias?.text?.replace("`", "") +
            " " +
            if (importDirective.isAllUnder) "*" else ""

    val sortedImports = importList.imports.sortedBy(::canonicalText).distinctBy(::canonicalText)
    val importsWithComments = commentList + sortedImports

    return code.replaceRange(
        importList.startOffset,
        importList.endOffset,
        importsWithComments.joinToString(separator = "\n") { imprt -> imprt.text } + "\n")
  }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.facebook.ktfmt.format

data class FormattingOptions(
    val style: Style = Style.FACEBOOK,

    /** ktfmt breaks lines longer than maxWidth. */
    val maxWidth: Int = DEFAULT_MAX_WIDTH,

    /**
     * blockIndent is the size of the indent used when a new block is opened, in spaces.
     *
     * For example,
     * ```
     * fun f() {
     *   //
     * }
     * ```
     */
    val blockIndent: Int = 2,

    /**
     * continuationIndent is the size of the indent used when a line is broken because it's too
     * long, in spaces.
     *
     * For example,
     * ```
     * val foo = bar(
     *     1)
     * ```
     */
    val continuationIndent: Int = 4,

    /** Whether ktfmt should remove imports that are not used. */
    val removeUnusedImports: Boolean = true,

    /**
     * Print the Ops generated by KotlinInputAstVisitor to help reason about formatting (i.e.,
     * newline) decisions
     */
    val debuggingPrintOpsAfterFormatting: Boolean = false
) {

  companion object {
    const val DEFAULT_MAX_WIDTH: Int = 100
  }

  enum class Style {
    FACEBOOK,
    DROPBOX,
    GOOGLE
  }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.facebook.ktfmt.format

import com.google.common.collect.DiscreteDomain
import com.google.common.collect.ImmutableList
import com.google.common.collect.ImmutableMap
import com.google.common.collect.ImmutableRangeMap
import com.google.common.collect.Iterables.getLast
import com.google.common.collect.Range
import com.google.common.collect.RangeSet
import com.google.common.collect.TreeRangeSet
import com.google.googlejavaformat.Input
import com.google.googlejavaformat.Newlines
import com.google.googlejavaformat.java.FormatterException
import com.google.googlejavaformat.java.JavaOutput
import org.jetbrains.kotlin.com.intellij.openapi.util.text.StringUtil
import org.jetbrains.kotlin.lexer.KtTokens
import org.jetbrains.kotlin.psi.KtFile

// TODO: share the code with JavaInput instead of copy-pasting here.
/**
 * KotlinInput is for Kotlin what JavaInput is for Java.
 *
 *
KotlinInput is duplicating most of JavaInput's code, but uses the Kotlin compiler as a lexer
 * instead of Javac. This is required because some valid Kotlin programs are not valid Java
 * programs, e.g., "a..b".
 *
 *
See javadoc for JavaInput.
 */
class KotlinInput(private val text: String, file: KtFile) : Input() {
  private val tokens: ImmutableList<Token> // The Tokens for this input.
  private val positionToColumnMap: ImmutableMap<Int, Int> // Map Tok position to column.
  private val positionTokenMap: ImmutableRangeMap<Int, Token> // Map position to Token.
  private var kN = 0 // The number of numbered toks (tokens or comments), excluding the EOF.
  private val kToToken: Array<Token?>

  init {
    setLines(ImmutableList.copyOf(Newlines.lineIterator(text)))
    val toks = buildToks(file, text)
    positionToColumnMap = makePositionToColumnMap(toks)
    tokens = buildTokens(toks)
    positionTokenMap = buildTokenPositionsMap(tokens)

    // adjust kN for EOF
    kToToken = arrayOfNulls(kN + 1)
    for (token in tokens) {
      for (tok in token.toksBefore) {
        if (tok.index < 0) {
          continue
        }
        kToToken[tok.index] = token
      }
      kToToken[token.tok.index] = token
      for (tok in token.toksAfter) {
        if (tok.index < 0) {
          continue
        }
        kToToken[tok.index] = token
      }
    }
  }

  @Throws(FormatterException::class)
  fun characterRangesToTokenRanges(characterRanges: Collection<Range<Int>>): RangeSet<Int> {
    val tokenRangeSet = TreeRangeSet.create<Int>()
    for (characterRange0 in characterRanges) {
      val characterRange = characterRange0.canonical(DiscreteDomain.integers())
      tokenRangeSet.add(
          characterRangeToTokenRange(
              characterRange.lowerEndpoint(),
              characterRange.upperEndpoint() - characterRange.lowerEndpoint()))
    }
    return tokenRangeSet
  }

  /**
   * Convert from an offset and length flag pair to a token range.
   *
   * @param offset the `0`-based offset in characters
   * @param length the length in characters
   * @return the `0`-based [Range] of tokens
   * @throws FormatterException
   */
  @Throws(FormatterException::class)
  internal fun characterRangeToTokenRange(offset: Int, length: Int): Range<Int> {
    val requiredLength = offset + length
    if (requiredLength > text.length) {
      throw FormatterException(
          String.format(
              "error: invalid length %d, offset + length (%d) is outside the file",
              length,
              requiredLength))
    }
    val expandedLength =
        when {
          length < 0 -> return EMPTY_RANGE
          length == 0 -> 1 // 0 stands for "format the line under the cursor"
          else -> length
        }
    val enclosed =
        getPositionTokenMap()
            .subRangeMap(Range.closedOpen(offset, offset + expandedLength))
            .asMapOfRanges()
            .values
    return if (enclosed.isEmpty()) {
      EMPTY_RANGE
    } else
        Range.closedOpen(
            enclosed.iterator().next().tok.index, getLast(enclosed).getTok().getIndex() + 1)
  }

  private fun makePositionToColumnMap(toks: List<KotlinTok>) =
      ImmutableMap.copyOf(toks.map { it.position to it.column }.toMap())

  private fun buildToks(file: KtFile, fileText: String): ImmutableList<KotlinTok> {
    val tokenizer = Tokenizer(fileText, file)
    file.accept(tokenizer)
    val toks = tokenizer.toks
    toks.add(KotlinTok(tokenizer.index, "", "", fileText.length, 0, true, KtTokens.EOF))
    kN = tokenizer.index
    computeRanges(toks)
    return ImmutableList.copyOf(toks)
  }

  private fun buildTokens(toks: List<KotlinTok>): ImmutableList<Token> {
    val tokens = ImmutableList.builder<Token>()
    var k = 0
    val kN = toks.size

    // Remaining non-tokens before the token go here.
    var toksBefore: ImmutableList.Builder<KotlinTok> = ImmutableList.builder()

    OUTERMOST@ while (k < kN) {
      while (!toks[k].isToken) {
        val tok = toks[k++]
        toksBefore.add(tok)
        if (isParamComment(tok)) {
          while (toks[k].isNewline) {
            // drop newlines after parameter comments
            k++
          }
        }
      }
      val tok = toks[k++]

      // Non-tokens starting on the same line go here too.
      val toksAfter = ImmutableList.builder<KotlinTok>()
      OUTER@ while (k < kN && !toks[k].isToken) {
        // Don't attach inline comments to certain leading tokens, e.g. for `f(/*flag1=*/true).
        //
        // Attaching inline comments to the right token is hard, and this barely
        // scratches the surface. But it's enough to do a better job with parameter
        // name comments.
        //
        // TODO(cushon): find a better strategy.
        if (toks[k].isSlashStarComment && (tok.text == "(" || tok.text == "<" || tok.text == "."))
            break@OUTER
        if (toks[k].isJavadocComment && tok.text == ";") break@OUTER
        if (isParamComment(toks[k])) {
          tokens.add(KotlinToken(toksBefore.build(), tok, toksAfter.build()))
          toksBefore = ImmutableList.builder<KotlinTok>().add(toks[k++])
          // drop newlines after parameter comments
          while (toks[k].isNewline) {
            k++
          }
          continue@OUTERMOST
        }
        val nonTokenAfter = toks[k++]
        toksAfter.add(nonTokenAfter)
        if (Newlines.containsBreaks(nonTokenAfter.text)) {
          break
        }
      }
      tokens.add(KotlinToken(toksBefore.build(), tok, toksAfter.build()))
      toksBefore = ImmutableList.builder()
    }
    return tokens.build()
  }

  private fun buildTokenPositionsMap(tokens: ImmutableList<Token>): ImmutableRangeMap<Int, Token> {
    val tokenLocations = ImmutableRangeMap.builder<Int, Token>()
    for (token in tokens) {
      val end = JavaOutput.endTok(token)
      val endPosition = end.position + (if (end.text.isNotEmpty()) end.length() - 1 else 0)
      tokenLocations.put(Range.closed(JavaOutput.startTok(token).position, endPosition), token)
    }

    return tokenLocations.build()
  }

  private fun isParamComment(tok: Tok): Boolean {
    return tok.isSlashStarComment && tok.text.matches("/\\*[A-Za-z0-9\\s_\\-]+=\\s*\\*/".toRegex())
  }

  override fun getkN(): Int = kN

  override fun getToken(k: Int): Token? = kToToken[k]

  override fun getTokens(): ImmutableList<out Token> = tokens

  override fun getPositionTokenMap(): ImmutableRangeMap<Int, out Token> = positionTokenMap

  override fun getPositionToColumnMap(): ImmutableMap<Int, Int> = positionToColumnMap

  override fun getText(): String = text

  override fun getLineNumber(inputPosition: Int) =
      StringUtil.offsetToLineColumn(text, inputPosition).line + 1

  override fun getColumnNumber(inputPosition: Int) =
      StringUtil.offsetToLineColumn(text, inputPosition).column
}