// Copyright (c) Facebook, Inc. and its 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

import com.google.common.base.MoreObjects
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 com.intellij.openapi.util.text.StringUtil
import com.intellij.psi.PsiComment
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiWhiteSpace
import com.intellij.psi.impl.source.tree.LeafPsiElement
import java.util.regex.Pattern
import org.jetbrains.kotlin.lexer.KtToken
import org.jetbrains.kotlin.lexer.KtTokens
import org.jetbrains.kotlin.psi.KtFile
import org.jetbrains.kotlin.psi.KtTreeVisitorVoid
import org.jetbrains.kotlin.psi.psiUtil.endOffset
import org.jetbrains.kotlin.psi.psiUtil.startOffset

// TODO: share the code with JavaInput instead of copy-pasting here.
/**
 * KotlinInput is for Kotlin what JavaInput is for Java.
 *
 * <p>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".
 *
 * <p>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)
    val tokenLocations = ImmutableRangeMap.builder<Int, Token>()
    for (token in tokens) {
      val end = JavaOutput.endTok(token)
      var upper = end.position
      if (end.text.isNotEmpty()) {
        upper += end.length() - 1
      }
      tokenLocations.put(Range.closed(JavaOutput.startTok(token).position, upper), token)
    }
    positionTokenMap = tokenLocations.build()

    // 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> {
    var length = length
    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))
    }
    when {
      length < 0 -> return EMPTY_RANGE
      length == 0 -> // 0 stands for "format the line under the cursor"
        length = 1
    }
    val enclosed = getPositionTokenMap()
        .subRangeMap(Range.closedOpen(offset, offset + length))
        .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<Int, Int> {
    val builder = ImmutableMap.builder<Int, Int>()
    for (tok in toks) {
      builder.put(tok.position, tok.column)
    }
    return builder.build()
  }

  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 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
}

class KotlinTok(
    private val index: Int,
    private val originalText: String,
    private val text: String,
    private val position: Int,
    private val columnI: Int,
    val isToken: Boolean,
    private val kind: KtToken
) : Input.Tok {

  override fun getIndex(): Int = index

  override fun getText(): String = text

  override fun getOriginalText(): String = originalText

  override fun length(): Int = originalText.length

  override fun getPosition(): Int = position

  override fun getColumn(): Int = columnI

  override fun isNewline(): Boolean = Newlines.isNewline(text)

  override fun isSlashSlashComment(): Boolean = text.startsWith("//")

  override fun isSlashStarComment(): Boolean = text.startsWith("/*")

  override fun isJavadocComment(): Boolean = text.startsWith("/**") && text.length > 4

  override fun isComment(): Boolean = isSlashSlashComment || isSlashStarComment

  fun kind(): KtToken = kind

  override fun toString(): String {
    return MoreObjects.toStringHelper(this)
        .add("index", index)
        .add("text", text)
        .add("position", position)
        .add("columnI", columnI)
        .add("isToken", isToken)
        .toString()
  }
}

class KotlinToken(
    private val toksBefore: ImmutableList<KotlinTok>,
    private val kotlinTok: KotlinTok,
    private val toksAfter: ImmutableList<KotlinTok>
) : Input.Token {

  override fun getTok(): KotlinTok = kotlinTok

  override fun getToksBefore(): ImmutableList<out Input.Tok> = toksBefore

  override fun getToksAfter(): ImmutableList<out Input.Tok> = toksAfter

  override fun toString(): String {
    return MoreObjects.toStringHelper(this)
        .add("tok", kotlinTok)
        .add("toksBefore", toksBefore)
        .add("toksAfter", toksAfter)
        .toString()
  }
}

internal val WHITESPACE_NEWLINE_REGEX: Pattern = Pattern.compile("\\R|( )+")

/**
 * Tokenizer traverses a Kotlin parse tree (which blessedly contains whitespaces and comments, unlike Javac) and
 * constructs a list of 'Tok's.
 *
 * <p>The google-java-format infra expects newline Toks to be separate from maximal-whitespace Toks, but Kotlin emits
 * them together. So, we split them using Java's \R regex matcher. We don't use 'split' et al. because we want Toks for
 * the newlines themselves.
 */
class Tokenizer(private val fileText: String, val file: KtFile) : KtTreeVisitorVoid() {
  val toks = mutableListOf<KotlinTok>()
  var index = 0

  override fun visitElement(element: PsiElement) {
    when (element) {
      is PsiComment -> {
        val startIndex = element.startOffset
        toks.add(
            KotlinTok(
                index,
                fileText.substring(startIndex, element.endOffset),
                element.text,
                startIndex,
                0,
                false,
                KtTokens.EOF))
        index++
        return
      }
      is LeafPsiElement -> {
        val elementText = element.text
        val startIndex = element.startOffset
        val endIndex = element.endOffset
        if (element is PsiWhiteSpace) {
          val matcher = WHITESPACE_NEWLINE_REGEX.matcher(elementText)
          while (matcher.find()) {
            val text = matcher.group()
            toks.add(
                KotlinTok(
                    -1,
                    fileText.substring(startIndex + matcher.start(), startIndex + matcher.end()),
                    text,
                    startIndex + matcher.start(),
                    0,
                    false,
                    KtTokens.EOF))
          }
        } else {
          toks.add(
              KotlinTok(
                  index,
                  fileText.substring(startIndex, endIndex),
                  elementText,
                  startIndex,
                  0,
                  true,
                  KtTokens.EOF))
          index++
        }
      }
    }
    super.visitElement(element)
  }
}