## Error in SQL Delight IntelliJ Plugin

**java.lang.NullPointerException** in **SqlDelightFoldingBuilder.kt:79**
nextSibling must not be null

[View on Bugsnag](https://app.bugsnag.com/square-inc/sql-delight-intellij-plugin/errors/60967a6f7d8e8d000751a733?event_id=60967a6f007a892c354a0000&i=gh&m=ci)

## Stacktrace

    SqlDelightFoldingBuilder.kt:79 - com.squareup.sqldelight.intellij.lang.SqlDelightFoldingBuilder.toCreateTableDescriptor
    SqlDelightFoldingBuilder.kt:56 - com.squareup.sqldelight.intellij.lang.SqlDelightFoldingBuilder.createFoldingDescriptors
    SqlDelightFoldingBuilder.kt:39 - com.squareup.sqldelight.intellij.lang.SqlDelightFoldingBuilder.buildFoldRegions

[View full stacktrace](https://app.bugsnag.com/square-inc/sql-delight-intellij-plugin/errors/60967a6f7d8e8d000751a733?event_id=60967a6f007a892c354a0000&i=gh&m=ci)

*Created automatically via Bugsnag*

diff --git a/sqldelight-idea-plugin/src/main/kotlin/com/squareup/sqldelight/intellij/lang/SqlDelightFoldingBuilder.kt b/sqldelight-idea-plugin/src/main/kotlin/com/squareup/sqldelight/intellij/lang/SqlDelightFoldingBuilder.kt
index cb35f7e56..f4057b916 100644
--- a/sqldelight-idea-plugin/src/main/kotlin/com/squareup/sqldelight/intellij/lang/SqlDelightFoldingBuilder.kt
+++ b/sqldelight-idea-plugin/src/main/kotlin/com/squareup/sqldelight/intellij/lang/SqlDelightFoldingBuilder.kt
@@ -76,7 +76,7 @@ class SqlDelightFoldingBuilder : FoldingBuilder, DumbAware {
   private fun PsiElement.toCreateTableDescriptor(createTableStmt: PsiElement?): FoldingDescriptor? {
     val openingBraceElement = createTableStmt?.node?.findChildByType(SqlTypes.LP) ?: return null
     val start = openingBraceElement.startOffset
-    val end = nextSibling.endOffset
+    val end = nextSibling?.endOffset ?: 0
     if (start >= end) return null
     return FoldingDescriptor(this, start, end, null, "(...);")
   }

java.lang.NullPointerException in SqlDelightFoldingBuilder.kt:79
## Error in SQL Delight IntelliJ Plugin

**java.lang.NullPointerException** in **SqlDelightFoldingBuilder.kt:79**
nextSibling must not be null

[View on Bugsnag](https://app.bugsnag.com/square-inc/sql-delight-intellij-plugin/errors/60967a6f7d8e8d000751a733?event_id=60967a6f007a892c354a0000&i=gh&m=ci)

## Stacktrace

    SqlDelightFoldingBuilder.kt:79 - com.squareup.sqldelight.intellij.lang.SqlDelightFoldingBuilder.toCreateTableDescriptor
    SqlDelightFoldingBuilder.kt:56 - com.squareup.sqldelight.intellij.lang.SqlDelightFoldingBuilder.createFoldingDescriptors
    SqlDelightFoldingBuilder.kt:39 - com.squareup.sqldelight.intellij.lang.SqlDelightFoldingBuilder.buildFoldRegions

[View full stacktrace](https://app.bugsnag.com/square-inc/sql-delight-intellij-plugin/errors/60967a6f7d8e8d000751a733?event_id=60967a6f007a892c354a0000&i=gh&m=ci)

*Created automatically via Bugsnag*

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight

/**
 * Marshal and map the type `T` to and from a database type `S` which is one of
 * `Long`, `Double`, `String`, `byte[]`.
 */
interface ColumnAdapter<T : Any, S> {
  /**
   * @return [databaseValue] decoded as type [T].
   */
  fun decode(databaseValue: S): T

  /**
   * @return [value] encoded as database type [S].
   */
  fun encode(value: T): S
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight

/** A [ColumnAdapter] which maps the enum class `T` to a string in the database. */
class EnumColumnAdapter<T : Enum<T>> @PublishedApi internal constructor(
  private val enumValues: Array<out T>
) : ColumnAdapter<T, String> {
  override fun decode(databaseValue: String): T = enumValues.first { it.name == databaseValue }

  override fun encode(value: T) = value.name
}

/** A [ColumnAdapter] which maps the enum class `T` to a string in the database. */
@Suppress("FunctionName") // Emulating a constructor.
inline fun <reified T : Enum<T>> EnumColumnAdapter(): EnumColumnAdapter<T> {
  return EnumColumnAdapter(enumValues())
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight

import com.squareup.sqldelight.db.SqlCursor
import com.squareup.sqldelight.db.SqlDriver
import com.squareup.sqldelight.db.use
import com.squareup.sqldelight.internal.QueryLock
import com.squareup.sqldelight.internal.copyOnWriteList
import com.squareup.sqldelight.internal.withLock

/**
 * A listenable, typed query generated by SQLDelight.
 *
 * @param RowType the type that this query can map its result set to.
 */
@Suppress("FunctionName") // Emulating a constructor.
fun <RowType : Any> Query(
  identifier: Int,
  queries: MutableList<Query<*>>,
  driver: SqlDriver,
  query: String,
  mapper: (SqlCursor) -> RowType
): Query<RowType> {
  return Query(identifier, queries, driver, "unknown", "unknown", query, mapper)
}

/**
 * A listenable, typed query generated by SQLDelight.
 *
 * @param RowType the type that this query can map its result set to.
 */
@Suppress("FunctionName") // Emulating a constructor.
fun <RowType : Any> Query(
  identifier: Int,
  queries: MutableList<Query<*>>,
  driver: SqlDriver,
  fileName: String,
  label: String,
  query: String,
  mapper: (SqlCursor) -> RowType
): Query<RowType> {
  return SimpleQuery(identifier, queries, driver, fileName, label, query, mapper)
}

private class SimpleQuery<out RowType : Any>(
  private val identifier: Int,
  queries: MutableList<Query<*>>,
  private val driver: SqlDriver,
  private val fileName: String,
  private val label: String,
  private val query: String,
  mapper: (SqlCursor) -> RowType
) : Query<RowType>(queries, mapper) {
  override fun execute(): SqlCursor {
    return driver.executeQuery(identifier, query, 0)
  }

  override fun toString() = "$fileName:$label"
}

/**
 * A listenable, typed query generated by SQLDelight.
 *
 * @param RowType the type that this query can map its result set to.
 *
 * @property mapper The mapper this [Query] was created with, which can convert a row in the SQL
 *   cursor returned by [execute] to [RowType].
 */
abstract class Query<out RowType : Any>(
  private val queries: MutableList<Query<*>>,
  val mapper: (SqlCursor) -> RowType
) {
  private val listenerLock = QueryLock()
  private val listeners = copyOnWriteList<Listener>()

  /**
   * Notify listeners that their current result set is staled.
   *
   * Called internally by SQLDelight when it detects a possible staling of the result set. Emits
   * some false positives but never misses a true positive.
   */
  fun notifyDataChanged() {
    listenerLock.withLock {
      listeners.forEach(Listener::queryResultsChanged)
    }
  }

  /**
   * Register a listener to be notified of future changes in the result set.
   */
  fun addListener(listener: Listener) {
    listenerLock.withLock {
      if (listeners.isEmpty()) queries.add(this)
      listeners.add(listener)
    }
  }

  /**
   * Remove a listener to no longer be notified of future changes in the result set.
   */
  fun removeListener(listener: Listener) {
    listenerLock.withLock {
      listeners.remove(listener)
      if (listeners.isEmpty()) queries.remove(this)
    }
  }

  /**
   * Execute the underlying statement.
   *
   * @return the cursor for the statement's result set.
   */
  abstract fun execute(): SqlCursor

  /**
   * @return The result set of the underlying SQL statement as a list of [RowType].
   */
  fun executeAsList(): List<RowType> {
    val result = mutableListOf<RowType>()
    execute().use {
      while (it.next()) result.add(mapper(it))
    }
    return result
  }

  /**
   * @return The only row of the result set for the underlying SQL statement as a non null
   *   [RowType].
   *
   * @throws NullPointerException if when executed this query has no rows in its result set.
   * @throws IllegalStateException if when executed this query has multiple rows in its result set.
   */
  fun executeAsOne(): RowType {
    return executeAsOneOrNull()
      ?: throw NullPointerException("ResultSet returned null for $this")
  }

  /**
   * @return The first row of the result set for the underlying SQL statement as a non null
   *   [RowType] or null if the result set has no rows.
   *
   * @throws IllegalStateException if when executed this query has multiple rows in its result set.
   */
  fun executeAsOneOrNull(): RowType? {
    execute().use {
      if (!it.next()) return null
      val item = mapper(it)
      check(!it.next()) { "ResultSet returned more than 1 row for $this" }
      return item
    }
  }

  /**
   * An interface for listening to changes in the result set of a query.
   */
  interface Listener {
    /**
     * Called whenever the query this listener was attached to is dirtied.
     *
     * Calls are made synchronously on the thread where the updated occurred, after the update applied successfully.
     */
    fun queryResultsChanged()
  }
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight

import com.squareup.sqldelight.Transacter.Transaction
import com.squareup.sqldelight.db.SqlDriver
import com.squareup.sqldelight.internal.currentThreadId

interface TransactionCallbacks {
  fun afterCommit(function: () -> Unit)
  fun afterRollback(function: () -> Unit)
}

interface TransactionWithReturn<R> : TransactionCallbacks {
  /**
   * Rolls back this transaction.
   */
  fun rollback(returnValue: R): Nothing

  /**
   * Begin an inner transaction.
   */
  fun <R> transaction(body: TransactionWithReturn<R>.() -> R): R
}

interface TransactionWithoutReturn : TransactionCallbacks {
  /**
   * Rolls back this transaction.
   */
  fun rollback(): Nothing

  /**
   * Begin an inner transaction.
   */
  fun transaction(body: TransactionWithoutReturn.() -> Unit)
}

/**
 * A transaction-aware [SqlDriver] wrapper which can begin a [Transaction] on the current connection.
 */
interface Transacter {
  /**
   * Starts a [Transaction] and runs [body] in that transaction.
   *
   * @throws IllegalStateException if [noEnclosing] is true and there is already an active
   *   [Transaction] on this thread.
   */
  fun <R> transactionWithResult(
    noEnclosing: Boolean = false,
    bodyWithReturn: TransactionWithReturn<R>.() -> R
  ): R

  /**
   * Starts a [Transaction] and runs [body] in that transaction.
   *
   * @throws IllegalStateException if [noEnclosing] is true and there is already an active
   *   [Transaction] on this thread.
   */
  fun transaction(
    noEnclosing: Boolean = false,
    body: TransactionWithoutReturn.() -> Unit
  )

  /**
   * A SQL transaction. Can be created through the driver via [SqlDriver.newTransaction] or
   * through an implementation of [Transacter] by calling [Transacter.transaction].
   *
   * A transaction is expected never to escape the thread it is created on, or more specifically,
   * never to escape the lambda scope of [Transacter.transaction] and [Transacter.transactionWithResult].
   */
  abstract class Transaction : TransactionCallbacks {
    private val ownerThreadId = currentThreadId()

    internal val postCommitHooks = mutableListOf<() -> Unit>()
    internal val postRollbackHooks = mutableListOf<() -> Unit>()
    internal val queriesFuncs = mutableMapOf<Int, () -> List<Query<*>>>()

    internal var successful: Boolean = false
    internal var childrenSuccessful: Boolean = true
    internal var transacter: Transacter? = null

    /**
     * The parent transaction, if there is any.
     */
    protected abstract val enclosingTransaction: Transaction?

    internal fun enclosingTransaction() = enclosingTransaction

    /**
     * Signal to the underlying SQL driver that this transaction should be finished.
     *
     * @param successful Whether the transaction completed successfully or not.
     */
    protected abstract fun endTransaction(successful: Boolean)

    internal fun endTransaction() {
      checkThreadConfinement()
      endTransaction(successful && childrenSuccessful)
    }
    /**
     * Queues [function] to be run after this transaction successfully commits.
     */

    override fun afterCommit(function: () -> Unit) {
      checkThreadConfinement()
      postCommitHooks.add(function)
    }

    /**
     * Queues [function] to be run after this transaction rolls back.
     */
    override fun afterRollback(function: () -> Unit) {
      checkThreadConfinement()
      postRollbackHooks.add(function)
    }

    internal fun checkThreadConfinement() = check(ownerThreadId == currentThreadId()) {
      """
        Transaction objects (`TransactionWithReturn` and `TransactionWithoutReturn`) must be used
        only within the transaction lambda scope.
      """.trimIndent()
    }
  }
}

private class RollbackException(val value: Any? = null) : Throwable()

private class TransactionWrapper<R>(
  val transaction: Transaction
) : TransactionWithoutReturn, TransactionWithReturn<R> {
  override fun rollback(): Nothing {
    transaction.checkThreadConfinement()
    throw RollbackException()
  }
  override fun rollback(returnValue: R): Nothing {
    transaction.checkThreadConfinement()
    throw RollbackException(returnValue)
  }

  /**
   * Queues [function] to be run after this transaction successfully commits.
   */
  override fun afterCommit(function: () -> Unit) {
    transaction.afterCommit(function)
  }

  /**
   * Queues [function] to be run after this transaction rolls back.
   */
  override fun afterRollback(function: () -> Unit) {
    transaction.afterRollback(function)
  }

  override fun transaction(body: TransactionWithoutReturn.() -> Unit) {
    transaction.transacter!!.transaction(false, body)
  }

  override fun <R> transaction(body: TransactionWithReturn<R>.() -> R): R {
    return transaction.transacter!!.transactionWithResult(false, body)
  }
}

/**
 * A transaction-aware [SqlDriver] wrapper which can begin a [Transaction] on the current connection.
 */
abstract class TransacterImpl(private val driver: SqlDriver) : Transacter {
  /**
   * For internal use, notifies the listeners of [queryList] that their underlying result set has
   * changed.
   */
  protected fun notifyQueries(identifier: Int, queryList: () -> List<Query<*>>) {
    val transaction = driver.currentTransaction()
    if (transaction != null) {
      if (!transaction.queriesFuncs.containsKey(identifier)) {
        transaction.queriesFuncs[identifier] = queryList
      }
    } else {
      queryList.invoke().forEach { it.notifyDataChanged() }
    }
  }

  /**
   * For internal use, creates a string in the format (?, ?, ?) where there are [count] offset.
   */
  protected fun createArguments(count: Int): String {
    if (count == 0) return "()"

    return buildString(count + 2) {
      append("(?")
      repeat(count - 1) {
        append(",?")
      }
      append(')')
    }
  }

  override fun transaction(
    noEnclosing: Boolean,
    body: TransactionWithoutReturn.() -> Unit
  ) {
    transactionWithWrapper<Unit?>(noEnclosing, body)
  }

  override fun <R> transactionWithResult(
    noEnclosing: Boolean,
    bodyWithReturn: TransactionWithReturn<R>.() -> R
  ): R {
    return transactionWithWrapper(noEnclosing, bodyWithReturn)
  }

  private fun <R> transactionWithWrapper(noEnclosing: Boolean, wrapperBody: TransactionWrapper<R>.() -> R): R {
    val transaction = driver.newTransaction()
    val enclosing = transaction.enclosingTransaction()

    check(enclosing == null || !noEnclosing) { "Already in a transaction" }

    var thrownException: Throwable? = null
    var returnValue: R? = null

    try {
      transaction.transacter = this
      returnValue = TransactionWrapper<R>(transaction).wrapperBody()
      transaction.successful = true
    } catch (e: Throwable) {
      thrownException = e
    } finally {
      transaction.endTransaction()
      if (enclosing == null) {
        if (!transaction.successful || !transaction.childrenSuccessful) {
          // TODO: If this throws, and we threw in [body] then create a composite exception.
          try {
            transaction.postRollbackHooks.forEach { it.invoke() }
          } catch (rollbackException: Throwable) {
            thrownException?.let {
              throw Throwable("Exception while rolling back from an exception.\nOriginal exception: $thrownException\nwith cause ${thrownException.cause}\n\nRollback exception: $rollbackException", rollbackException)
            }
            throw rollbackException
          }
          transaction.postRollbackHooks.clear()
        } else {
          transaction.queriesFuncs
            .flatMap { (_, queryListSupplier) -> queryListSupplier.invoke() }
            .distinct()
            .forEach { it.notifyDataChanged() }

          transaction.queriesFuncs.clear()
          transaction.postCommitHooks.forEach { it.invoke() }
          transaction.postCommitHooks.clear()
        }
      } else {
        enclosing.childrenSuccessful = transaction.successful && transaction.childrenSuccessful
        enclosing.postCommitHooks.addAll(transaction.postCommitHooks)
        enclosing.postRollbackHooks.addAll(transaction.postRollbackHooks)
        enclosing.queriesFuncs.putAll(transaction.queriesFuncs)
      }

      if (enclosing == null && thrownException is RollbackException) {
        // We can safely cast to R here because the rollback exception is always created with the
        // correct type.
        @Suppress("UNCHECKED_CAST")
        return thrownException.value as R
      } else if (thrownException != null) {
        throw thrownException
      } else {
        // We can safely cast to R here because any code path that led here will have set the
        // returnValue to the result of the block
        @Suppress("UNCHECKED_CAST")
        return returnValue as R
      }
    }
  }
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.db

/**
 * A type that can be closed.
 */
expect interface Closeable {
  /**
   * Close any resources backed by this object.
   */
  fun close()
}

/**
 * Run [body] on the receiver and call [Closeable.close] before returning or throwing.
 */
expect inline fun <T : Closeable?, R> T.use(body: (T) -> R): R

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.db

/**
 * Represents a SQL result set which can be iterated through with [next]. Initially the cursor will
 * not point to any row, and calling [next] once will iterate to the first row.
 */
interface SqlCursor : Closeable {
  /**
   * Move to the next row in the result set.
   *
   * @return true if the cursor successfully moved to a new row, false if there was no row to
   *   iterate to.
   */
  fun next(): Boolean

  /**
   * @return The string or null value of column [index] for the current row of the result set.
   */
  fun getString(index: Int): String?

  /**
   * @return The int or null value of column [index] for the current row of the result set.
   */
  fun getLong(index: Int): Long?

  /**
   * @return The bytes or null value of column [index] for the current row of the result set.
   */
  fun getBytes(index: Int): ByteArray?

  /**
   * @return The double or null value of column [index] for the current row of the result set.
   */
  fun getDouble(index: Int): Double?
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.db

import com.squareup.sqldelight.Transacter

/**
 * Maintains connections to an underlying SQL database and provides APIs for managing transactions
 * and executing SQL statements.
 */
interface SqlDriver : Closeable {
  /**
   * Execute a SQL statement and return a [SqlCursor] that iterates the result set.
   *
   * @param [identifier] An opaque, unique value that can be used to implement any driver-side
   *   caching of prepared statements. If [identifier] is null, a fresh statement is required.
   * @param [sql] The SQL string to be executed.
   * @param [parameters] The number of bindable parameters [sql] contains.
   * @param [binders] A lambda which is called before execution to bind any parameters to the SQL
   *   statement.
   */
  fun executeQuery(
    identifier: Int?,
    sql: String,
    parameters: Int,
    binders: (SqlPreparedStatement.() -> Unit)? = null
  ): SqlCursor

  /**
   * Execute a SQL statement.
   *
   * @param [identifier] An opaque, unique value that can be used to implement any driver-side
   *   caching of prepared statements. If [identifier] is null, a fresh statement is required.
   * @param [sql] The SQL string to be executed.
   * @param [parameters] The number of bindable parameters [sql] contains.
   * @param [binders] A lambda which is called before execution to bind any parameters to the SQL
   *   statement.
   */
  fun execute(
    identifier: Int?,
    sql: String,
    parameters: Int,
    binders: (SqlPreparedStatement.() -> Unit)? = null
  )

  /**
   * Start a new [Transacter.Transaction] on the database.
   *
   * It's up to the implementor how this method behaves for different connection/threading patterns.
   */
  fun newTransaction(): Transacter.Transaction

  /**
   * The currently open [Transacter.Transaction] on the database.
   *
   * It's up to the implementor how this method behaves for different connection/threading patterns.
   */
  fun currentTransaction(): Transacter.Transaction?

  /**
   * API for creating and migrating a SQL database.
   */
  interface Schema {
    /**
     * The version of this schema.
     */
    val version: Int

    /**
     * Use [driver] to create the schema from scratch. Assumes no existing database state.
     */
    fun create(driver: SqlDriver)

    /**
     * Use [driver] to migrate from schema [oldVersion] to [newVersion].
     */
    fun migrate(driver: SqlDriver, oldVersion: Int, newVersion: Int)
  }
}

/**
 * Represents a block of code [block] that should be executed during a migration after the migration
 * has finished migrating to [afterVersion].
 */
class AfterVersion(
  internal val afterVersion: Int,
  internal val block: () -> Unit
)

/**
 * Run [SqlDriver.Schema.migrate] normally but execute [callbacks] during the migration whenever
 * the it finished upgrading to a version specified by [AfterVersion.afterVersion].
 */
fun SqlDriver.Schema.migrateWithCallbacks(
  driver: SqlDriver,
  oldVersion: Int,
  newVersion: Int,
  vararg callbacks: AfterVersion
) {
  var lastVersion = oldVersion

  // For each callback within the [oldVersion..newVersion) range, alternate between migrating
  // the schema and invoking each callback.
  callbacks.filter { it.afterVersion in oldVersion until newVersion }
    .sortedBy { it.afterVersion }
    .forEach { callback ->
      migrate(driver, oldVersion = lastVersion, newVersion = callback.afterVersion + 1)
      callback.block()
      lastVersion = callback.afterVersion + 1
    }

  // If there were no callbacks, or the newVersion is higher than the highest callback,
  // complete the migration.
  if (lastVersion < newVersion) {
    migrate(driver, lastVersion, newVersion)
  }
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.db

/**
 * Represents a SQL statement that has been prepared by a driver to be executed.
 *
 * This type is not thread safe unless otherwise specified by the driver emitting these.
 *
 * Prepared statements should not be cached by client code. Drivers can implement caching by using
 * the integer identifier passed to [SqlDriver.execute] or [SqlDriver.executeQuery].
 * Client code can pass the same identifier to that method to request that the prepared statement
 * is cached.
 */
interface SqlPreparedStatement {

  /**
   * Bind [bytes] to the underlying statement at [index].
   */
  fun bindBytes(index: Int, bytes: ByteArray?)

  /**
   * Bind [long] to the underlying statement at [index].
   */
  fun bindLong(index: Int, long: Long?)

  /**
   * Bind [double] to the underlying statement at [index].
   */
  fun bindDouble(index: Int, double: Double?)

  /**
   * Bind [string] to the underlying statement at [index].
   */
  fun bindString(index: Int, string: String?)
}

package com.squareup.sqldelight.internal

import kotlin.reflect.KProperty

expect class AtomicBoolean(value: Boolean) {
  fun get(): Boolean
  fun set(value: Boolean)
}

internal operator fun AtomicBoolean.getValue(thisRef: Any?, prop: KProperty<*>): Boolean {
  return get()
}

internal operator fun AtomicBoolean.setValue(thisRef: Any?, prop: KProperty<*>, value: Boolean) {
  set(value)
}

expect class Atomic<V>(value: V) {
  fun get(): V
  fun set(value: V)
}

operator fun <T> Atomic<T>.getValue(thisRef: Any?, prop: KProperty<*>): T {
  return get()
}

operator fun <T> Atomic<T>.setValue(thisRef: Any?, prop: KProperty<*>, value: T) {
  set(value)
}

package com.squareup.sqldelight.internal

import com.squareup.sqldelight.Query

expect fun copyOnWriteList(): MutableList<Query<*>>

internal expect fun <T> copyOnWriteList(): MutableList<T>

internal expect class QueryLock()

internal expect inline fun <T> QueryLock.withLock(block: () -> T): T

internal expect fun currentThreadId(): Long

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.logs

import com.squareup.sqldelight.Transacter
import com.squareup.sqldelight.db.SqlCursor
import com.squareup.sqldelight.db.SqlDriver
import com.squareup.sqldelight.db.SqlPreparedStatement

class LogSqliteDriver(
  private val sqlDriver: SqlDriver,
  private val logger: (String) -> Unit
) : SqlDriver {

  override fun currentTransaction(): Transacter.Transaction? {
    return sqlDriver.currentTransaction()
  }

  override fun execute(
    identifier: Int?,
    sql: String,
    parameters: Int,
    binders: (SqlPreparedStatement.() -> Unit)?
  ) {
    logger("EXECUTE\n $sql")
    logParameters(binders)
    sqlDriver.execute(identifier, sql, parameters, binders)
  }

  override fun executeQuery(
    identifier: Int?,
    sql: String,
    parameters: Int,
    binders: (SqlPreparedStatement.() -> Unit)?
  ): SqlCursor {
    logger("QUERY\n $sql")
    logParameters(binders)
    return sqlDriver.executeQuery(identifier, sql, parameters, binders)
  }

  override fun newTransaction(): Transacter.Transaction {
    logger("TRANSACTION BEGIN")
    val transaction = sqlDriver.newTransaction()
    transaction.afterCommit { logger("TRANSACTION COMMIT") }
    transaction.afterRollback { logger("TRANSACTION ROLLBACK") }
    return transaction
  }

  override fun close() {
    logger("CLOSE CONNECTION")
    sqlDriver.close()
  }

  private fun logParameters(binders: (SqlPreparedStatement.() -> Unit)?) {
    binders?.let { func ->
      val parametersInterceptor = StatementParameterInterceptor()
      parametersInterceptor.func()
      val logParameters = parametersInterceptor.getAndClearParameters()
      if (logParameters.isNotEmpty()) logger(" $logParameters")
    }
  }
}

class StatementParameterInterceptor : SqlPreparedStatement {
  private val values = mutableListOf<Any?>()

  override fun bindBytes(
    index: Int,
    bytes: ByteArray?
  ) {
    values.add(bytes)
  }

  override fun bindDouble(
    index: Int,
    double: Double?
  ) {
    values.add(double)
  }

  override fun bindLong(
    index: Int,
    long: Long?
  ) {
    values.add(long)
  }

  override fun bindString(
    index: Int,
    string: String?
  ) {
    values.add(string)
  }

  fun getAndClearParameters(): List<Any?> {
    val list = values.toList()
    values.clear()
    return list
  }
}

package com.squareup.sqldelight.db

actual interface Closeable {
  actual fun close()
}

actual inline fun <T : Closeable?, R> T.use(body: (T) -> R): R {
  var exception: Throwable? = null
  try {
    return body(this)
  } catch (e: Throwable) {
    exception = e
    throw e
  } finally {
    when {
      this == null -> {}
      exception == null -> close()
      else ->
        try {
          close()
        } catch (closeException: Throwable) {
          // Nothing to do...
        }
    }
  }
}

package com.squareup.sqldelight.internal

actual open class Atomic<V> actual constructor(private var value: V) {
  actual fun get() = value
  actual fun set(value: V) { this.value = value }
}

actual class AtomicBoolean actual constructor(value: Boolean) : Atomic<Boolean>(value)

package com.squareup.sqldelight.internal

import com.squareup.sqldelight.Query

actual fun copyOnWriteList(): MutableList<Query<*>> {
  return mutableListOf()
}

internal actual fun <T> copyOnWriteList(): MutableList<T> {
  return mutableListOf()
}

internal actual class QueryLock

internal actual inline fun <T> QueryLock.withLock(block: () -> T): T {
  return block()
}

internal actual fun currentThreadId(): Long = 0

package com.squareup.sqldelight.db

import kotlin.io.use as kotlinIoUse

actual typealias Closeable = java.io.Closeable

actual inline fun <T : Closeable?, R> T.use(body: (T) -> R): R = kotlinIoUse(body)

package com.squareup.sqldelight.internal

import java.util.concurrent.atomic.AtomicReference

actual typealias AtomicBoolean = java.util.concurrent.atomic.AtomicBoolean

actual typealias Atomic<V> = AtomicReference<V>

package com.squareup.sqldelight.internal

import com.squareup.sqldelight.Query
import java.util.concurrent.CopyOnWriteArrayList

actual fun copyOnWriteList(): MutableList<Query<*>> {
  return CopyOnWriteArrayList()
}

@JvmName("copyOnWriteListGeneric")
internal actual fun <T> copyOnWriteList(): MutableList<T> {
  return CopyOnWriteArrayList()
}

internal actual class QueryLock

internal actual inline fun <T> QueryLock.withLock(block: () -> T): T {
  synchronized(this) {
    return block()
  }
}

internal actual fun currentThreadId(): Long = Thread.currentThread().id

package com.squareup.sqldelight.internal

import co.touchlab.stately.concurrency.value
import kotlin.native.concurrent.AtomicReference
import kotlin.native.concurrent.freeze

actual class Atomic<V> actual constructor(value: V) {
  private val atomicRef = AtomicReference(value)

  actual fun get() = atomicRef.value
  actual fun set(value: V) {
    atomicRef.value = value.freeze()
  }
}

actual class AtomicBoolean actual constructor(value: Boolean) {
  private val atomicBoolean = co.touchlab.stately.concurrency.AtomicBoolean(value)

  actual fun get() = atomicBoolean.value
  actual fun set(value: Boolean) {
    atomicBoolean.value = value
  }
}

package com.squareup.sqldelight.internal

import co.touchlab.stately.collections.frozenCopyOnWriteList
import co.touchlab.stately.concurrency.Lock
import co.touchlab.stately.concurrency.withLock
import com.squareup.sqldelight.Query
import kotlin.native.concurrent.Worker

actual fun copyOnWriteList(): MutableList<Query<*>> {
  return frozenCopyOnWriteList()
}

internal actual fun <T> copyOnWriteList(): MutableList<T> {
  return frozenCopyOnWriteList()
}

internal actual class QueryLock {
  internal val lock = Lock()
}

internal actual inline fun <T> QueryLock.withLock(block: () -> T) = lock.withLock(block)

internal actual fun currentThreadId(): Long = Worker.current.id.toLong()

package com.squareup.sqldelight.core

import com.squareup.sqldelight.VERSION

object GradleCompatibility {
  fun validate(propertiesFile: SqlDelightPropertiesFile): CompatibilityReport {
    val (minimumIdeVersion, currentGradleVersion) = try {
      SemVer(propertiesFile.minimumSupportedVersion) to SemVer(propertiesFile.currentVersion)
    } catch (e: Throwable) {
      // If we can't even read the properties file version, it is not compatibile
      return CompatibilityReport.Incompatible(
        reason = "Plugin 'SQLDelight' is incompatible with the current version of the SQLDelight Gradle plugin. Upgrade the version of com.squareup.sqldelight:gradle-plugin.",
      )
    }

    val currentVersion = SemVer(VERSION)

    if (currentVersion < minimumIdeVersion) {
      return CompatibilityReport.Incompatible(
        reason = "Plugin 'SQLDelight' is no longer compatible with the current version of the SQLDelight Gradle plugin. Use version $minimumIdeVersion or later.",
      )
    }

    val minimumGradleVersion = SemVer(MINIMUM_SUPPORTED_VERSION)

    if (currentGradleVersion < minimumGradleVersion) {
      return CompatibilityReport.Incompatible(
        reason = "Gradle plugin 'com.squareup.sqldelight:gradle-plugin' is incompatible with the current SQLDelight IntelliJ Plugin version. Use version $minimumGradleVersion or later."
      )
    }

    return CompatibilityReport.Compatible
  }

  private data class SemVer(
    val major: Int,
    val minor: Int,
    val patch: Int,
    val snapshot: Boolean,
  ) {
    constructor(semVerText: String) : this(
      major = semVerText.split('.')[0].toInt(),
      minor = semVerText.split('.')[1].toInt(),
      patch = semVerText.split('.')[2].removeSuffix("-SNAPSHOT").toInt(),
      snapshot = semVerText.endsWith("-SNAPSHOT"),
    )

    operator fun compareTo(other: SemVer): Int {
      val majorCompare = major.compareTo(other.major)
      if (majorCompare != 0) return majorCompare
      val minorCompare = minor.compareTo(other.minor)
      if (minorCompare != 0) return minorCompare
      return patch.compareTo(other.patch)
    }

    override fun toString() = "$major.$minor.$patch"
  }

  sealed class CompatibilityReport {
    object Compatible : CompatibilityReport()
    data class Incompatible(
      val reason: String
    ) : CompatibilityReport()
  }
}

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.core

import com.alecstrong.sql.psi.core.DialectPreset
import com.alecstrong.sql.psi.core.SqlAnnotationHolder
import com.alecstrong.sql.psi.core.SqlCoreEnvironment
import com.alecstrong.sql.psi.core.SqlFileBase
import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.alecstrong.sql.psi.core.psi.SqlStmt
import com.intellij.core.CoreApplicationEnvironment
import com.intellij.mock.MockModule
import com.intellij.openapi.module.Module
import com.intellij.openapi.project.Project
import com.intellij.openapi.roots.ModuleExtension
import com.intellij.openapi.roots.ModuleRootManager
import com.intellij.openapi.roots.impl.ModuleRootManagerImpl
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.FileTypeFileViewProviders
import com.intellij.psi.PsiDirectory
import com.intellij.psi.PsiDocumentManager
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiErrorElement
import com.intellij.psi.PsiFileSystemItem
import com.intellij.psi.PsiManager
import com.intellij.psi.util.PsiTreeUtil
import com.squareup.sqldelight.core.compiler.SqlDelightCompiler
import com.squareup.sqldelight.core.lang.DatabaseFileType
import com.squareup.sqldelight.core.lang.DatabaseFileViewProviderFactory
import com.squareup.sqldelight.core.lang.MigrationFile
import com.squareup.sqldelight.core.lang.MigrationFileType
import com.squareup.sqldelight.core.lang.MigrationParserDefinition
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.SqlDelightFileType
import com.squareup.sqldelight.core.lang.SqlDelightParserDefinition
import com.squareup.sqldelight.core.lang.SqlDelightQueriesFile
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.lang.util.sqFile
import com.squareup.sqldelight.core.psi.SqlDelightImportStmt
import org.picocontainer.MutablePicoContainer
import java.io.File
import java.util.StringTokenizer
import kotlin.math.log10
import kotlin.system.measureTimeMillis

/**
 * Mocks an intellij environment for compiling sqldelight files without an instance of intellij
 * running.
 */
class SqlDelightEnvironment(
  /**
   * The package name to be used for the generated SqlDelightDatabase class.
   */
  private val properties: SqlDelightDatabaseProperties,
  /**
   * The package name to be used for the generated SqlDelightDatabase class.
   */
  private val compilationUnit: SqlDelightCompilationUnit,
  /**
   * If true, fail migrations during compilation when there are errors.
   */
  private val verifyMigrations: Boolean,
  moduleName: String,
  private val sourceFolders: List<File> = compilationUnit.sourceFolders
    .filter { it.folder.exists() && !it.dependency }
    .map { it.folder },
  private val dependencyFolders: List<File> = compilationUnit.sourceFolders
    .filter { it.folder.exists() && it.dependency }
    .map { it.folder },
) : SqlCoreEnvironment(sourceFolders, dependencyFolders),
  SqlDelightProjectService {
  val project: Project = projectEnvironment.project
  val module = MockModule(project, projectEnvironment.parentDisposable)
  private val moduleName = SqlDelightFileIndex.sanitizeDirectoryName(moduleName)

  init {
    (project.picoContainer as MutablePicoContainer).registerComponentInstance(SqlDelightProjectService::class.java.name, this)

    CoreApplicationEnvironment.registerExtensionPoint(
      module.extensionArea,
      ModuleExtension.EP_NAME, ModuleExtension::class.java
    )
    module.picoContainer.registerComponentInstance(
      ModuleRootManager::class.java.name,
      ModuleRootManagerImpl(module)
    )

    initializeApplication {
      registerFileType(MigrationFileType, MigrationFileType.defaultExtension)
      registerParserDefinition(MigrationParserDefinition())
      registerFileType(SqlDelightFileType, SqlDelightFileType.defaultExtension)
      registerParserDefinition(SqlDelightParserDefinition())
      registerFileType(DatabaseFileType, DatabaseFileType.defaultExtension)
      FileTypeFileViewProviders.INSTANCE.addExplicitExtension(DatabaseFileType, DatabaseFileViewProviderFactory())
    }
  }

  override fun module(vFile: VirtualFile) = module

  override fun fileIndex(module: Module): SqlDelightFileIndex = FileIndex()

  override fun resetIndex() = throw UnsupportedOperationException()

  override var dialectPreset: DialectPreset
    get() = properties.dialectPreset
    set(_) { throw UnsupportedOperationException() }

  override fun forSourceFiles(action: (SqlFileBase) -> Unit) {
    super.forSourceFiles {
      if (it.fileType != MigrationFileType ||
        verifyMigrations ||
        properties.deriveSchemaFromMigrations
      ) {
        action(it)
      }
    }
  }

  /**
   * Run the SQLDelight compiler and return the error or success status.
   */
  fun generateSqlDelightFiles(logger: (String) -> Unit): CompilationStatus {
    val errors = sortedMapOf<Int, MutableList<String>>()
    annotate(object : SqlAnnotationHolder {
      override fun createErrorAnnotation(element: PsiElement, s: String) {
        val key = element.sqFile().order ?: Integer.MAX_VALUE
        errors.putIfAbsent(key, ArrayList())
        errors[key]!!.add(errorMessage(element, s))
      }
    })
    if (errors.isNotEmpty()) return CompilationStatus.Failure(errors.values.flatten())

    val writer = writer@{ fileName: String ->
      val file = File(fileName)
      if (!file.exists()) {
        file.parentFile.mkdirs()
        file.createNewFile()
      }
      return@writer file.writer()
    }

    var sourceFile: SqlDelightFile? = null
    var topMigrationFile: MigrationFile? = null
    forSourceFiles {
      if (it is MigrationFile && properties.deriveSchemaFromMigrations) {
        if (topMigrationFile == null || it.order > topMigrationFile!!.order) topMigrationFile = it
        if (sourceFile == null) sourceFile = it
      }

      if (it !is SqlDelightQueriesFile) return@forSourceFiles
      logger("----- START ${it.name} ms -------")
      val timeTaken = measureTimeMillis {
        SqlDelightCompiler.writeInterfaces(module, it, writer)
        sourceFile = it
      }
      logger("----- END ${it.name} in $timeTaken ms ------")
    }

    topMigrationFile?.let { migrationFile ->
      logger("----- START ${migrationFile.name} ms -------")
      val timeTaken = measureTimeMillis {
        SqlDelightCompiler.writeInterfaces(
          file = migrationFile,
          output = writer,
          includeAll = true
        )
      }
      logger("----- END ${migrationFile.name} in $timeTaken ms ------")
    }

    sourceFile?.let {
      SqlDelightCompiler.writeDatabaseInterface(module, it, moduleName, writer)
      if (it is SqlDelightQueriesFile)
        SqlDelightCompiler.writeImplementations(module, it, moduleName, writer)
    }

    return CompilationStatus.Success()
  }

  fun forMigrationFiles(body: (MigrationFile) -> Unit) {
    val psiManager = PsiManager.getInstance(projectEnvironment.project)
    val migrationFiles: Collection<MigrationFile> = sourceFolders
      .map { localFileSystem.findFileByPath(it.absolutePath)!! }
      .map { psiManager.findDirectory(it)!! }
      .flatMap { directory: PsiDirectory -> directory.findChildrenOfType<MigrationFile>().asIterable() }
    migrationFiles.sortedBy { it.version }
      .forEach {
        val errorElements = ArrayList<PsiErrorElement>()
        PsiTreeUtil.processElements(it) { element ->
          when (element) {
            is PsiErrorElement -> errorElements.add(element)
            // Uncomment when sqm files understand their state of the world.
            // is SqlAnnotatedElement -> element.annotate(annotationHolder)
          }
          return@processElements true
        }
        if (errorElements.isNotEmpty()) {
          throw SqlDelightException(
            "Error Reading ${it.name}:\n\n" +
              errorElements.joinToString(separator = "\n") { errorMessage(it, it.errorDescription) }
          )
        }
        body(it)
      }
  }

  private fun errorMessage(element: PsiElement, message: String): String {
    return "${element.containingFile.virtualFile.path} " +
      "line ${element.lineStart}:${element.charPositionInLine} - $message\n${detailText(element)}"
  }

  private fun detailText(element: PsiElement) = try {
    val context = context(element) ?: element
    val result = StringBuilder()
    val tokenizer = StringTokenizer(context.text, "\n", false)

    val maxDigits = (log10(context.lineEnd.toDouble()) + 1).toInt()
    for (line in context.lineStart..context.lineEnd) {
      if (!tokenizer.hasMoreTokens()) break
      result.append(("%0${maxDigits}d    %s\n").format(line, tokenizer.nextToken()))
      if (element.lineStart == element.lineEnd && element.lineStart == line) {
        // If its an error on a single line highlight where on the line.
        result.append(("%${maxDigits}s    ").format(""))
        if (element.charPositionInLine > 0) {
          result.append(("%${element.charPositionInLine}s").format(""))
        }
        result.append(("%s\n").format("^".repeat(element.textLength)))
      }
    }

    result.toString()
  } catch (e: Exception) {
    // If there is an exception while trying to print an error, just give back the unformatted error
    // and print the stack trace for more debugging.
    e.printStackTrace()
    element.text
  }

  private val PsiElement.charPositionInLine: Int
    get() {
      val file = PsiDocumentManager.getInstance(project).getDocument(containingFile)!!
      return textOffset - file.getLineStartOffset(file.getLineNumber(textOffset))
    }

  private val PsiElement.lineStart: Int
    get() {
      val file = PsiDocumentManager.getInstance(project).getDocument(containingFile)!!
      return file.getLineNumber(textOffset) + 1
    }

  private val PsiElement.lineEnd: Int
    get() {
      val file = PsiDocumentManager.getInstance(project).getDocument(containingFile)!!
      return file.getLineNumber(textOffset + textLength) + 1
    }

  private fun context(element: PsiElement?): PsiElement? =
    when (element) {
      null -> element
      is SqlCreateTableStmt -> element
      is SqlStmt -> element
      is SqlDelightImportStmt -> element
      else -> context(element.parent)
    }

  sealed class CompilationStatus {
    class Success : CompilationStatus()
    class Failure(val errors: List<String>) : CompilationStatus()
  }

  private inner class FileIndex : SqlDelightFileIndex {
    override val contentRoot
      get() = throw UnsupportedOperationException("Content root only usable from IDE")

    override val packageName = properties.packageName
    override val className = properties.className
    override val dependencies = properties.dependencies
    override val isConfigured = true
    override val deriveSchemaFromMigrations = properties.deriveSchemaFromMigrations

    override fun outputDirectory(file: SqlDelightFile) = outputDirectories()
    override fun outputDirectories(): List<String> {
      return listOf(compilationUnit.outputDirectoryFile.absolutePath)
    }

    private val virtualDirectoriesWithDependencies: List<VirtualFile> by lazy {
      return@lazy (sourceFolders + dependencyFolders)
        .map { localFileSystem.findFileByPath(it.absolutePath)!! }
    }

    private val directoriesWithDependencies: List<PsiDirectory> by lazy {
      val psiManager = PsiManager.getInstance(projectEnvironment.project)
      return@lazy virtualDirectoriesWithDependencies.map { psiManager.findDirectory(it)!! }
    }

    private val virtualDirectories: List<VirtualFile> by lazy {
      return@lazy sourceFolders
        .map { localFileSystem.findFileByPath(it.absolutePath)!! }
    }

    private val directories: List<PsiDirectory> by lazy {
      val psiManager = PsiManager.getInstance(projectEnvironment.project)
      return@lazy virtualDirectories.map { psiManager.findDirectory(it)!! }
    }

    override fun packageName(file: SqlDelightFile): String {
      fun PsiFileSystemItem.relativePathUnder(ancestor: PsiDirectory): List<String>? {
        if (this.virtualFile.path == ancestor.virtualFile.path) return emptyList()
        parent?.let {
          return it.relativePathUnder(ancestor)?.plus(name)
        }
        return null
      }

      for (sourceFolder in sourceFolders(file)) {
        val path = file.parent!!.relativePathUnder(sourceFolder)
        if (path != null) return path.joinToString(separator = ".") { SqlDelightFileIndex.sanitizeDirectoryName(it) }
      }

      throw IllegalStateException(
        "Tried to find package name of file ${file.virtualFile!!.path} when" +
          " it is not under any of the source folders $sourceFolders"
      )
    }

    override fun sourceFolders(file: VirtualFile, includeDependencies: Boolean) =
      if (includeDependencies) virtualDirectoriesWithDependencies else virtualDirectories

    override fun sourceFolders(file: SqlDelightFile, includeDependencies: Boolean) =
      if (includeDependencies) directoriesWithDependencies else directories
  }
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.core

class SqlDelightException(message: String) : IllegalStateException(message)

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.core

import com.intellij.openapi.module.Module
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiDirectory
import com.squareup.sqldelight.core.lang.SqlDelightFile

interface SqlDelightFileIndex {
  /**
   * @return true if this index is configured to be used by SqlDelight.
   */
  val isConfigured: Boolean

  /**
   * @return The package name for the whole source set. This is equivalent to the package name
   * found in the manifest file for the current variant.
   */
  val packageName: String

  /**
   * @return The package name for the generated type which holds all the query files.
   */
  val className: String

  /**
   * @return The list of fully qualified classnames of databases this module depends on.
   */
  val dependencies: List<SqlDelightDatabaseName>

  /**
   * @return The content root for the [Module] backing this index.
   */
  val contentRoot: VirtualFile

  val deriveSchemaFromMigrations: Boolean

  /**
   * @return The package name for a given SqlDelight file. Equal to the relative path under its
   * fixture's sqldelight directory.
   */
  fun packageName(file: SqlDelightFile): String

  /**
   * @return A list of output directory paths generated code should be placed in, relative to
   *   [contentRoot], for the given [file].
   */
  fun outputDirectory(file: SqlDelightFile): List<String>

  /**
   * @return A list of all SQLDelight output directories.
   */
  fun outputDirectories(): List<String>

  /**
   * @return The source roots of sqldelight files for [file].
   */
  fun sourceFolders(file: VirtualFile, includeDependencies: Boolean = true): Collection<VirtualFile>

  /**
   * @return The source roots of sqldelight files for [file].
   */
  fun sourceFolders(file: SqlDelightFile, includeDependencies: Boolean = true): Collection<PsiDirectory>

  companion object {
    fun getInstance(module: Module) = SqlDelightProjectService.getInstance(module.project).fileIndex(module)

    fun sanitizeDirectoryName(name: String): String {
      return name.filter { it.isLetterOrDigit() }
    }
  }
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.core

import com.alecstrong.sql.psi.core.DialectPreset
import com.intellij.openapi.components.ServiceManager
import com.intellij.openapi.module.Module
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.VirtualFile

interface SqlDelightProjectService {
  var dialectPreset: DialectPreset

  fun module(vFile: VirtualFile): Module?

  fun fileIndex(module: Module): SqlDelightFileIndex

  fun resetIndex()

  companion object {
    fun getInstance(project: Project): SqlDelightProjectService {
      return ServiceManager.getService(project, SqlDelightProjectService::class.java)!!
    }
  }
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.core

import com.alecstrong.sql.psi.core.DialectPreset
import java.io.File
import java.io.Serializable

val MINIMUM_SUPPORTED_VERSION = "1.5.0"

interface SqlDelightPropertiesFile : Serializable {
  val databases: List<SqlDelightDatabaseProperties>
  val minimumSupportedVersion: String
  val currentVersion: String
}

interface SqlDelightDatabaseProperties : Serializable {
  val packageName: String
  val compilationUnits: List<SqlDelightCompilationUnit>
  val className: String
  val dependencies: List<SqlDelightDatabaseName>
  val dialectPresetName: String
  val deriveSchemaFromMigrations: Boolean
  val rootDirectory: File
}

val SqlDelightDatabaseProperties.dialectPreset: DialectPreset
  get() = DialectPreset.valueOf(dialectPresetName)

/**
 * A compilation unit represents the group of .sq files which will be compiled all at once. A
 * single database can have multiple compilation units, depending on which gradle task is invoked.
 *
 * For example, a multiplatform module has separate compilation units for ios and android. An
 * android module has separate compilation units for different variants. Only one compilation unit
 * will be worked on during compilation time.
 */
interface SqlDelightCompilationUnit : Serializable {
  val name: String
  val sourceFolders: List<SqlDelightSourceFolder>
  val outputDirectoryFile: File
}

interface SqlDelightSourceFolder : Serializable {
  val folder: File
  val dependency: Boolean
}

interface SqlDelightDatabaseName : Serializable {
  val packageName: String
  val className: String
}

package com.squareup.sqldelight.gradle

import com.squareup.sqldelight.VERSION
import com.squareup.sqldelight.core.SqlDelightCompilationUnit
import com.squareup.sqldelight.core.SqlDelightDatabaseProperties
import com.squareup.sqldelight.core.SqlDelightEnvironment
import com.squareup.sqldelight.core.lang.util.rawSqlText
import org.gradle.api.file.DirectoryProperty
import org.gradle.api.file.FileTree
import org.gradle.api.provider.Property
import org.gradle.api.tasks.CacheableTask
import org.gradle.api.tasks.Input
import org.gradle.api.tasks.InputFiles
import org.gradle.api.tasks.Nested
import org.gradle.api.tasks.OutputDirectory
import org.gradle.api.tasks.PathSensitive
import org.gradle.api.tasks.PathSensitivity
import org.gradle.api.tasks.SkipWhenEmpty
import org.gradle.api.tasks.TaskAction
import org.gradle.workers.WorkAction
import org.gradle.workers.WorkParameters
import java.io.File

@Suppress("UnstableApiUsage") // Worker API.
@CacheableTask
abstract class GenerateMigrationOutputTask : SqlDelightWorkerTask() {
  @Suppress("unused") // Required to invalidate the task on version updates.
  @Input val pluginVersion = VERSION

  @get:OutputDirectory
  var outputDirectory: File? = null

  @Input val projectName: Property<String> = project.objects.property(String::class.java)

  @Nested lateinit var properties: SqlDelightDatabasePropertiesImpl
  @Nested lateinit var compilationUnit: SqlDelightCompilationUnitImpl
  @Input lateinit var migrationOutputExtension: String

  @TaskAction
  fun generateSchemaFile() {
    workQueue().submit(GenerateMigration::class.java) {
      it.outputDirectory.set(outputDirectory)
      it.moduleName.set(projectName)
      it.properties.set(properties)
      it.migrationExtension.set(migrationOutputExtension)
      it.compilationUnit.set(compilationUnit)
    }
  }

  @InputFiles
  @SkipWhenEmpty
  @PathSensitive(PathSensitivity.RELATIVE)
  override fun getSource(): FileTree {
    return super.getSource()
  }

  interface GenerateSchemaWorkParameters : WorkParameters {
    val outputDirectory: DirectoryProperty
    val moduleName: Property<String>
    val properties: Property<SqlDelightDatabaseProperties>
    val compilationUnit: Property<SqlDelightCompilationUnit>
    val migrationExtension: Property<String>
  }

  abstract class GenerateMigration : WorkAction<GenerateSchemaWorkParameters> {

    private val sourceFolders: List<File>
      get() = parameters.compilationUnit.get().sourceFolders.map { it.folder }

    override fun execute() {
      val properties = parameters.properties.get()
      val environment = SqlDelightEnvironment(
        sourceFolders = sourceFolders.filter { it.exists() },
        dependencyFolders = emptyList(),
        moduleName = parameters.moduleName.get(),
        properties = properties,
        verifyMigrations = false,
        compilationUnit = parameters.compilationUnit.get(),
      )

      val outputDirectory = parameters.outputDirectory.get().asFile
      val migrationExtension = parameters.migrationExtension.get()

      // Clear out the output directory.
      outputDirectory.listFiles()?.forEach { it.delete() }

      // Generate the new files.
      environment.forMigrationFiles { migrationFile ->
        val output = File(
          outputDirectory,
          "${migrationFile.virtualFile!!.nameWithoutExtension}$migrationExtension"
        )
        output.writeText(
          migrationFile.sqlStmtList?.stmtList.orEmpty()
            .filterNotNull().joinToString(separator = "\n\n") { "${it.rawSqlText()};" }
        )
      }
    }
  }
}

package com.squareup.sqldelight.gradle

import com.squareup.sqldelight.VERSION
import com.squareup.sqldelight.core.SqlDelightCompilationUnit
import com.squareup.sqldelight.core.SqlDelightDatabaseProperties
import com.squareup.sqldelight.core.SqlDelightEnvironment
import com.squareup.sqldelight.core.lang.SqlDelightQueriesFile
import com.squareup.sqldelight.core.lang.util.forInitializationStatements
import org.gradle.api.file.DirectoryProperty
import org.gradle.api.file.FileTree
import org.gradle.api.provider.Property
import org.gradle.api.tasks.CacheableTask
import org.gradle.api.tasks.Input
import org.gradle.api.tasks.InputFiles
import org.gradle.api.tasks.Nested
import org.gradle.api.tasks.OutputDirectory
import org.gradle.api.tasks.PathSensitive
import org.gradle.api.tasks.PathSensitivity
import org.gradle.api.tasks.SkipWhenEmpty
import org.gradle.api.tasks.TaskAction
import org.gradle.workers.WorkAction
import org.gradle.workers.WorkParameters
import java.io.File
import java.sql.Connection
import java.sql.DriverManager
import java.sql.SQLException

@Suppress("UnstableApiUsage") // Worker API
@CacheableTask
abstract class GenerateSchemaTask : SqlDelightWorkerTask() {
  @Suppress("unused") // Required to invalidate the task on version updates.
  @Input val pluginVersion = VERSION

  @get:OutputDirectory
  var outputDirectory: File? = null

  @Input val projectName: Property<String> = project.objects.property(String::class.java)

  @Nested lateinit var properties: SqlDelightDatabasePropertiesImpl
  @Nested lateinit var compilationUnit: SqlDelightCompilationUnitImpl

  @Input var verifyMigrations: Boolean = false

  @TaskAction
  fun generateSchemaFile() {
    workQueue().submit(GenerateSchema::class.java) {
      it.outputDirectory.set(outputDirectory)
      it.moduleName.set(projectName)
      it.properties.set(properties)
      it.verifyMigrations.set(verifyMigrations)
      it.compilationUnit.set(compilationUnit)
    }
  }

  @InputFiles
  @SkipWhenEmpty
  @PathSensitive(PathSensitivity.RELATIVE)
  override fun getSource(): FileTree {
    return super.getSource()
  }

  interface GenerateSchemaWorkParameters : WorkParameters {
    val outputDirectory: DirectoryProperty
    val moduleName: Property<String>
    val properties: Property<SqlDelightDatabaseProperties>
    val compilationUnit: Property<SqlDelightCompilationUnit>
    val verifyMigrations: Property<Boolean>
  }

  abstract class GenerateSchema : WorkAction<GenerateSchemaWorkParameters> {

    private val sourceFolders: List<File>
      get() = parameters.compilationUnit.get().sourceFolders.map { it.folder }

    override fun execute() {
      val environment = SqlDelightEnvironment(
        sourceFolders = sourceFolders.filter { it.exists() },
        dependencyFolders = emptyList(),
        moduleName = parameters.moduleName.get(),
        properties = parameters.properties.get(),
        verifyMigrations = parameters.verifyMigrations.get(),
        compilationUnit = parameters.compilationUnit.get(),
      )

      var maxVersion = 1
      environment.forMigrationFiles { migrationFile ->
        maxVersion = maxOf(maxVersion, migrationFile.version + 1)
      }

      val outputDirectory = parameters.outputDirectory.get().asFile
      File("$outputDirectory/$maxVersion.db").apply {
        if (exists()) delete()
      }
      createConnection("$outputDirectory/$maxVersion.db").use { connection ->
        val sourceFiles = ArrayList<SqlDelightQueriesFile>()
        environment.forSourceFiles { file ->
          if (file is SqlDelightQueriesFile) sourceFiles.add(file)
        }
        sourceFiles.forInitializationStatements { sqlText ->
          connection.prepareStatement(sqlText).execute()
        }
      }
    }

    private fun createConnection(path: String): Connection {
      return try {
        DriverManager.getConnection("jdbc:sqlite:$path")
      } catch (e: SQLException) {
        DriverManager.getConnection("jdbc:sqlite:$path")
      }
    }
  }
}

package com.squareup.sqldelight.gradle

import com.alecstrong.sql.psi.core.DialectPreset
import com.squareup.sqldelight.core.SqlDelightCompilationUnit
import com.squareup.sqldelight.core.SqlDelightDatabaseName
import com.squareup.sqldelight.core.SqlDelightDatabaseProperties
import com.squareup.sqldelight.core.SqlDelightPropertiesFile
import com.squareup.sqldelight.core.SqlDelightSourceFolder
import org.gradle.api.tasks.Input
import org.gradle.api.tasks.Internal
import org.gradle.api.tasks.Nested
import java.io.File

data class SqlDelightPropertiesFileImpl(
  override val databases: List<SqlDelightDatabasePropertiesImpl>,
  override val minimumSupportedVersion: String,
  override val currentVersion: String,
) : SqlDelightPropertiesFile

data class SqlDelightDatabasePropertiesImpl(
  @Input override val packageName: String,
  @Nested override val compilationUnits: List<SqlDelightCompilationUnitImpl>,
  @Input override val className: String,
  @Nested override val dependencies: List<SqlDelightDatabaseNameImpl>,
  @Input override val dialectPresetName: String = DialectPreset.SQLITE_3_18.name,
  @Input override val deriveSchemaFromMigrations: Boolean = false,
  // Only used by intellij plugin to help with resolution.
  @Internal override val rootDirectory: File
) : SqlDelightDatabaseProperties

data class SqlDelightDatabaseNameImpl(
  @Input override val packageName: String,
  @Input override val className: String
) : SqlDelightDatabaseName

data class SqlDelightCompilationUnitImpl(
  @Input override val name: String,
  @Nested override val sourceFolders: List<SqlDelightSourceFolderImpl>,
  // Output directory is already cached [SqlDelightTask.outputDirectory].
  @Internal override val outputDirectoryFile: File,
) : SqlDelightCompilationUnit

data class SqlDelightSourceFolderImpl(
  // Sources are already cached [SqlDelightTask.getSources]
  @Internal override val folder: File,
  @Input override val dependency: Boolean = false
) : SqlDelightSourceFolder

package com.squareup.sqldelight.gradle

import com.alecstrong.sql.psi.core.DialectPreset
import com.android.builder.model.AndroidProject.FD_GENERATED
import com.squareup.sqldelight.core.lang.MigrationFileType
import com.squareup.sqldelight.core.lang.SqlDelightFileType
import com.squareup.sqldelight.gradle.kotlin.Source
import com.squareup.sqldelight.gradle.kotlin.sources
import groovy.lang.GroovyObject
import org.gradle.api.GradleException
import org.gradle.api.Project
import java.io.File

class SqlDelightDatabase(
  val project: Project,
  var name: String,
  var packageName: String? = null,
  var schemaOutputDirectory: File? = null,
  var sourceFolders: Collection<String>? = null,
  var dialect: String? = null,
  var deriveSchemaFromMigrations: Boolean = false,
  var verifyMigrations: Boolean = false,
  var migrationOutputDirectory: File? = null,
  var migrationOutputFileFormat: String = ".sql"
) {
  private val generatedSourcesDirectory
    get() = File(project.buildDir, "$FD_GENERATED/sqldelight/code/$name")

  private val sources by lazy { sources() }
  private val dependencies = mutableListOf<SqlDelightDatabase>()

  private var recursionGuard = false

  fun methodMissing(name: String, args: Any): Any {
    return (project as GroovyObject).invokeMethod(name, args)
  }

  @Suppress("unused") // Public API used in gradle files.
  fun dependency(dependencyProject: Project) {
    project.evaluationDependsOn(dependencyProject.path)

    val dependency = dependencyProject.extensions.findByType(SqlDelightExtension::class.java)
      ?: throw IllegalStateException("Cannot depend on a module with no sqldelight plugin.")
    val database = dependency.databases.singleOrNull { it.name == name }
      ?: throw IllegalStateException("No database named $name in $dependencyProject")
    check(database.packageName != packageName) { "Detected a schema that already has the package name $packageName in project $dependencyProject" }
    dependencies.add(database)
  }

  internal fun getProperties(): SqlDelightDatabasePropertiesImpl {
    val packageName = requireNotNull(packageName) { "property packageName for $name database must be provided" }

    check(!recursionGuard) { "Found a circular dependency in $project with database $name" }
    recursionGuard = true

    val dialectMapping = mapOf(
      "sqlite:3.18" to DialectPreset.SQLITE_3_18,
      "sqlite:3.24" to DialectPreset.SQLITE_3_24,
      "sqlite:3.25" to DialectPreset.SQLITE_3_25,
      "mysql" to DialectPreset.MYSQL,
      "postgresql" to DialectPreset.POSTGRESQL,
      "hsql" to DialectPreset.HSQL
    )

    if (dialect == null) throw GradleException(
      """
      A dialect is needed for SQLDelight. For example for sqlite:

      sqldelight {
        $name {
          dialect = "sqlite:3.18"
        }
      }
      """.trimIndent()
    )

    val dialect = dialectMapping[dialect]
      ?: throw GradleException("The dialect $dialect is not supported. Supported dialects: ${dialectMapping.keys.joinToString()}.")

    try {
      return SqlDelightDatabasePropertiesImpl(
        packageName = packageName,
        compilationUnits = sources.map { source ->
          SqlDelightCompilationUnitImpl(
            name = source.name,
            sourceFolders = sourceFolders(source).sortedBy { it.folder.absolutePath },
            outputDirectoryFile = source.outputDir,
          )
        },
        rootDirectory = project.projectDir,
        className = name,
        dependencies = dependencies.map { SqlDelightDatabaseNameImpl(it.packageName!!, it.name) },
        dialectPresetName = dialect.name,
        deriveSchemaFromMigrations = deriveSchemaFromMigrations
      )
    } finally {
      recursionGuard = false
    }
  }

  private fun sourceFolders(source: Source): List<SqlDelightSourceFolderImpl> {
    val sourceFolders = sourceFolders ?: listOf("sqldelight")

    val relativeSourceFolders = sourceFolders.flatMap { folder ->
      source.sourceSets.map {
        SqlDelightSourceFolderImpl(
          folder = File(project.projectDir, "src/$it/$folder"),
          dependency = false
        )
      }
    }

    return relativeSourceFolders + dependencies.flatMap { dependency ->
      val dependencySource = source.closestMatch(dependency.sources)
        ?: return@flatMap emptyList<SqlDelightSourceFolderImpl>()
      val compilationUnit = dependency.getProperties().compilationUnits
        .single { it.name == dependencySource.name }

      return@flatMap compilationUnit.sourceFolders.map {
        SqlDelightSourceFolderImpl(
          folder = File(project.projectDir, project.relativePath(it.folder.absolutePath)),
          dependency = true
        )
      }
    }.distinct()
  }

  internal fun registerTasks() {
    sources.forEach { source ->
      // Add the source dependency on the generated code.
      source.sourceDirectorySet.srcDir(source.outputDir.toRelativeString(project.projectDir))

      val allFiles = sourceFolders(source)
      val sourceFiles = project.files(*allFiles.filter { !it.dependency }.map { it.folder }.toTypedArray())
      val dependencyFiles = project.files(*allFiles.filter { it.dependency }.map { it.folder }.toTypedArray())

      // Register the sqldelight generating task.
      val task = project.tasks.register("generate${source.name.capitalize()}${name}Interface", SqlDelightTask::class.java) {
        it.projectName.set(project.name)
        it.properties = getProperties()
        it.compilationUnit = getProperties().compilationUnits.single { it.name == source.name }
        it.outputDirectory = source.outputDir
        it.source(sourceFiles + dependencyFiles)
        it.include("**${File.separatorChar}*.${SqlDelightFileType.defaultExtension}")
        it.include("**${File.separatorChar}*.${MigrationFileType.defaultExtension}")
        it.group = SqlDelightPlugin.GROUP
        it.description = "Generate ${source.name} Kotlin interface for $name"
        it.verifyMigrations = verifyMigrations
      }

      project.tasks.named("generateSqlDelightInterface").configure {
        it.dependsOn(task)
      }

      // Register the task as a dependency of source compilation.
      source.registerTaskDependency(task)

      if (!deriveSchemaFromMigrations) {
        addMigrationTasks(sourceFiles.files + dependencyFiles.files, source)
      } else if (migrationOutputDirectory != null) {
        addMigrationOutputTasks(sourceFiles.files + dependencyFiles.files, source)
      }
    }
  }

  private fun addMigrationTasks(
    sourceSet: Collection<File>,
    source: Source
  ) {
    val verifyMigrationTask =
      project.tasks.register("verify${source.name.capitalize()}${name}Migration", VerifyMigrationTask::class.java) {
        it.projectName.set(project.name)
        it.compilationUnit = getProperties().compilationUnits.single { it.name == source.name }
        it.source(sourceSet)
        it.include("**${File.separatorChar}*.${SqlDelightFileType.defaultExtension}")
        it.include("**${File.separatorChar}*.${MigrationFileType.defaultExtension}")
        it.workingDirectory = File(project.buildDir, "sqldelight/migration_verification/${source.name.capitalize()}$name")
        it.group = SqlDelightPlugin.GROUP
        it.description = "Verify ${source.name} $name migrations and CREATE statements match."
        it.properties = getProperties()
        it.verifyMigrations = verifyMigrations
      }

    if (schemaOutputDirectory != null) {
      project.tasks.register("generate${source.name.capitalize()}${name}Schema", GenerateSchemaTask::class.java) {
        it.projectName.set(project.name)
        it.compilationUnit = getProperties().compilationUnits.single { it.name == source.name }
        it.outputDirectory = schemaOutputDirectory
        it.source(sourceSet)
        it.include("**${File.separatorChar}*.${SqlDelightFileType.defaultExtension}")
        it.include("**${File.separatorChar}*.${MigrationFileType.defaultExtension}")
        it.group = SqlDelightPlugin.GROUP
        it.description = "Generate a .db file containing the current $name schema for ${source.name}."
        it.properties = getProperties()
        it.verifyMigrations = verifyMigrations
      }
    }
    project.tasks.named("check").configure {
      it.dependsOn(verifyMigrationTask)
    }
    project.tasks.named("verifySqlDelightMigration").configure {
      it.dependsOn(verifyMigrationTask)
    }
  }

  private fun addMigrationOutputTasks(
    sourceSet: Collection<File>,
    source: Source
  ) {
    project.tasks.register("generate${source.name.capitalize()}${name}Migrations", GenerateMigrationOutputTask::class.java) {
      it.projectName.set(project.name)
      it.compilationUnit = getProperties().compilationUnits.single { it.name == source.name }
      it.source(sourceSet)
      it.include("**${File.separatorChar}*.${MigrationFileType.defaultExtension}")
      it.migrationOutputExtension = migrationOutputFileFormat
      it.outputDirectory = migrationOutputDirectory
      it.group = SqlDelightPlugin.GROUP
      it.description = "Generate valid sql migration files for ${source.name} $name."
      it.properties = getProperties()
    }
  }

  private val Source.outputDir get() =
    if (sources.size > 1) File(generatedSourcesDirectory, name)
    else generatedSourcesDirectory
}

package com.squareup.sqldelight.gradle

import groovy.lang.Closure
import org.gradle.api.Project
import org.gradle.util.ConfigureUtil

open class SqlDelightExtension {
  internal val databases = mutableListOf<SqlDelightDatabase>()
  internal var configuringDatabase: SqlDelightDatabase? = null
  internal lateinit var project: Project

  var linkSqlite = true

  fun methodMissing(name: String, args: Any): Any {
    configuringDatabase?.methodMissing(name, args)?.let { return it }

    val closure = (args as? Array<*>)?.getOrNull(0) as? Closure<*>
      ?: throw IllegalStateException(
        """
        |Expected a closure for database names:
        |
        |sqldelight {
        |  $name {
        |    packageName = "com.sample"
        |    sourceSet = files("src/main/sqldelight")
        |  }
        |}
      """.trimMargin()
      )

    val database = SqlDelightDatabase(project, name = name)
    configuringDatabase = database
    ConfigureUtil.configure(closure, database)
    configuringDatabase = null

    check(!databases.any { it.name == database.name }) { "There is already a database defined for ${database.name}" }

    databases.add(database)
    return Unit
  }

  /**
   * Supports configuration in Kotlin script build files.
   *
   * sqldelight {
   *   database("MyDatabase") {
   *     packageName = "com.example"
   *     sourceSet = files("src/main/sqldelight")
   *   }
   * }
   */
  fun database(name: String, config: SqlDelightDatabase.() -> Unit) {
    val database = SqlDelightDatabase(project, name = name).apply(config)

    check(!databases.any { it.name == database.name }) { "There is already a database defined for ${database.name}" }

    databases.add(database)
  }
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.gradle

import com.squareup.sqldelight.VERSION
import com.squareup.sqldelight.core.MINIMUM_SUPPORTED_VERSION
import com.squareup.sqldelight.core.SqlDelightPropertiesFile
import com.squareup.sqldelight.gradle.android.packageName
import com.squareup.sqldelight.gradle.android.sqliteVersion
import com.squareup.sqldelight.gradle.kotlin.linkSqlite
import org.gradle.api.Plugin
import org.gradle.api.Project
import org.gradle.tooling.provider.model.ToolingModelBuilder
import org.gradle.tooling.provider.model.ToolingModelBuilderRegistry
import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
import org.jetbrains.kotlin.gradle.plugin.sources.DefaultKotlinSourceSet
import java.util.concurrent.atomic.AtomicBoolean
import javax.inject.Inject

abstract class SqlDelightPlugin : Plugin<Project> {
  private val android = AtomicBoolean(false)
  private val kotlin = AtomicBoolean(false)

  @get:Inject
  abstract val registry: ToolingModelBuilderRegistry

  private lateinit var extension: SqlDelightExtension

  override fun apply(project: Project) {
    extension = project.extensions.create("sqldelight", SqlDelightExtension::class.java)
    extension.project = project

    val androidPluginHandler = { _: Plugin<*> ->
      android.set(true)
      project.afterEvaluate {
        project.setupSqlDelightTasks(afterAndroid = true)
      }
    }
    project.plugins.withId("com.android.application", androidPluginHandler)
    project.plugins.withId("com.android.library", androidPluginHandler)
    project.plugins.withId("com.android.instantapp", androidPluginHandler)
    project.plugins.withId("com.android.feature", androidPluginHandler)
    project.plugins.withId("com.android.dynamic-feature", androidPluginHandler)

    val kotlinPluginHandler = { _: Plugin<*> -> kotlin.set(true) }
    project.plugins.withId("org.jetbrains.kotlin.multiplatform", kotlinPluginHandler)
    project.plugins.withId("org.jetbrains.kotlin.android", kotlinPluginHandler)
    project.plugins.withId("org.jetbrains.kotlin.jvm", kotlinPluginHandler)
    project.plugins.withId("kotlin2js", kotlinPluginHandler)

    project.afterEvaluate {
      project.setupSqlDelightTasks(afterAndroid = false)
    }
  }

  private fun Project.setupSqlDelightTasks(afterAndroid: Boolean) {
    if (android.get() && !afterAndroid) return

    check(kotlin.get()) {
      "SQL Delight Gradle plugin applied in " +
        "project '${project.path}' but no supported Kotlin plugin was found"
    }

    val isMultiplatform = project.plugins.hasPlugin("org.jetbrains.kotlin.multiplatform")

    // Add the runtime dependency.
    if (isMultiplatform) {
      val sourceSets =
        project.extensions.getByType(KotlinMultiplatformExtension::class.java).sourceSets
      val sourceSet = (sourceSets.getByName("commonMain") as DefaultKotlinSourceSet)
      project.configurations.getByName(sourceSet.apiConfigurationName).dependencies.add(
        project.dependencies.create("com.squareup.sqldelight:runtime:$VERSION")
      )
    } else {
      project.configurations.getByName("api").dependencies.add(
        project.dependencies.create("com.squareup.sqldelight:runtime-jvm:$VERSION")
      )
    }

    if (extension.linkSqlite) {
      project.linkSqlite()
    }

    extension.run {
      if (databases.isEmpty() && android.get() && !isMultiplatform) {
        // Default to a database for android named "Database" to keep things simple.
        databases.add(
          SqlDelightDatabase(
            project = project,
            name = "Database",
            packageName = project.packageName(),
            dialect = project.sqliteVersion()
          )
        )
      } else if (databases.isEmpty()) {
        logger.warn("SQLDelight Gradle plugin was applied but there are no databases set up.")
      }

      project.tasks.register("generateSqlDelightInterface") {
        it.group = GROUP
        it.description = "Aggregation task which runs every interface generation task for every given source"
      }

      project.tasks.register("verifySqlDelightMigration") {
        it.group = GROUP
        it.description = "Aggregation task which runs every migration task for every given source"
      }

      databases.forEach { database ->
        if (database.packageName == null && android.get() && !isMultiplatform) {
          database.packageName = project.packageName()
        }
        if (database.dialect == null && android.get() && !isMultiplatform) {
          database.dialect = project.sqliteVersion()
        }
        if (database.dialect == null) {
          database.dialect = "sqlite:3.18"
        }
        database.registerTasks()
      }

      val properties = SqlDelightPropertiesFileImpl(
        databases = databases.map { it.getProperties() },
        currentVersion = VERSION,
        minimumSupportedVersion = MINIMUM_SUPPORTED_VERSION,
      )
      registry.register(PropertiesModelBuilder(properties))
    }
  }

  class PropertiesModelBuilder(
    private val properties: SqlDelightPropertiesFile
  ) : ToolingModelBuilder {
    override fun canBuild(modelName: String): Boolean {
      return modelName == SqlDelightPropertiesFile::class.java.name
    }

    override fun buildAll(modelName: String, project: Project): Any {
      return properties
    }
  }

  internal companion object {
    const val GROUP = "sqldelight"
  }
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.gradle

import com.squareup.sqldelight.VERSION
import com.squareup.sqldelight.core.SqlDelightCompilationUnit
import com.squareup.sqldelight.core.SqlDelightDatabaseProperties
import com.squareup.sqldelight.core.SqlDelightEnvironment
import com.squareup.sqldelight.core.SqlDelightEnvironment.CompilationStatus.Failure
import com.squareup.sqldelight.core.SqlDelightException
import org.gradle.api.file.DirectoryProperty
import org.gradle.api.file.FileTree
import org.gradle.api.logging.LogLevel.ERROR
import org.gradle.api.logging.LogLevel.INFO
import org.gradle.api.logging.Logging
import org.gradle.api.provider.Property
import org.gradle.api.tasks.CacheableTask
import org.gradle.api.tasks.Input
import org.gradle.api.tasks.InputFiles
import org.gradle.api.tasks.Nested
import org.gradle.api.tasks.OutputDirectory
import org.gradle.api.tasks.PathSensitive
import org.gradle.api.tasks.PathSensitivity
import org.gradle.api.tasks.SkipWhenEmpty
import org.gradle.api.tasks.TaskAction
import org.gradle.workers.WorkAction
import org.gradle.workers.WorkParameters
import java.io.File

@Suppress("UnstableApiUsage") // Worker API
@CacheableTask
abstract class SqlDelightTask : SqlDelightWorkerTask() {
  @Suppress("unused") // Required to invalidate the task on version updates.
  @Input val pluginVersion = VERSION

  @get:OutputDirectory
  var outputDirectory: File? = null

  @Input val projectName: Property<String> = project.objects.property(String::class.java)

  @Nested lateinit var properties: SqlDelightDatabasePropertiesImpl
  @Nested lateinit var compilationUnit: SqlDelightCompilationUnitImpl

  @Input var verifyMigrations: Boolean = false

  @TaskAction
  fun generateSqlDelightFiles() {
    workQueue().submit(GenerateInterfaces::class.java) {
      it.outputDirectory.set(outputDirectory)
      it.projectName.set(projectName)
      it.properties.set(properties)
      it.verifyMigrations.set(verifyMigrations)
      it.compilationUnit.set(compilationUnit)
    }
  }

  @InputFiles
  @SkipWhenEmpty
  @PathSensitive(PathSensitivity.RELATIVE)
  override fun getSource(): FileTree {
    return super.getSource()
  }

  interface GenerateInterfacesWorkParameters : WorkParameters {
    val outputDirectory: DirectoryProperty
    val projectName: Property<String>
    val properties: Property<SqlDelightDatabaseProperties>
    val compilationUnit: Property<SqlDelightCompilationUnit>
    val verifyMigrations: Property<Boolean>
  }

  abstract class GenerateInterfaces : WorkAction<GenerateInterfacesWorkParameters> {
    private val logger = Logging.getLogger(SqlDelightTask::class.java)

    override fun execute() {
      parameters.outputDirectory.get().asFile.deleteRecursively()
      val environment = SqlDelightEnvironment(
        compilationUnit = parameters.compilationUnit.get(),
        properties = parameters.properties.get(),
        moduleName = parameters.projectName.get(),
        verifyMigrations = parameters.verifyMigrations.get(),
      )

      val generationStatus = environment.generateSqlDelightFiles { info ->
        logger.log(INFO, info)
      }

      when (generationStatus) {
        is Failure -> {
          logger.log(ERROR, "")
          generationStatus.errors.forEach { logger.log(ERROR, it) }
          throw SqlDelightException(
            "Generation failed; see the generator error output for details."
          )
        }
      }
    }
  }
}

package com.squareup.sqldelight.gradle

import org.gradle.api.services.BuildService
import org.gradle.api.services.BuildServiceParameters
import org.gradle.api.tasks.Input
import org.gradle.api.tasks.SourceTask
import org.gradle.workers.WorkQueue
import org.gradle.workers.WorkerExecutor
import javax.inject.Inject

/**
 * Common API for interacting with gradle workers
 * in tasks
 */
abstract class SqlDelightWorkerTask : SourceTask() {

  @get:Inject
  internal abstract val workerExecutor: WorkerExecutor

  /**
   * If true, [WorkerExecutor.classLoaderIsolation] is used. Otherwise,
   * [WorkerExecutor.noIsolation] is used for running the task's action
   */
  @get:Input
  internal var useClassLoaderIsolation: Boolean = true

  internal fun workQueue(): WorkQueue = if (useClassLoaderIsolation) {
    workerExecutor.classLoaderIsolation()
  } else {
    workerExecutor.noIsolation()
  }

  /**
   * Makes the task use [WorkerExecutor.noIsolation]
   * instead of [WorkerExecutor.classLoaderIsolation]
   */
  @Suppress("unused")
  fun disableClassLoaderIsolation() {
    useClassLoaderIsolation = false
    usesService(
      project.gradle.sharedServices.registerIfAbsent(
        SqlDelightWorkerTaskSerialService::class.toString(),
        SqlDelightWorkerTaskSerialService::class.java
      ) { it.maxParallelUsages.set(1) }
    )
  }
}

private abstract class SqlDelightWorkerTaskSerialService : BuildService<BuildServiceParameters.None>

package com.squareup.sqldelight.gradle

import com.squareup.sqldelight.VERSION
import com.squareup.sqldelight.core.SqlDelightCompilationUnit
import com.squareup.sqldelight.core.SqlDelightDatabaseProperties
import com.squareup.sqldelight.core.SqlDelightEnvironment
import com.squareup.sqldelight.core.lang.SqlDelightQueriesFile
import com.squareup.sqldelight.core.lang.util.forInitializationStatements
import com.squareup.sqldelight.core.lang.util.rawSqlText
import com.squareup.sqlite.migrations.CatalogDatabase
import com.squareup.sqlite.migrations.DatabaseFilesCollector
import com.squareup.sqlite.migrations.ObjectDifferDatabaseComparator
import org.gradle.api.file.DirectoryProperty
import org.gradle.api.file.FileTree
import org.gradle.api.logging.Logging
import org.gradle.api.provider.Property
import org.gradle.api.tasks.CacheableTask
import org.gradle.api.tasks.Input
import org.gradle.api.tasks.InputFiles
import org.gradle.api.tasks.Internal
import org.gradle.api.tasks.Nested
import org.gradle.api.tasks.PathSensitive
import org.gradle.api.tasks.PathSensitivity
import org.gradle.api.tasks.SkipWhenEmpty
import org.gradle.api.tasks.TaskAction
import org.gradle.workers.WorkAction
import org.gradle.workers.WorkParameters
import org.gradle.workers.WorkerExecutor
import java.io.File
import javax.inject.Inject

@Suppress("UnstableApiUsage") // Worker API
@CacheableTask
abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
  @Suppress("unused") // Required to invalidate the task on version updates.
  @Input val pluginVersion = VERSION

  @get:Inject
  abstract override val workerExecutor: WorkerExecutor

  @Input val projectName: Property<String> = project.objects.property(String::class.java)

  /** Directory where the database files are copied for the migration scripts to run against. */
  @Internal lateinit var workingDirectory: File

  @Nested lateinit var properties: SqlDelightDatabasePropertiesImpl
  @Nested lateinit var compilationUnit: SqlDelightCompilationUnitImpl

  @Input var verifyMigrations: Boolean = false

  @TaskAction
  fun verifyMigrations() {
    workQueue().submit(VerifyMigrationAction::class.java) {
      it.workingDirectory.set(workingDirectory)
      it.projectName.set(projectName)
      it.properties.set(properties)
      it.verifyMigrations.set(verifyMigrations)
      it.compilationUnit.set(compilationUnit)
    }
  }

  @InputFiles
  @SkipWhenEmpty
  @PathSensitive(PathSensitivity.RELATIVE)
  override fun getSource(): FileTree {
    return super.getSource()
  }

  interface VerifyMigrationWorkParameters : WorkParameters {
    val workingDirectory: DirectoryProperty
    val projectName: Property<String>
    val properties: Property<SqlDelightDatabaseProperties>
    val compilationUnit: Property<SqlDelightCompilationUnit>
    val verifyMigrations: Property<Boolean>
  }

  abstract class VerifyMigrationAction : WorkAction<VerifyMigrationWorkParameters> {
    private val logger = Logging.getLogger(VerifyMigrationTask::class.java)

    private val sourceFolders: List<File>
      get() = parameters.compilationUnit.get().sourceFolders.map { it.folder }

    private val environment by lazy {
      SqlDelightEnvironment(
        sourceFolders = sourceFolders.filter { it.exists() },
        dependencyFolders = emptyList(),
        moduleName = parameters.projectName.get(),
        properties = parameters.properties.get(),
        verifyMigrations = parameters.verifyMigrations.get(),
        compilationUnit = parameters.compilationUnit.get(),
      )
    }

    override fun execute() {
      parameters.workingDirectory.get().asFile.deleteRecursively()

      val catalog = createCurrentDb()
      DatabaseFilesCollector.forDatabaseFiles(sourceFolders) {
        checkMigration(it, catalog)
      }

      checkForGaps()
    }

    private fun createCurrentDb(): CatalogDatabase {
      val sourceFiles = ArrayList<SqlDelightQueriesFile>()
      environment.forSourceFiles { file ->
        if (file is SqlDelightQueriesFile) sourceFiles.add(file)
      }
      val initStatements = ArrayList<String>()
      sourceFiles.forInitializationStatements { sqlText ->
        initStatements.add(sqlText)
      }
      return CatalogDatabase.withInitStatements(initStatements)
    }

    private fun checkMigration(dbFile: File, currentDb: CatalogDatabase) {
      val actualCatalog = createActualDb(dbFile)
      val databaseComparator = ObjectDifferDatabaseComparator(
        circularReferenceExceptionLogger = {
          logger.debug(it)
        }
      )
      val diffReport = databaseComparator.compare(currentDb, actualCatalog).let { diff ->
        buildString(diff::printTo)
      }

      check(diffReport.isEmpty()) {
        "Error migrating from ${dbFile.name}, fresh database looks" +
          " different from migration database:\n$diffReport"
      }
    }

    private fun createActualDb(dbFile: File): CatalogDatabase {
      val version = dbFile.nameWithoutExtension.toInt()
      val copy = dbFile.copyTo(File(parameters.workingDirectory.get().asFile, dbFile.name))
      val initStatements = ArrayList<String>()
      environment.forMigrationFiles {
        check(it.name.any { it in '0'..'9' }) {
          "Migration files must have an integer value somewhere in their filename but ${it.name} does not."
        }
        if (version > it.version) return@forMigrationFiles
        it.sqlStmtList!!.stmtList.forEach {
          initStatements.add(it.rawSqlText())
        }
      }
      return CatalogDatabase.fromFile(copy.absolutePath, initStatements).also { copy.delete() }
    }

    private fun checkForGaps() {
      var lastMigrationVersion: Int? = null
      environment.forMigrationFiles {
        val actual = it.version
        val expected = lastMigrationVersion?.plus(1) ?: actual
        check(actual == expected) {
          "Gap in migrations detected. Expected migration $expected, got $actual."
        }

        lastMigrationVersion = actual
      }
    }
  }
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij

import com.intellij.openapi.vfs.LocalFileSystem
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiDirectory
import com.intellij.psi.PsiManager
import com.intellij.psi.util.PsiTreeUtil
import com.squareup.sqldelight.core.SqlDelightCompilationUnit
import com.squareup.sqldelight.core.SqlDelightDatabaseProperties
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.intellij.util.isAncestorOf
import java.io.File

class FileIndex(
  private val properties: SqlDelightDatabaseProperties,
  override val contentRoot: VirtualFile =
    LocalFileSystem.getInstance().findFileByPath(properties.rootDirectory.absolutePath)!!
) : SqlDelightFileIndex {
  override val isConfigured = true
  override val packageName = properties.packageName
  override val className = properties.className
  override val dependencies = properties.dependencies
  override val deriveSchemaFromMigrations = properties.deriveSchemaFromMigrations

  override fun packageName(file: SqlDelightFile): String {
    val original = if (file.parent == null) {
      file.originalFile as SqlDelightFile
    } else {
      file
    }
    val folder = sourceFolders(original, includeDependencies = false)
      .firstOrNull { PsiTreeUtil.isAncestor(it, original, false) } ?: return ""
    val folderPath = folder.virtualFile.path
    val filePath = original.virtualFile!!.path
    val startIndex = folderPath.length + 1
    return filePath.substring(startIndex, filePath.indexOf(original.name, startIndex) - 1).replace('/', '.')
  }

  override fun outputDirectory(sqFile: SqlDelightFile): List<String> {
    val file = sqFile.virtualFile ?: return emptyList()
    val compilationUnits = properties.compilationUnits.filter { compilationUnit ->
      compilationUnit.sourceFolders.any { it.folder.localVirtualFile()?.isAncestorOf(file) == true }
    }
    return compilationUnits.map { it.outputDirectoryPath }.distinct()
  }

  override fun outputDirectories(): List<String> {
    return properties.compilationUnits.map { it.outputDirectoryPath }.distinct()
  }

  private val SqlDelightCompilationUnit.outputDirectoryPath get() =
    outputDirectoryFile.relativeTo(File(contentRoot.path))
      .path.replace(File.separatorChar, '/').trimEnd('/')

  override fun sourceFolders(
    file: VirtualFile,
    includeDependencies: Boolean
  ): Collection<VirtualFile> {
    return properties.compilationUnits.map {
      it.sourceFolders
        .filter { includeDependencies || !it.dependency }
        .mapNotNull { sourceFolder -> sourceFolder.folder.localVirtualFile() }
    }.fold(emptySet()) { currentSources: Collection<VirtualFile>, sourceSet ->
      if (sourceSet.any { it.isAncestorOf(file) }) {
        // File is in this source set.
        if (currentSources.isEmpty()) {
          return@fold sourceSet
        } else {
          // File also in another source set! The files available sources is the intersection.
          return@fold currentSources.intersect(sourceSet)
        }
      }
      return@fold currentSources
    }
  }

  override fun sourceFolders(
    file: SqlDelightFile,
    includeDependencies: Boolean
  ): Collection<PsiDirectory> {
    return sourceFolders(file.virtualFile!!, includeDependencies)
      .map { PsiManager.getInstance(file.project).findDirectory(it)!! }
  }

  private fun File.localVirtualFile() = contentRoot.findFileByRelativePath(
    this.relativeTo(
      File(contentRoot.path)
    ).path.replace(File.separatorChar, '/').trimEnd('/')
  )
}

package com.squareup.sqldelight.intellij

import com.intellij.openapi.externalSystem.model.task.ExternalSystemTaskId
import com.intellij.openapi.externalSystem.model.task.ExternalSystemTaskNotificationListenerAdapter
import com.squareup.sqldelight.core.SqlDelightProjectService
import org.jetbrains.kotlin.idea.framework.GRADLE_SYSTEM_ID

class GradleSystemListener : ExternalSystemTaskNotificationListenerAdapter() {
  override fun onSuccess(id: ExternalSystemTaskId) {
    if (id.projectSystemId == GRADLE_SYSTEM_ID) {
      // Gradle sync just finished, reset the file index.
      id.findProject()?.let { project ->
        SqlDelightProjectService.getInstance(project).resetIndex()
      }
    }
  }
}

package com.squareup.sqldelight.intellij

import com.intellij.openapi.diagnostic.Logger
import timber.log.Timber

internal class LoggerTree(
  private val logger: Logger
) : Timber.Tree() {
  override fun v(message: String?, vararg args: Any?) {
    logger.info(String.format(message!!, args))
  }

  override fun v(t: Throwable?, message: String?, vararg args: Any?) {
    logger.info(String.format(message!!, args), t!!)
  }

  override fun v(t: Throwable?) {
    logger.info(t!!)
  }

  override fun d(message: String?, vararg args: Any?) {
    logger.debug(String.format(message!!, args))
  }

  override fun d(t: Throwable?, message: String?, vararg args: Any?) {
    logger.debug(String.format(message!!, args), t!!)
  }

  override fun d(t: Throwable?) {
    logger.debug(t!!)
  }

  override fun i(message: String?, vararg args: Any?) {
    logger.info(String.format(message!!, args))
  }

  override fun i(t: Throwable?, message: String?, vararg args: Any?) {
    logger.info(String.format(message!!, args), t)
  }

  override fun i(t: Throwable?) {
    logger.info(t!!)
  }

  override fun w(message: String?, vararg args: Any?) {
    logger.warn(String.format(message!!, args))
  }

  override fun w(t: Throwable?, message: String?, vararg args: Any?) {
    logger.warn(String.format(message!!, args), t!!)
  }

  override fun w(t: Throwable?) {
    logger.warn(t!!)
  }

  override fun e(message: String?, vararg args: Any?) {
    logger.error(String.format(message!!, args))
  }

  override fun e(t: Throwable?, message: String?, vararg args: Any?) {
    logger.error(String.format(message!!, args), t!!)
  }

  override fun e(t: Throwable?) {
    logger.error(t!!)
  }

  override fun wtf(message: String?, vararg args: Any?) {
    logger.error(String.format(message!!, args))
  }

  override fun wtf(t: Throwable?, message: String?, vararg args: Any?) {
    logger.error(String.format(message!!, args), t!!)
  }

  override fun wtf(t: Throwable?) {
    logger.error(t!!)
  }

  override fun log(priority: Int, message: String?, vararg args: Any?) {
    throw UnsupportedOperationException()
  }

  override fun log(priority: Int, t: Throwable?, message: String?, vararg args: Any?) {
    throw UnsupportedOperationException()
  }

  override fun log(priority: Int, t: Throwable?) {
    throw UnsupportedOperationException()
  }

  override fun log(priority: Int, tag: String?, message: String, t: Throwable?) {
    throw UnsupportedOperationException()
  }
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij

import com.alecstrong.sql.psi.core.DialectPreset
import com.intellij.openapi.Disposable
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.diagnostic.Logger
import com.intellij.openapi.module.Module
import com.intellij.openapi.project.Project
import com.intellij.openapi.roots.ProjectRootManager
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.openapi.vfs.VirtualFileManager
import com.intellij.openapi.vfs.newvfs.BulkFileListener
import com.intellij.openapi.vfs.newvfs.events.VFileCreateEvent
import com.intellij.openapi.vfs.newvfs.events.VFileDeleteEvent
import com.intellij.openapi.vfs.newvfs.events.VFileEvent
import com.intellij.openapi.vfs.newvfs.events.VFileMoveEvent
import com.intellij.psi.PsiDocumentManager
import com.intellij.psi.PsiManager
import com.intellij.psi.impl.PsiDocumentManagerImpl
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.SqlDelightProjectService
import com.squareup.sqldelight.core.compiler.SqlDelightCompiler
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.SqlDelightFileType
import com.squareup.sqldelight.intellij.gradle.FileIndexMap
import com.squareup.sqldelight.intellij.util.GeneratedVirtualFile
import timber.log.Timber
import java.io.PrintStream

class ProjectService(val project: Project) : SqlDelightProjectService, Disposable {
  private var fileIndexes = FileIndexMap()
  private val loggingTree = LoggerTree(Logger.getInstance("SQLDelight[${project.name}]"))

  init {
    Timber.plant(loggingTree)

    if (!ApplicationManager.getApplication().isUnitTestMode) {
      project.messageBus.connect()
        .subscribe(
          VirtualFileManager.VFS_CHANGES,
          object : BulkFileListener {
            override fun after(events: MutableList<out VFileEvent>) {
              events.filter { it.file?.fileType == SqlDelightFileType }.forEach { event ->
                if (event is VFileCreateEvent || event is VFileMoveEvent) {
                  PsiManager.getInstance(project).findViewProvider(event.file!!)
                    ?.contentsSynchronized()
                }
                if (event is VFileCreateEvent || event is VFileDeleteEvent || event is VFileMoveEvent) {
                  event.file?.let { generateDatabaseOnSync(it) }
                }
              }
            }
          }
        )
    }
  }

  override fun dispose() {
    Timber.uproot(loggingTree)
  }

  override fun resetIndex() {
    fileIndexes = FileIndexMap()
  }

  private fun generateDatabaseOnSync(vFile: VirtualFile) {
    val module = module(vFile) ?: return
    if (fileIndex(module) !is FileIndex) return

    val file = PsiManager.getInstance(project).findFile(vFile) as SqlDelightFile? ?: return

    val fileAppender = { filePath: String ->
      val vFile: VirtualFile by GeneratedVirtualFile(filePath, module)
      PrintStream(vFile.getOutputStream(this))
    }

    SqlDelightCompiler.writeDatabaseInterface(module, file, module.name, fileAppender)
  }

  override var dialectPreset: DialectPreset = DialectPreset.SQLITE_3_18
    set(value) {
      Timber.i("Setting dialect from $field to $value")
      val invalidate = field != value
      field = value
      if (invalidate) {
        val files = mutableListOf<VirtualFile>()
        ProjectRootManager.getInstance(project).fileIndex.iterateContent { vFile ->
          if (vFile.fileType != SqlDelightFileType) {
            return@iterateContent true
          }
          files += vFile
          return@iterateContent true
        }
        Timber.i("Invalidating ${files.size} files")
        ApplicationManager.getApplication().invokeLater {
          Timber.i("Reparsing ${files.size} files")
          (PsiDocumentManager.getInstance(project) as PsiDocumentManagerImpl)
            .reparseFiles(files, true)
        }
      }
    }

  override fun module(vFile: VirtualFile): Module? {
    return ProjectRootManager.getInstance(project).fileIndex.getModuleForFile(vFile)
  }

  override fun fileIndex(module: Module): SqlDelightFileIndex {
    return fileIndexes[module]
  }
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij

import com.intellij.codeInsight.completion.CompletionParameters
import com.intellij.codeInsight.completion.CompletionResultSet
import com.intellij.codeInsight.completion.CompletionUtil.findReferenceOrAlphanumericPrefix
import com.intellij.codeInsight.completion.InsertHandler
import com.intellij.codeInsight.completion.InsertionContext
import com.intellij.codeInsight.completion.JavaClassNameCompletionContributor
import com.intellij.codeInsight.completion.JavaPsiClassReferenceElement
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.psi.JavaTypeMixin
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.psi.SqlDelightImportStmt
import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
import kotlin.math.max

class SqlDelightClassCompletionContributor : JavaClassNameCompletionContributor() {
  private val insertHandler = AutoImportInsertionHandler()

  override fun fillCompletionVariants(
    parameters: CompletionParameters,
    resultSet: CompletionResultSet
  ) {
    if (parameters.position.getNonStrictParentOfType<JavaTypeMixin>() == null) return

    val result = resultSet.withPrefixMatcher(findReferenceOrAlphanumericPrefix(parameters))
    JavaClassNameCompletionContributor.addAllClasses(
      parameters, parameters.invocationCount <= 1,
      result.prefixMatcher
    ) { lookupElement ->
      if (lookupElement is JavaPsiClassReferenceElement) {
        lookupElement.setInsertHandler(insertHandler)
      }
      resultSet.addElement(lookupElement)
    }
  }
}

private class AutoImportInsertionHandler : InsertHandler<JavaPsiClassReferenceElement> {
  override fun handleInsert(context: InsertionContext, item: JavaPsiClassReferenceElement) {
    val qname = item.qualifiedName
    val imports = (context.file as SqlDelightFile).sqlStmtList
      ?.findChildrenOfType<SqlDelightImportStmt>().orEmpty()
    val ref = imports.map { it.javaType }.find { it.textMatches(qname) }
    val refEnd = context.trackOffset(context.tailOffset, false)

    if (ref == null) {
      if (imports.isEmpty()) {
        context.document.insertString(0, "import $qname;\n\n")
      } else {
        context.insertAndOrganizeImports(imports, "import $qname;")
      }
    }

    context.tailOffset = context.getOffset(refEnd)

    context.commitDocument()
  }

  private fun InsertionContext.insertAndOrganizeImports(
    imports: Collection<SqlDelightImportStmt>,
    newImport: String
  ) {
    val newImports = arrayListOf(newImport)
    var endOffset = 0
    for (importElement in imports) {
      newImports.add(importElement.text)
      endOffset = max(endOffset, importElement.textOffset + importElement.textLength)
    }
    document.replaceString(0, endOffset, newImports.sorted().joinToString("\n"))
  }
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij

import com.intellij.codeInspection.ProblemHighlightType
import com.intellij.lang.annotation.AnnotationHolder
import com.intellij.lang.annotation.Annotator
import com.intellij.lang.annotation.HighlightSeverity
import com.intellij.openapi.module.ModuleUtil
import com.intellij.psi.PsiClass
import com.intellij.psi.PsiElement
import com.intellij.psi.search.GlobalSearchScope
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.psi.ImportStmtMixin
import com.squareup.sqldelight.core.lang.psi.JavaTypeMixin
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.intellij.intentions.AddImportIntention
import com.squareup.sqldelight.intellij.util.PsiClassSearchHelper

class SqlDelightClassNameElementAnnotator : Annotator {
  override fun annotate(element: PsiElement, holder: AnnotationHolder) {
    if (element !is JavaTypeMixin || element.reference.resolve() != null) {
      return
    }

    val project = element.project
    val sqlDelightFile = element.containingFile as SqlDelightFile
    val module = ModuleUtil.findModuleForFile(sqlDelightFile) ?: return
    val scope = GlobalSearchScope.moduleWithDependenciesAndLibrariesScope(module, false)
    val outerClassElement = element.firstChild
    val outerClassName = outerClassElement.text
    val classes = PsiClassSearchHelper.getClassesByShortName(outerClassName, project, scope)

    val hasImport = sqlDelightFile.hasImport(outerClassName)
    val psiElement = if (hasImport) {
      missingNestedClass(classes, element)
    } else {
      outerClassElement
    }

    val needsQuickFix = classes.isNotEmpty() && !hasImport

    holder.newAnnotation(HighlightSeverity.ERROR, "Unresolved reference: ${psiElement.text}")
      .range(psiElement)
      .highlightType(ProblemHighlightType.LIKE_UNKNOWN_SYMBOL)
      .apply {
        if (needsQuickFix) {
          withFix(AddImportIntention(outerClassName))
        }
      }
      .create()
  }

  private fun SqlDelightFile.hasImport(outerClassName: String): Boolean {
    return sqlStmtList?.findChildrenOfType<ImportStmtMixin>()
      .orEmpty()
      .any { it.javaType.text.endsWith(outerClassName) }
  }

  private fun missingNestedClass(classes: List<PsiClass>, javaTypeMixin: JavaTypeMixin): PsiElement {
    val elementText = javaTypeMixin.text
    val className = classes.map { clazz -> findMissingNestedClassName(clazz, elementText) }
      .maxByOrNull { it.length }
      ?.substringBefore(".") ?: return javaTypeMixin.firstChild
    return javaTypeMixin.findChildrenOfType<PsiElement>().first { it.textMatches(className) }
  }

  private fun findMissingNestedClassName(psiClass: PsiClass, className: String): String {
    val nestedClassName = className.substringBefore(".")
    if (psiClass.name != nestedClassName) {
      return nestedClassName
    }
    val nextName = className.removePrefix(nestedClassName).removePrefix(".")
    val lookupString = nextName.substringBefore(".")
    val nextClass = psiClass.innerClasses.firstOrNull { clazz ->
      clazz.textMatches(lookupString)
    }
    return nextClass?.let { findMissingNestedClassName(it, nextName) } ?: nextName
  }
}

package com.squareup.sqldelight.intellij

import com.intellij.openapi.editor.colors.TextAttributesKey
import com.intellij.openapi.fileTypes.SyntaxHighlighter
import com.intellij.openapi.options.colors.AttributesDescriptor
import com.intellij.openapi.options.colors.ColorDescriptor
import com.intellij.openapi.options.colors.ColorSettingsPage
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlightVisitor.Companion.SQL_ALIAS
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlightVisitor.Companion.SQL_COLUMN_NAME
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlightVisitor.Companion.SQL_FUNCTION_NAME
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlightVisitor.Companion.SQL_STMT_IDENTIFIER_NAME
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlightVisitor.Companion.SQL_TABLE_NAME
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlightVisitor.Companion.SQL_TYPE_NAME
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlighter
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlighter.Companion.SQLITE_COMMA
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlighter.Companion.SQLITE_DOC
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlighter.Companion.SQLITE_DOT
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlighter.Companion.SQLITE_KEYWORD
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlighter.Companion.SQLITE_LINE_COMMENT
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlighter.Companion.SQLITE_NUMBER
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlighter.Companion.SQLITE_OPERATOR
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlighter.Companion.SQLITE_PAREN
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlighter.Companion.SQLITE_SEMICOLON
import com.squareup.sqldelight.intellij.lang.SqlDelightHighlighter.Companion.SQLITE_STRING
import javax.swing.Icon

class SqlDelightColorSettingsPage : ColorSettingsPage {

  private companion object {
    val attributes = mapOf(
      "Column" to SQL_COLUMN_NAME,
      "Comma" to SQLITE_COMMA,
      "Comment" to SQLITE_LINE_COMMENT,
      "Doc" to SQLITE_DOC,
      "Dot" to SQLITE_DOT,
      "Function" to SQL_FUNCTION_NAME,
      "Keyword" to SQLITE_KEYWORD,
      "Number token" to SQLITE_NUMBER,
      "Operator" to SQLITE_OPERATOR,
      "Parentheses" to SQLITE_PAREN,
      "Semicolon" to SQLITE_SEMICOLON,
      "Statement identifier" to SQL_STMT_IDENTIFIER_NAME,
      "String token" to SQLITE_STRING,
      "Table or column alias" to SQL_ALIAS,
      "Table" to SQL_TABLE_NAME,
      "Type" to SQL_TYPE_NAME,
    ).map { AttributesDescriptor(it.key, it.value) }.toTypedArray()

    val additionalTags = mapOf(
      "alias" to SQL_ALIAS,
      "column" to SQL_COLUMN_NAME,
      "function" to SQL_FUNCTION_NAME,
      "stmt_id" to SQL_STMT_IDENTIFIER_NAME,
      "table" to SQL_TABLE_NAME,
      "type" to SQL_TYPE_NAME,
    )
  }

  override fun getAttributeDescriptors(): Array<AttributesDescriptor> {
    return attributes
  }

  override fun getColorDescriptors(): Array<ColorDescriptor> {
    return ColorDescriptor.EMPTY_ARRAY
  }

  override fun getDisplayName(): String = "SQLDelight"

  override fun getIcon(): Icon? {
    return null
  }

  override fun getHighlighter(): SyntaxHighlighter = SqlDelightHighlighter()

  override fun getAdditionalHighlightingTagToDescriptorMap(): Map<String, TextAttributesKey> {
    return additionalTags
  }

  override fun getDemoText(): String {
    return """
      |/**
      |* Doc comment
      |*/
      |CREATE TABLE crm.
product
 (
      |  <column>id</column> <type>INTEGER</type> PRIMARY KEY,
      |  <column>title</column> <type>TEXT</type>
      |);
      |
      |-- Line comment
      |INSERT INTO
product
 VALUES (1, 'Product1');
      |
      |<stmt_id>count</stmt_id>:
      |SELECT <function>COUNT</function>(*) FROM crm.
product
;
      |
      |<stmt_id>selectById</stmt_id>:
      |SELECT <column>id</column> AS <alias>ProductID</alias>, <column>title</column> AS <alias>ProductName</alias>
      |  FROM crm.
product
 WHERE <column>id</column> = ?;
      |""".trimMargin()
  }
}

package com.squareup.sqldelight.intellij

import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.alecstrong.sql.psi.core.psi.SqlStmt
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.intellij.codeInsight.daemon.impl.CollectHighlightsUtil
import com.intellij.codeInsight.editorActions.CopyPastePostProcessor
import com.intellij.codeInsight.editorActions.ReferenceData
import com.intellij.codeInsight.editorActions.ReferenceTransferableData
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.editor.RangeMarker
import com.intellij.openapi.editor.ReadOnlyModificationException
import com.intellij.openapi.project.DumbService
import com.intellij.openapi.project.Project
import com.intellij.openapi.util.Ref
import com.intellij.psi.PsiClass
import com.intellij.psi.PsiDocumentManager
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import com.intellij.psi.util.PsiTreeUtil
import com.intellij.psi.util.parentOfType
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.psi.JavaTypeMixin
import com.squareup.sqldelight.core.lang.util.findChildOfType
import com.squareup.sqldelight.core.psi.SqlDelightImportStmtList
import java.awt.datatransfer.Transferable
import java.awt.datatransfer.UnsupportedFlavorException
import java.io.IOException

class SqlDelightCopyPasteProcessor : CopyPastePostProcessor<ReferenceTransferableData>() {

  override fun collectTransferableData(
    file: PsiFile,
    editor: Editor,
    startOffsets: IntArray,
    endOffsets: IntArray
  ): List<ReferenceTransferableData> {
    if (file !is SqlDelightFile) {
      return emptyList()
    }

    val result = mutableListOf<ReferenceData>()
    var refOffset = 0

    for (j in startOffsets.indices) {
      refOffset += startOffsets[j]
      val elementsInRange = CollectHighlightsUtil.getElementsInRange(file, startOffsets[j], endOffsets[j])
      for (element in elementsInRange) {
        val data = referenceData(element, refOffset)
        if (data != null) {
          result += data
        }
      }
      refOffset -= endOffsets[j] + 1
    }
    return listOf(ReferenceTransferableData(result.toTypedArray()))
  }

  private fun referenceData(
    element: PsiElement,
    refOffset: Int
  ): ReferenceData? {
    if (element !is JavaTypeMixin) {
      return null
    }

    val resolvedClass = element.reference.resolve() as? PsiClass ?: return null
    val qualifiedName = resolvedClass.qualifiedName ?: return null
    val range = element.textRange
    return ReferenceData(
      range.startOffset - refOffset,
      range.endOffset - refOffset,
      qualifiedName,
      null
    )
  }

  override fun extractTransferableData(content: Transferable): List<ReferenceTransferableData> {
    try {
      val dataFlavor = ReferenceData.getDataFlavor() ?: return emptyList()
      val referenceData = content.getTransferData(dataFlavor) as ReferenceTransferableData?
      return listOfNotNull(referenceData)
    } catch (ignored: UnsupportedFlavorException) {
    } catch (ignored: IOException) {
    }
    return super.extractTransferableData(content)
  }

  override fun processTransferableData(
    project: Project,
    editor: Editor,
    bounds: RangeMarker,
    caretOffset: Int,
    indented: Ref<Boolean>,
    values: List<ReferenceTransferableData>
  ) {
    if (DumbService.getInstance(project).isDumb) {
      return
    }

    val references = values.first().data
    if (references.isNullOrEmpty()) {
      return
    }

    val document = editor.document
    val documentManager = PsiDocumentManager.getInstance(project)
    val file = documentManager.getPsiFile(document)
    if (file !is SqlDelightFile) {
      return
    }
    val elementAtCaret = file.findElementAt(caretOffset)
    val insideCreateTableStmt = elementAtCaret?.parentOfType<SqlCreateTableStmt>() != null
    val hasCreateTableSibling = elementAtCaret?.parentOfType<SqlStmt>() != null &&
      findCreateTableSiblingCatching(elementAtCaret) != null
    if (!insideCreateTableStmt && !hasCreateTableSibling) {
      return
    }
    documentManager.commitAllDocuments()

    WriteCommandAction.writeCommandAction(project).run<ReadOnlyModificationException> {
      val qClassNames = references.map(ReferenceData::qClassName)
      val importStmtList = file.findChildOfType<SqlDelightImportStmtList>()?.importStmtList.orEmpty()
      val oldImports = importStmtList.map { it.javaType.text }
      val newImports = (qClassNames + oldImports).distinct()
        .sorted()
        .joinToString("\n") { "import $it;" }
      if (oldImports.isEmpty()) {
        document.insertString(0, "$newImports\n\n")
      } else {
        val endOffset = importStmtList.maxOfOrNull { it.textOffset + it.textLength } ?: 0
        document.replaceString(0, endOffset, newImports)
      }
    }
  }

  private fun findCreateTableSiblingCatching(elementAtCaret: PsiElement?): PsiElement? {
    return elementAtCaret?.runCatching {
      PsiTreeUtil.findSiblingBackward(this, SqlTypes.CREATE_TABLE_STMT, false, null)
    }?.getOrNull()
  }
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij

import com.bugsnag.Bugsnag
import com.bugsnag.Severity
import com.intellij.diagnostic.AbstractMessage
import com.intellij.ide.plugins.PluginManagerCore
import com.intellij.openapi.application.ApplicationInfo
import com.intellij.openapi.diagnostic.ErrorReportSubmitter
import com.intellij.openapi.diagnostic.IdeaLoggingEvent
import com.intellij.openapi.diagnostic.SubmittedReportInfo
import com.intellij.util.Consumer
import com.squareup.sqldelight.VERSION
import java.awt.Component

class SqlDelightErrorHandler : ErrorReportSubmitter() {
  val bugsnag = Bugsnag(BUGSNAG_KEY, false)

  init {
    bugsnag.setAutoCaptureSessions(false)
    bugsnag.startSession()
    bugsnag.setAppVersion(VERSION)
    bugsnag.setProjectPackages("com.squareup.sqldelight")
    bugsnag.addCallback {
      it.addToTab("Device", "osVersion", System.getProperty("os.version"))
      it.addToTab("Device", "JRE", System.getProperty("java.version"))
      it.addToTab("Device", "IDE Version", ApplicationInfo.getInstance().fullVersion)
      it.addToTab("Device", "IDE Build #", ApplicationInfo.getInstance().build)
      PluginManagerCore.getPlugins().forEach { plugin ->
        it.addToTab("Plugins", plugin.name, "${plugin.pluginId} : ${plugin.version}")
      }
    }
  }

  override fun getReportActionText() = "Send to Square"
  override fun submit(
    events: Array<out IdeaLoggingEvent>,
    additionalInfo: String?,
    parentComponent: Component,
    consumer: Consumer<SubmittedReportInfo>
  ): Boolean {
    for (event in events) {
      if (BUGSNAG_KEY.isNotBlank()) {
        val throwable = (event.data as? AbstractMessage)?.throwable ?: event.throwable
        bugsnag.notify(throwable, Severity.ERROR) {
          it.addToTab("Data", "message", event.message)
          it.addToTab("Data", "additional info", additionalInfo)
          it.addToTab("Data", "stacktrace", event.throwableText)
        }
      }
    }
    return true
  }
}

package com.squareup.sqldelight.intellij

import com.alecstrong.sql.psi.core.DialectPreset
import com.alecstrong.sql.psi.core.DialectPreset.SQLITE_3_18
import com.alecstrong.sql.psi.core.DialectPreset.SQLITE_3_24
import com.alecstrong.sql.psi.core.DialectPreset.SQLITE_3_25
import com.intellij.icons.AllIcons
import com.intellij.ide.FileIconProvider
import com.intellij.lang.Language
import com.intellij.lang.LanguageUtil
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.VirtualFile
import com.squareup.sqldelight.core.SqlDelightProjectService
import com.squareup.sqldelight.core.lang.MigrationLanguage
import com.squareup.sqldelight.core.lang.SqlDelightLanguage
import javax.swing.Icon

class SqlDelightFileIconProvider : FileIconProvider {

  override fun getIcon(file: VirtualFile, flags: Int, project: Project?): Icon? =
    getIcon(LanguageUtil.getFileLanguage(file), project)

  companion object {
    private val supportedLanguages = setOf(SqlDelightLanguage, MigrationLanguage)

    fun getIcon(language: Language?, project: Project?): Icon? {
      return if (project != null && language in supportedLanguages) {
        when (SqlDelightProjectService.getInstance(project).dialectPreset) {
          SQLITE_3_18, SQLITE_3_24, SQLITE_3_25 -> AllIcons.Providers.Sqlite
          DialectPreset.MYSQL -> AllIcons.Providers.Mysql
          DialectPreset.POSTGRESQL -> AllIcons.Providers.Postgresql
          DialectPreset.HSQL -> AllIcons.Providers.Hsqldb
        }
      } else {
        null
      }
    }
  }
}

package com.squareup.sqldelight.intellij

import com.intellij.openapi.vfs.VirtualFile
import com.squareup.sqldelight.core.SqlDelightDatabaseName
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.lang.SqlDelightFile

internal class SqlDelightFileIndexImpl : SqlDelightFileIndex {
  override val isConfigured
    get() = false
  override val packageName = ""
  override val className
    get() = throw UnsupportedOperationException()
  override val contentRoot
    get() = throw UnsupportedOperationException()
  override val dependencies: List<SqlDelightDatabaseName>
    get() = throw UnsupportedOperationException()
  override val deriveSchemaFromMigrations = false

  override fun outputDirectory(file: SqlDelightFile) = throw UnsupportedOperationException()
  override fun outputDirectories() = throw UnsupportedOperationException()

  override fun packageName(file: SqlDelightFile) = ""
  override fun sourceFolders(
    file: VirtualFile,
    includeDependencies: Boolean
  ) = listOf(file.parent)
  override fun sourceFolders(
    file: SqlDelightFile,
    includeDependencies: Boolean
  ) = listOfNotNull(file.parent)
}

package com.squareup.sqldelight.intellij

import com.alecstrong.sql.psi.core.psi.SqlColumnName
import com.intellij.find.findUsages.FindUsagesHandler
import com.intellij.find.findUsages.FindUsagesHandlerFactory
import com.intellij.find.findUsages.FindUsagesOptions
import com.intellij.openapi.module.ModuleUtil
import com.intellij.openapi.roots.ProjectRootManager
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiManager
import com.intellij.psi.util.PsiTreeUtil
import com.intellij.usageView.UsageInfo
import com.intellij.util.Processor
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.psi.StmtIdentifierMixin
import com.squareup.sqldelight.core.lang.queriesName
import com.squareup.sqldelight.core.psi.SqlDelightStmtIdentifier
import org.jetbrains.kotlin.idea.findUsages.KotlinFindUsagesHandlerFactory
import org.jetbrains.kotlin.idea.findUsages.KotlinReferenceUsageInfo
import org.jetbrains.kotlin.psi.KtClass
import org.jetbrains.kotlin.psi.KtFile
import org.jetbrains.kotlin.psi.KtNamedDeclaration
import org.jetbrains.kotlin.psi.KtNamedFunction

class SqlDelightFindUsagesHandlerFactory : FindUsagesHandlerFactory() {
  override fun canFindUsages(element: PsiElement): Boolean {
    val module = ModuleUtil.findModuleForPsiElement(element) ?: return false
    val isConfigured = SqlDelightFileIndex.getInstance(module).isConfigured
    return (element is SqlDelightStmtIdentifier || element is SqlColumnName) && isConfigured
  }

  override fun createFindUsagesHandler(
    element: PsiElement,
    forHighlightUsages: Boolean
  ): FindUsagesHandler = SqlDelightIdentifierHandler(element)
}

private class SqlDelightIdentifierHandler(
  private val element: PsiElement
) : FindUsagesHandler(element) {
  private val factory = KotlinFindUsagesHandlerFactory(element.project)
  private val kotlinHandlers = when (element) {
    is StmtIdentifierMixin -> element.generatedMethods()
    is SqlColumnName -> element.generatedProperties()
    else -> throw AssertionError()
  }
    .map { factory.createFindUsagesHandler(it, false) }

  private val kotlinFindUsagesOptions: FindUsagesOptions = when (element) {
    is StmtIdentifierMixin -> factory.findFunctionOptions
    is SqlColumnName -> factory.findPropertyOptions
    else -> throw AssertionError()
  }

  override fun getPrimaryElements() = arrayOf(element)

  override fun processElementUsages(
    sourceElement: PsiElement,
    processor: Processor<in UsageInfo>,
    options: FindUsagesOptions
  ): Boolean {
    val generatedFiles = element.generatedVirtualFiles()
    val ignoringFileProcessor = Processor<UsageInfo> { t ->
      if (t is KotlinReferenceUsageInfo && t.virtualFile in generatedFiles) {
        return@Processor true
      }
      processor.process(t)
    }
    super.processElementUsages(element, ignoringFileProcessor, options)
    return kotlinHandlers.all {
      it.processElementUsages(
        it.primaryElements.single(),
        ignoringFileProcessor,
        kotlinFindUsagesOptions
      )
    }
  }
}

internal fun PsiElement.generatedQueryFiles(): List<VirtualFile> {
  val queriesName = (containingFile as SqlDelightFile).let { file ->
    file.virtualFile?.queriesName
  }
  return generatedVirtualFiles().filter { it.nameWithoutExtension == queriesName }
}

internal fun StmtIdentifierMixin.generatedMethods(): Collection<KtNamedDeclaration> {
  val files = generatedQueryFiles().map {
    PsiManager.getInstance(project).findFile(it) as KtFile
  }
  return files.flatMap { file ->
    PsiTreeUtil.findChildrenOfType(file, KtNamedFunction::class.java).filter {
      it.name == identifier()?.text
    }
  }
}

internal fun PsiElement.generatedVirtualFiles(): List<VirtualFile> {
  val module = ModuleUtil.findModuleForPsiElement(this) ?: return emptyList()
  val fileIndex = SqlDelightFileIndex.getInstance(module)
  val generatedDirectories = (containingFile as SqlDelightFile).generatedDirectories.orEmpty()
    .map { fileIndex.contentRoot.findFileByRelativePath(it)!! }
  val rootManager = ProjectRootManager.getInstance(project)
  return generatedDirectories.flatMap { dir ->
    val mut = mutableListOf<VirtualFile>()
    rootManager.fileIndex.iterateContentUnderDirectory(dir) { file ->
      mut += file
      true
    }
    mut.toList()
  }
}

internal fun PsiElement.generatedKtFiles(): List<KtFile> {
  val psiManager = PsiManager.getInstance(project)
  return generatedVirtualFiles().mapNotNull { psiManager.findFile(it) }
    .filterIsInstance<KtFile>()
}

internal fun SqlColumnName.generatedProperties(): Collection<KtNamedDeclaration> {
  return generatedKtFiles().asSequence()
    .flatMap { it.declarations.filterIsInstance<KtClass>() }
    .filter(KtClass::isData)
    .flatMap(KtClass::getPrimaryConstructorParameters)
    .filter { it.name == name }
    .toList()
}

package com.squareup.sqldelight.intellij

import com.alecstrong.sql.psi.core.psi.SqlColumnAlias
import com.alecstrong.sql.psi.core.psi.SqlColumnName
import com.alecstrong.sql.psi.core.psi.SqlCteTableName
import com.alecstrong.sql.psi.core.psi.SqlTableAlias
import com.alecstrong.sql.psi.core.psi.SqlTableName
import com.alecstrong.sql.psi.core.psi.SqlViewName
import com.intellij.lang.findUsages.FindUsagesProvider
import com.intellij.psi.PsiElement
import com.squareup.sqldelight.core.lang.psi.StmtIdentifierMixin

class SqlDelightFindUsagesProvider : FindUsagesProvider {
  override fun getWordsScanner() = null
  override fun getNodeText(element: PsiElement, useFullName: Boolean) = element.text
  override fun getDescriptiveName(element: PsiElement) = element.text
  override fun getHelpId(psiElement: PsiElement) = null

  override fun getType(element: PsiElement): String {
    return when (element) {
      is StmtIdentifierMixin -> "query"
      is SqlTableName -> "table"
      is SqlColumnName -> "column"
      is SqlTableAlias -> "table alias"
      is SqlColumnAlias -> "column alias"
      is SqlCteTableName -> "common table"
      is SqlViewName -> "view"
      else -> throw IllegalArgumentException("Unexpected type $element")
    }
  }

  override fun canFindUsagesFor(element: PsiElement): Boolean {
    return when (element) {
      is StmtIdentifierMixin, is SqlTableName, is SqlColumnName, is SqlTableAlias,
      is SqlColumnAlias, is SqlViewName, is SqlCteTableName -> true
      else -> false
    }
  }
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij

import com.alecstrong.sql.psi.core.psi.QueryElement
import com.alecstrong.sql.psi.core.psi.SqlStmt
import com.intellij.codeInsight.navigation.actions.GotoDeclarationHandler
import com.intellij.openapi.actionSystem.DataContext
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiManager
import com.intellij.psi.PsiMethod
import com.intellij.psi.PsiNamedElement
import com.intellij.psi.PsiReference
import com.intellij.psi.search.FileTypeIndex
import com.intellij.psi.search.GlobalSearchScope
import com.intellij.psi.search.searches.ReferencesSearch
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.lang.QUERIES_SUFFIX_NAME
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.SqlDelightFileType
import com.squareup.sqldelight.core.lang.queriesName
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.psi.SqlDelightStmtIdentifier
import com.squareup.sqldelight.intellij.util.isAncestorOf
import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
import org.jetbrains.kotlin.idea.util.projectStructure.getModule
import org.jetbrains.kotlin.psi.KtNameReferenceExpression
import org.jetbrains.kotlin.psi.KtNamedFunction
import org.jetbrains.kotlin.psi.KtParameter
import org.jetbrains.kotlin.psi.psiUtil.containingClass
import org.jetbrains.kotlin.psi.psiUtil.getChildOfType
import org.jetbrains.kotlin.psi.psiUtil.getNextSiblingIgnoringWhitespaceAndComments
import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance

class SqlDelightGotoDeclarationHandler : GotoDeclarationHandler {
  override fun getGotoDeclarationTargets(
    sourceElement: PsiElement?,
    offset: Int,
    editor: Editor
  ): Array<PsiElement> {
    if (sourceElement == null) {
      return emptyArray()
    }

    val targetData = targetData(sourceElement) ?: return emptyArray()
    val function = targetData.function
    val elementFile = targetData.containingFile

    val module = elementFile.getModule(function.project) ?: return emptyArray()

    // Only handle files under the generated sqlite directory.
    val fileIndex = SqlDelightFileIndex.getInstance(module)
    if (!fileIndex.isConfigured) return emptyArray()

    fileIndex.outputDirectories().forEach { dir ->
      val outputDirectory = fileIndex.contentRoot.findFileByRelativePath(dir) ?: return emptyArray()
      if (!outputDirectory.isAncestorOf(elementFile)) return emptyArray()

      var result = emptyArray<PsiElement>()
      val scope = GlobalSearchScope.moduleWithDependentsScope(module)
      val manager = PsiManager.getInstance(function.project)
      FileTypeIndex.getFiles(SqlDelightFileType, scope)
        .asSequence()
        .filter { it.queriesName == elementFile.nameWithoutExtension }
        .map { manager.findFile(it) as SqlDelightFile }
        .filter { it.sqlStmtList != null }
        .forEach inner@{ sqlDelightFile ->
          val identifier = sqlDelightFile.sqlStmtList!!
            .findChildrenOfType<SqlDelightStmtIdentifier>()
            .mapNotNull { it.identifier() }
            .first { it.textMatches(function.name!!) }
          result = if (targetData.parameter != null) {
            val sqlStmt = identifier.getParentOfType<SqlDelightStmtIdentifier>(true)
              ?.getNextSiblingIgnoringWhitespaceAndComments() as? SqlStmt
              ?: return@inner

            sqlStmt.getChildOfType<QueryElement>()?.queryExposed()
              .orEmpty()
              .flatMap { it.columns }
              .filter { it.element.textMatches(targetData.parameter.name!!) }
              .mapNotNull { it.element.reference?.resolve() }
              .toTypedArray()
          } else {
            arrayOf(identifier)
          }
        }
      return result
    }

    return emptyArray()
  }

  private fun targetData(psiElement: PsiElement): TargetData? {
    val reference = when (psiElement.parent) {
      is PsiReference -> psiElement.parent as PsiReference
      is KtNameReferenceExpression -> {
        psiElement.parent.references.firstIsInstance<KtSimpleNameReference>()
      }
      else -> return null
    }
    return when (val resolved = reference.resolve()) {
      is KtParameter -> {
        val ktClass = resolved.containingClass() ?: return null
        val element = ReferencesSearch.search(ktClass)
          .mapping { it.element }
          .firstOrNull {
            it.containingFile.virtualFile.nameWithoutExtension.endsWith(QUERIES_SUFFIX_NAME)
          } ?: return null
        TargetData(
          parameter = resolved,
          function = element.getParentOfType(true)!!,
          containingFile = element.containingFile.virtualFile
        )
      }
      is KtNamedFunction, is PsiMethod -> TargetData(
        parameter = null,
        function = resolved as PsiNamedElement,
        containingFile = resolved.containingFile.virtualFile
      )
      else -> null
    }
  }

  override fun getActionText(context: DataContext): String? = null

  data class TargetData(
    val parameter: PsiNamedElement?,
    val function: PsiNamedElement,
    val containingFile: VirtualFile
  )
}

package com.squareup.sqldelight.intellij

import com.alecstrong.sql.psi.core.psi.SqlInsertStmt
import com.alecstrong.sql.psi.core.psi.SqlInsertStmtValues
import com.alecstrong.sql.psi.core.psi.SqlValuesExpression
import com.intellij.codeInsight.hints.InlayInfo
import com.intellij.codeInsight.hints.InlayParameterHintsProvider
import com.intellij.psi.PsiElement
import com.intellij.psi.util.parentOfType
import com.intellij.refactoring.suggested.startOffset

class SqlDelightInlayParameterHintsProvider : InlayParameterHintsProvider {

  override fun getParameterHints(element: PsiElement): List<InlayInfo> {
    if (element !is SqlValuesExpression || element.parent !is SqlInsertStmtValues) {
      return emptyList()
    }
    return element.parentOfType<SqlInsertStmt>()?.columnNameList.orEmpty()
      .zip(element.exprList) { col, expr ->
        InlayInfo(col.name, expr.startOffset)
      }
  }

  override fun getDefaultBlackList(): Set<String> {
    return emptySet()
  }
}

package com.squareup.sqldelight.intellij

import com.intellij.codeInsight.template.TemplateContextType
import com.intellij.psi.PsiFile
import com.squareup.sqldelight.core.lang.MigrationFileType
import com.squareup.sqldelight.core.lang.SqlDelightFileType

class SqlDelightLiveTemplateContextType : TemplateContextType("SQLDELIGHT", "SqlDelight") {

  private val supportedFileTypes = setOf(SqlDelightFileType, MigrationFileType)

  override fun isInContext(file: PsiFile, offset: Int): Boolean = file.fileType in supportedFileTypes
}

package com.squareup.sqldelight.intellij

import com.alecstrong.sql.psi.core.psi.SqlBindExpr
import com.alecstrong.sql.psi.core.psi.SqlColumnDef
import com.alecstrong.sql.psi.core.psi.SqlInsertStmt
import com.alecstrong.sql.psi.core.psi.SqlInsertStmtValues
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.alecstrong.sql.psi.core.psi.SqlValuesExpression
import com.intellij.codeInsight.lookup.LookupElement
import com.intellij.lang.parameterInfo.CreateParameterInfoContext
import com.intellij.lang.parameterInfo.ParameterInfoContext
import com.intellij.lang.parameterInfo.ParameterInfoHandlerWithTabActionSupport
import com.intellij.lang.parameterInfo.ParameterInfoUIContext
import com.intellij.lang.parameterInfo.ParameterInfoUtils
import com.intellij.lang.parameterInfo.UpdateParameterInfoContext
import com.intellij.psi.tree.IElementType
import com.intellij.psi.util.PsiTreeUtil
import com.intellij.psi.util.parentOfType
import com.squareup.sqldelight.core.psi.SqlDelightColumnType

class SqlDelightParameterInfoHandler : ParameterInfoHandlerWithTabActionSupport<SqlValuesExpression, List<String>, SqlBindExpr> {

  override fun couldShowInLookup(): Boolean {
    return false
  }

  override fun getParametersForLookup(
    item: LookupElement,
    context: ParameterInfoContext
  ): Array<Any> {
    return emptyArray()
  }

  override fun findElementForParameterInfo(context: CreateParameterInfoContext): SqlValuesExpression? {
    val element = context.file.findElementAt(context.offset)
    val valuesExpr = element?.parentOfType<SqlValuesExpression>() ?: return null
    if (valuesExpr.parent !is SqlInsertStmtValues) {
      return null
    }

    val columns = element.parentOfType<SqlInsertStmt>()?.columnNameList.orEmpty()
      .mapNotNull { it.reference?.resolve()?.parent as? SqlColumnDef }
      .mapNotNull { columnDef ->
        val columnType = PsiTreeUtil.getChildOfType(columnDef, SqlDelightColumnType::class.java) ?: return@mapNotNull null
        val annotations = columnType.annotationList.joinToString(", ") { "@${it.text}" }
        val columnName = columnDef.columnName.text
        val type = columnType.javaTypeName?.text ?: columnType.typeName.text
        buildString {
          if (annotations.isNotBlank()) {
            append(annotations)
            append(" ")
          }
          append("$columnName: $type")
        }
      }
    context.itemsToShow = arrayOf(columns)
    return valuesExpr
  }

  override fun findElementForUpdatingParameterInfo(context: UpdateParameterInfoContext): SqlValuesExpression? {
    return context.file.findElementAt(context.offset)?.parentOfType()
  }

  override fun showParameterInfo(
    element: SqlValuesExpression,
    context: CreateParameterInfoContext
  ) {
    context.showHint(element, element.textRange.startOffset, this)
  }

  override fun updateParameterInfo(
    parameterOwner: SqlValuesExpression,
    context: UpdateParameterInfoContext
  ) {
    val offset = context.offset
    val index = ParameterInfoUtils.getCurrentParameterIndex(parameterOwner.node, offset, SqlTypes.COMMA)
    context.setCurrentParameter(index)
  }

  override fun updateUI(p: List<String>, context: ParameterInfoUIContext) {
    val text = p.joinToString(", ")
    val index = context.currentParameterIndex
    val s = p.getOrNull(index)
    var startIndex = -1
    var endIndex = -1
    if (s != null) {
      startIndex = text.indexOf(s)
      endIndex = startIndex + s.length
    }
    context.setupUIComponentPresentation(
      text,
      startIndex,
      endIndex,
      !context.isUIComponentEnabled,
      false,
      false,
      context.defaultParameterColor
    )
  }

  override fun getActualParameters(o: SqlValuesExpression): Array<SqlBindExpr> {
    return o.exprList
      .filterIsInstance<SqlBindExpr>()
      .toTypedArray()
  }

  override fun getActualParameterDelimiterType(): IElementType {
    return SqlTypes.COMMA
  }

  override fun getActualParametersRBraceType(): IElementType = SqlTypes.RP

  override fun getArgumentListAllowedParentClasses(): Set<Class<*>> {
    return emptySet()
  }

  override fun getArgListStopSearchClasses(): Set<Class<*>> {
    return emptySet()
  }

  override fun getArgumentListClass(): Class<SqlValuesExpression> {
    return SqlValuesExpression::class.java
  }
}

package com.squareup.sqldelight.intellij

import com.alecstrong.sql.psi.core.psi.NamedElement
import com.intellij.lang.refactoring.RefactoringSupportProvider
import com.intellij.psi.PsiElement

class SqlDelightRefactoringSupportProvider : RefactoringSupportProvider() {
  override fun isMemberInplaceRenameAvailable(element: PsiElement, context: PsiElement?): Boolean {
    return element is NamedElement
  }
}

package com.squareup.sqldelight.intellij

import com.alecstrong.sql.psi.core.psi.SqlTableName
import com.alecstrong.sql.psi.core.psi.SqlViewName
import com.intellij.psi.PsiClass
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiManager
import com.intellij.psi.PsiReference
import com.intellij.psi.search.GlobalSearchScope
import com.intellij.psi.search.SearchScope
import com.intellij.refactoring.listeners.RefactoringElementListener
import com.intellij.refactoring.rename.RenamePsiElementProcessor
import com.intellij.usageView.UsageInfo
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.psi.StmtIdentifier
import com.squareup.sqldelight.core.lang.psi.StmtIdentifierMixin
import org.jetbrains.kotlin.idea.util.module
import org.jetbrains.kotlin.psi.KtFile

class SqlDelightRenameProcessor : RenamePsiElementProcessor() {
  override fun canProcessElement(element: PsiElement): Boolean {
    if (element.module == null ||
      !SqlDelightFileIndex.getInstance(element.module!!).isConfigured
    ) {
      return false
    }
    return when (element) {
      is StmtIdentifier, is SqlTableName, is SqlViewName -> true
      else -> false
    }
  }

  override fun renameElement(
    element: PsiElement,
    newName: String,
    usages: Array<out UsageInfo>,
    listener: RefactoringElementListener?
  ) {
    val newTypeName = newName.capitalize()
    val currentTypeName = when (element) {
      is StmtIdentifierMixin -> element.identifier()!!.text
      is SqlTableName -> element.name
      is SqlViewName -> element.name
      else -> throw AssertionError()
    }.capitalize()
    element.generatedTypes(currentTypeName).forEach { type ->
      element.references(type)
        .filter { it.element.containingFile.virtualFile != type.containingFile }
        .forEach { reference ->
          val currentName = reference.element.text
          reference.handleElementRename(currentName.replace(currentTypeName, newTypeName))
        }
    }
    super.renameElement(element, newName, usages, listener)
  }

  override fun findReferences(
    element: PsiElement,
    searchScope: SearchScope,
    searchInCommentsAndStrings: Boolean
  ): Collection<PsiReference> {
    if (element !is StmtIdentifierMixin) return super.findReferences(element, searchScope, searchInCommentsAndStrings)
    return element.generatedMethods().flatMap { element.references(it) }
  }

  private fun PsiElement.references(element: PsiElement): Collection<PsiReference> {
    val processor = RenamePsiElementProcessor.forElement(element)
    return processor.findReferences(element, GlobalSearchScope.projectScope(element.project), false)
      .filter { it.element.containingFile.virtualFile !in generatedQueryFiles() }
  }

  private fun PsiElement.generatedTypes(name: String): List<PsiClass> {
    val paths = (containingFile as SqlDelightFile).generatedDirectories?.map { "$it/$name.kt" }
      ?: return emptyList()
    val module = module ?: return emptyList()
    return paths.flatMap { path ->
      val vFile = SqlDelightFileIndex.getInstance(module).contentRoot.findFileByRelativePath(path)
        ?: return emptyList()
      val file = PsiManager.getInstance(project).findFile(vFile) as? KtFile ?: return emptyList()
      file.classes.asList()
    }
  }
}

package com.example

class GeneratesTypeImpl : GeneratesType

package com.example

class KotlinClass {
  init {
    val queryWrapper = QueryWrapper()
    queryWrapper.mainQueries.someQuery()

    queryWrapper.mainQueries.multiQuery()
    queryWrapper.mainQueries.multiQuery(::MultiQuery)

    queryWrapper.mainQueries.generatesType(::GeneratesType)
    queryWrapper.mainQueries.generatesType(::GeneratesTypeImpl)
  }

  enum class InnerClass {
    VAL
  }
}

package com.example

class MainImpl : Main

package com.example

class ViewImpl : SomeView

package com.squareup.sqlite.migrations

import schemacrawler.schema.Catalog
import schemacrawler.schemacrawler.SchemaCrawlerOptions
import schemacrawler.schemacrawler.SchemaInfoLevelBuilder
import schemacrawler.utility.SchemaCrawlerUtility
import java.sql.Connection
import java.sql.DriverManager
import java.sql.SQLException

class CatalogDatabase private constructor(
  internal val catalog: Catalog
) : Database() {

  companion object {

    private val schemaCrawlerOptions = SchemaCrawlerOptions().apply {
      schemaInfoLevel = SchemaInfoLevelBuilder.maximum()
      routineTypes = emptyList() // SQLite does not support stored procedures ("routines" in JBDC)
    }

    fun withInitStatements(initStatements: List<String>): CatalogDatabase {
      return fromFile("", initStatements)
    }

    fun fromFile(path: String, initStatements: List<String>): CatalogDatabase {
      return createConnection(path).init(initStatements).use {
        CatalogDatabase(SchemaCrawlerUtility.getCatalog(it, schemaCrawlerOptions))
      }
    }

    private fun createConnection(path: String): Connection {
      return try {
        DriverManager.getConnection("jdbc:sqlite:$path")
      } catch (e: SQLException) {
        DriverManager.getConnection("jdbc:sqlite:$path")
      }
    }

    private fun Connection.init(initStatements: List<String>) = apply {
      initStatements.forEach { sqlText ->
        prepareStatement(sqlText).execute()
      }
    }
  }
}

package com.squareup.sqlite.migrations

abstract class Database

package com.squareup.sqlite.migrations

interface DatabaseComparator<T : Database> {
  fun compare(db1: T, db2: T): DatabaseDiff
}

package com.squareup.sqlite.migrations

interface DatabaseDiff {
  fun printTo(out: Appendable)
}

package com.squareup.sqlite.migrations

import java.io.File
import java.util.ArrayDeque

object DatabaseFilesCollector {
  inline fun forDatabaseFiles(sourceFolders: Iterable<File>, block: (File) -> Unit) {
    val folders = ArrayDeque<File>().apply { sourceFolders.forEach(this::addLast) }
    while (folders.isNotEmpty()) {
      val folder = folders.removeFirst()
      folder.listFiles()?.forEach { file ->
        if (file.name.endsWith(".db")) {
          block(file)
        } else if (file.isDirectory) {
          folders.addLast(file)
        }
      }
    }
  }
}

package com.squareup.sqlite.migrations

import de.danielbechler.diff.ObjectDifferBuilder
import de.danielbechler.diff.node.DiffNode

class ObjectDifferDatabaseComparator(
  private val circularReferenceExceptionLogger: ((String) -> Unit)? = null
) : DatabaseComparator<CatalogDatabase> {

  override fun compare(db1: CatalogDatabase, db2: CatalogDatabase): DatabaseDiff {
    return ObjectDifferDatabaseDiff(differBuilder().compare(db1.catalog, db2.catalog))
  }

  private fun differBuilder() = ObjectDifferBuilder.startBuilding().apply {
    filtering().omitNodesWithState(DiffNode.State.UNTOUCHED)
    filtering().omitNodesWithState(DiffNode.State.CIRCULAR)
    inclusion().exclude().apply {
      propertyName("fullName")
      propertyName("parent")
      propertyName("exportedForeignKeys")
      propertyName("importedForeignKeys")
      propertyName("deferrable")
      propertyName("initiallyDeferred")
      // Definition changes aren't important, its just things like comments or whitespace.
      propertyName("definition")
    }
    // Partial columns are used for unresolved columns to avoid cycles. Matching based on string
    // is fine for our purposes.
    comparison().ofType(Class.forName("schemacrawler.crawl.ColumnPartial")).toUseEqualsMethod()

    // Custom error handler for circular reference warnings which allows to override SL4J warning log
    val circularReferenceExceptionLogger = circularReferenceExceptionLogger
    if (circularReferenceExceptionLogger != null) {
      circularReferenceHandling().handleCircularReferenceExceptionsUsing { node ->
        // The same message as original CircularReferenceExceptionHandler
        val message = (
          "Detected circular reference in node at path ${node.path} " +
            "Going deeper would cause an infinite loop, so I'll stop looking at " +
            "this instance along the current path."
          )
        circularReferenceExceptionLogger(message)
      }
    }
  }.build()
}

package com.squareup.sqlite.migrations

import de.danielbechler.diff.node.DiffNode
import schemacrawler.schema.CrawlInfo
import schemacrawler.schema.JdbcDriverInfo

class ObjectDifferDatabaseDiff(
  private val diff: DiffNode
) : DatabaseDiff {

  override fun printTo(out: Appendable) = with(out) {
    diff.visit { node, visit ->
      if (CrawlInfo::class.java.isAssignableFrom(node.valueType) ||
        JdbcDriverInfo::class.java.isAssignableFrom(node.valueType)
      ) {
        visit.dontGoDeeper()
        return@visit
      }
      if (node.childCount() == 0) {
        append("${node.path} - ${node.state}\n")
      } else if (node.state == DiffNode.State.ADDED || node.state == DiffNode.State.REMOVED) {
        append("${node.path} - ${node.state}\n")
        visit.dontGoDeeper()
      }
    }
  }

  override fun toString(): String {
    val sb = StringBuilder()
    printTo(sb)
    return sb.toString()
  }
}

package com.squareup.sqldelight.android

import android.content.Context
import android.database.Cursor
import android.util.LruCache
import androidx.sqlite.db.SupportSQLiteDatabase
import androidx.sqlite.db.SupportSQLiteOpenHelper
import androidx.sqlite.db.SupportSQLiteProgram
import androidx.sqlite.db.SupportSQLiteQuery
import androidx.sqlite.db.SupportSQLiteStatement
import androidx.sqlite.db.framework.FrameworkSQLiteOpenHelperFactory
import com.squareup.sqldelight.Transacter
import com.squareup.sqldelight.db.AfterVersion
import com.squareup.sqldelight.db.SqlCursor
import com.squareup.sqldelight.db.SqlDriver
import com.squareup.sqldelight.db.SqlPreparedStatement
import com.squareup.sqldelight.db.migrateWithCallbacks

private val DEFAULT_CACHE_SIZE = 20

class AndroidSqliteDriver private constructor(
  private val openHelper: SupportSQLiteOpenHelper? = null,
  database: SupportSQLiteDatabase? = null,
  private val cacheSize: Int
) : SqlDriver {
  init {
    require((openHelper != null) xor (database != null))
  }

  private val transactions = ThreadLocal<Transacter.Transaction>()
  private val database by lazy {
    openHelper?.writableDatabase ?: database!!
  }

  constructor(
    openHelper: SupportSQLiteOpenHelper
  ) : this(openHelper = openHelper, database = null, cacheSize = DEFAULT_CACHE_SIZE)

  /**
   * @param [cacheSize] The number of compiled sqlite statements to keep in memory per connection.
   *   Defaults to 20.
   * @param [useNoBackupDirectory] Sets whether to use a no backup directory or not.
   */
  @JvmOverloads constructor(
    schema: SqlDriver.Schema,
    context: Context,
    name: String? = null,
    factory: SupportSQLiteOpenHelper.Factory = FrameworkSQLiteOpenHelperFactory(),
    callback: SupportSQLiteOpenHelper.Callback = AndroidSqliteDriver.Callback(schema),
    cacheSize: Int = DEFAULT_CACHE_SIZE,
    useNoBackupDirectory: Boolean = false
  ) : this(
    database = null,
    openHelper = factory.create(
      SupportSQLiteOpenHelper.Configuration.builder(context)
        .callback(callback)
        .name(name)
        .noBackupDirectory(useNoBackupDirectory)
        .build()
    ),
    cacheSize = cacheSize
  )

  @JvmOverloads constructor(
    database: SupportSQLiteDatabase,
    cacheSize: Int = DEFAULT_CACHE_SIZE
  ) : this(openHelper = null, database = database, cacheSize = cacheSize)

  private val statements = object : LruCache<Int, AndroidStatement>(cacheSize) {
    override fun entryRemoved(
      evicted: Boolean,
      key: Int,
      oldValue: AndroidStatement,
      newValue: AndroidStatement?
    ) {
      if (evicted) oldValue.close()
    }
  }

  override fun newTransaction(): Transacter.Transaction {
    val enclosing = transactions.get()
    val transaction = Transaction(enclosing)
    transactions.set(transaction)

    if (enclosing == null) {
      database.beginTransactionNonExclusive()
    }

    return transaction
  }

  override fun currentTransaction() = transactions.get()

  inner class Transaction(
    override val enclosingTransaction: Transacter.Transaction?
  ) : Transacter.Transaction() {
    override fun endTransaction(successful: Boolean) {
      if (enclosingTransaction == null) {
        if (successful) {
          database.setTransactionSuccessful()
          database.endTransaction()
        } else {
          database.endTransaction()
        }
      }
      transactions.set(enclosingTransaction)
    }
  }

  private fun <T> execute(
    identifier: Int?,
    createStatement: () -> AndroidStatement,
    binders: (SqlPreparedStatement.() -> Unit)?,
    result: AndroidStatement.() -> T
  ): T {
    var statement: AndroidStatement? = null
    if (identifier != null) {
      statement = statements.remove(identifier)
    }
    if (statement == null) {
      statement = createStatement()
    }
    try {
      if (binders != null) { statement.binders() }
      return statement.result()
    } finally {
      if (identifier != null) {
        statements.put(identifier, statement)?.close()
      } else {
        statement.close()
      }
    }
  }

  override fun execute(
    identifier: Int?,
    sql: String,
    parameters: Int,
    binders: (SqlPreparedStatement.() -> Unit)?
  ) = execute(identifier, { AndroidPreparedStatement(database.compileStatement(sql)) }, binders, AndroidStatement::execute)

  override fun executeQuery(
    identifier: Int?,
    sql: String,
    parameters: Int,
    binders: (SqlPreparedStatement.() -> Unit)?
  ) = execute(identifier, { AndroidQuery(sql, database, parameters) }, binders, AndroidStatement::executeQuery)

  override fun close() {
    statements.evictAll()
    return openHelper?.close() ?: database.close()
  }

  open class Callback(
    private val schema: SqlDriver.Schema,
    vararg callbacks: AfterVersion,
  ) : SupportSQLiteOpenHelper.Callback(schema.version) {
    private val callbacks = callbacks

    constructor(
      schema: SqlDriver.Schema
    ) : this(schema, *emptyArray())

    override fun onCreate(db: SupportSQLiteDatabase) {
      schema.create(AndroidSqliteDriver(openHelper = null, database = db, cacheSize = 1))
    }

    override fun onUpgrade(
      db: SupportSQLiteDatabase,
      oldVersion: Int,
      newVersion: Int
    ) {
      if (callbacks.isNotEmpty()) {
        schema.migrateWithCallbacks(AndroidSqliteDriver(openHelper = null, database = db, cacheSize = 1), oldVersion, newVersion, *callbacks)
      } else {
        schema.migrate(
          AndroidSqliteDriver(openHelper = null, database = db, cacheSize = 1),
          oldVersion, newVersion
        )
      }
    }
  }
}

internal interface AndroidStatement : SqlPreparedStatement {
  fun execute()
  fun executeQuery(): SqlCursor
  fun close()
}

private class AndroidPreparedStatement(
  private val statement: SupportSQLiteStatement
) : AndroidStatement {
  override fun bindBytes(index: Int, value: ByteArray?) {
    if (value == null) statement.bindNull(index) else statement.bindBlob(index, value)
  }

  override fun bindLong(index: Int, value: Long?) {
    if (value == null) statement.bindNull(index) else statement.bindLong(index, value)
  }

  override fun bindDouble(index: Int, value: Double?) {
    if (value == null) statement.bindNull(index) else statement.bindDouble(index, value)
  }

  override fun bindString(index: Int, value: String?) {
    if (value == null) statement.bindNull(index) else statement.bindString(index, value)
  }

  override fun executeQuery() = throw UnsupportedOperationException()

  override fun execute() {
    statement.execute()
  }

  override fun close() {
    statement.close()
  }
}

private class AndroidQuery(
  private val sql: String,
  private val database: SupportSQLiteDatabase,
  private val argCount: Int
) : SupportSQLiteQuery, AndroidStatement {
  private val binds: MutableMap<Int, (SupportSQLiteProgram) -> Unit> = LinkedHashMap()

  override fun bindBytes(index: Int, bytes: ByteArray?) {
    binds[index] = { if (bytes == null) it.bindNull(index) else it.bindBlob(index, bytes) }
  }

  override fun bindLong(index: Int, long: Long?) {
    binds[index] = { if (long == null) it.bindNull(index) else it.bindLong(index, long) }
  }

  override fun bindDouble(index: Int, double: Double?) {
    binds[index] = { if (double == null) it.bindNull(index) else it.bindDouble(index, double) }
  }

  override fun bindString(index: Int, string: String?) {
    binds[index] = { if (string == null) it.bindNull(index) else it.bindString(index, string) }
  }

  override fun execute() = throw UnsupportedOperationException()

  override fun executeQuery() = AndroidCursor(database.query(this))

  override fun bindTo(statement: SupportSQLiteProgram) {
    for (action in binds.values) {
      action(statement)
    }
  }

  override fun getSql() = sql

  override fun toString() = sql

  override fun getArgCount() = argCount

  override fun close() { }
}

private class AndroidCursor(
  private val cursor: Cursor
) : SqlCursor {
  override fun next() = cursor.moveToNext()
  override fun getString(index: Int) = if (cursor.isNull(index)) null else cursor.getString(index)
  override fun getLong(index: Int) = if (cursor.isNull(index)) null else cursor.getLong(index)
  override fun getBytes(index: Int) = if (cursor.isNull(index)) null else cursor.getBlob(index)
  override fun getDouble(index: Int) = if (cursor.isNull(index)) null else cursor.getDouble(index)
  override fun close() = cursor.close()
}

package com.example.sqldelight.hockey

import com.example.sqldelight.hockey.data.Db
import com.example.sqldelight.hockey.data.getInstance
import com.example.sqldelight.hockey.platform.defaultFormatter
import kotlinx.browser.document
import kotlinx.html.dom.create
import kotlinx.html.table
import kotlinx.html.tbody
import kotlinx.html.td
import kotlinx.html.th
import kotlinx.html.thead
import kotlinx.html.tr

fun main() {
  Db.getInstance().then { db ->

    val players = db.playerQueries.forTeam(-1).executeAsList()
    val playersTable = document.create.table {
      thead {
        tr {
          th { +"First Name" }
          th { +"Last Name" }
          th { +"Team Name" }
          th { +"Number" }
        }
      }
      tbody {
        players.forEach { player ->
          tr {
            td { +player.first_name }
            td { +player.last_name }
            td { +player.teamName }
            td { +player.number }
          }
        }
      }
    }
    document.getElementById("players")?.append(playersTable)

    val teams = db.teamQueries.selectAll().executeAsList()
    val teamsTable = document.create.table {
      thead {
        tr {
          th { +"Name" }
          th { +"Coach" }
          th { +"Founded" }
        }
      }
      tbody {
        teams.forEach { team ->
          tr {
            td { +team.name }
            td { +team.coach }
            td { +defaultFormatter.format(team.founded) }
          }
        }
      }
    }
    document.getElementById("teams")?.append(teamsTable)
  }
}

package com.squareup.sqldelight.core.compiler

import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.KModifier
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.kotlinpoet.TypeSpec
import com.squareup.kotlinpoet.asTypeName
import com.squareup.kotlinpoet.joinToCode
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.lang.DATABASE_SCHEMA_TYPE
import kotlin.reflect.KClass

internal class DatabaseExposerGenerator(
  val implementation: TypeSpec,
  val fileIndex: SqlDelightFileIndex
) {
  private val interfaceType = ClassName(fileIndex.packageName, fileIndex.className)

  fun exposedSchema(): PropertySpec {
    return PropertySpec.builder("schema", DATABASE_SCHEMA_TYPE)
      .addModifiers(KModifier.INTERNAL)
      .receiver(KClass::class.asTypeName().parameterizedBy(interfaceType))
      .getter(
        FunSpec.getterBuilder()
          .addStatement("return %N.Schema", implementation)
          .build()
      )
      .build()
  }

  fun exposedConstructor(): FunSpec {
    return implementation.primaryConstructor!!.let { constructor ->
      FunSpec.builder("newInstance")
        .addModifiers(KModifier.INTERNAL)
        .returns(ClassName(fileIndex.packageName, fileIndex.className))
        .receiver(KClass::class.asTypeName().parameterizedBy(interfaceType))
        .addParameters(constructor.parameters)
        .addStatement("return %N(%L)", implementation, constructor.parameters.map { CodeBlock.of("%N", it) }.joinToCode(", "))
        .build()
    }
  }
}

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.core.compiler

import com.alecstrong.sql.psi.core.psi.SqlStmt
import com.intellij.openapi.module.Module
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.INT
import com.squareup.kotlinpoet.KModifier.OPERATOR
import com.squareup.kotlinpoet.KModifier.OVERRIDE
import com.squareup.kotlinpoet.KModifier.PRIVATE
import com.squareup.kotlinpoet.ParameterSpec
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.kotlinpoet.TypeSpec
import com.squareup.sqldelight.core.SqlDelightException
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.compiler.integration.adapterProperty
import com.squareup.sqldelight.core.compiler.integration.needsAdapters
import com.squareup.sqldelight.core.lang.DATABASE_SCHEMA_TYPE
import com.squareup.sqldelight.core.lang.DRIVER_NAME
import com.squareup.sqldelight.core.lang.DRIVER_TYPE
import com.squareup.sqldelight.core.lang.MigrationFile
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.SqlDelightQueriesFile
import com.squareup.sqldelight.core.lang.TRANSACTER_IMPL_TYPE
import com.squareup.sqldelight.core.lang.TRANSACTER_TYPE
import com.squareup.sqldelight.core.lang.queriesImplType
import com.squareup.sqldelight.core.lang.queriesName
import com.squareup.sqldelight.core.lang.queriesType
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.lang.util.forInitializationStatements
import com.squareup.sqldelight.core.lang.util.rawSqlText

internal class DatabaseGenerator(
  module: Module,
  sourceFile: SqlDelightFile
) {
  private val sourceFolders = SqlDelightFileIndex.getInstance(module).sourceFolders(sourceFile)
  private val moduleFolders = SqlDelightFileIndex.getInstance(module)
    .sourceFolders(sourceFile, includeDependencies = false)
  private val fileIndex = SqlDelightFileIndex.getInstance(module)
  private val type = ClassName(fileIndex.packageName, fileIndex.className)

  fun interfaceType(): TypeSpec {
    val typeSpec = TypeSpec.interfaceBuilder(fileIndex.className)
      .addSuperinterface(TRANSACTER_TYPE)

    fileIndex.dependencies.forEach {
      typeSpec.addSuperinterface(ClassName(it.packageName, it.className))
    }

    val invoke = FunSpec.builder("invoke")
      .returns(ClassName(fileIndex.packageName, fileIndex.className))
      .addModifiers(OPERATOR)

    val invokeReturn = CodeBlock.builder()
      .add("return %T::class.newInstance(", type)

    // Database constructor parameter:
    // driver: SqlDriver
    val dbParameter = ParameterSpec.builder(DRIVER_NAME, DRIVER_TYPE).build()
    invoke.addParameter(dbParameter)
    invokeReturn.add("%N", dbParameter)

    moduleFolders.flatMap { it.findChildrenOfType<SqlDelightQueriesFile>() }
      .sortedBy { it.name }
      .forEach { file ->
        // queries property added to QueryWrapper type:
        // val dataQueries = DataQueries(this, driver, transactions)
        typeSpec.addProperty(file.queriesName, file.queriesType)
      }

    forAdapters {
      invoke.addParameter(it.name, it.type)
      invokeReturn.add(", %L", it.name)
    }

    return typeSpec
      .addType(
        TypeSpec.companionObjectBuilder()
          .addProperty(
            PropertySpec.builder("Schema", DATABASE_SCHEMA_TYPE)
              .getter(
                FunSpec.getterBuilder()
                  .addStatement("return %T::class.schema", type)
                  .build()
              )
              .build()
          )
          .addFunction(
            invoke
              .addCode(
                invokeReturn
                  .add(")")
                  .build()
              )
              .build()
          )
          .build()
      )
      .build()
  }

  private fun forAdapters(
    block: (PropertySpec) -> Unit
  ) {
    val queriesFile = sourceFolders.flatMap { it.findChildrenOfType<SqlDelightQueriesFile>() }
      .firstOrNull() ?: return
    queriesFile.tables(true)
      .toSet()
      .mapNotNull { if (it.needsAdapters()) it.adapterProperty() else null }
      .sortedBy { it.name }
      .forEach(block)
  }

  fun type(implementationPackage: String): TypeSpec {
    val typeSpec = TypeSpec.classBuilder("${fileIndex.className}Impl")
      .superclass(TRANSACTER_IMPL_TYPE)
      .addModifiers(PRIVATE)
      .addSuperclassConstructorParameter(DRIVER_NAME)

    val constructor = FunSpec.constructorBuilder()

    // Database constructor parameter:
    // driver: SqlDriver
    val dbParameter = ParameterSpec.builder(DRIVER_NAME, DRIVER_TYPE).build()
    constructor.addParameter(dbParameter)

    // Static on create function:
    // fun create(driver: SqlDriver)
    val createFunction = FunSpec.builder("create")
      .addModifiers(OVERRIDE)
      .addParameter(DRIVER_NAME, DRIVER_TYPE)

    val oldVersion = ParameterSpec.builder("oldVersion", INT).build()
    val newVersion = ParameterSpec.builder("newVersion", INT).build()

    val migrateFunction = FunSpec.builder("migrate")
      .addModifiers(OVERRIDE)
      .addParameter(DRIVER_NAME, DRIVER_TYPE)
      .addParameter(oldVersion)
      .addParameter(newVersion)

    forAdapters {
      typeSpec.addProperty(it)
      constructor.addParameter(it.name, it.type)
    }

    sourceFolders.flatMap { it.findChildrenOfType<SqlDelightQueriesFile>() }
      .sortedBy { it.name }
      .forEach { file ->
        // queries property added to QueryWrapper type:
        // val dataQueries = DataQueries(this, driver, transactions)
        typeSpec.addProperty(
          PropertySpec.builder(file.queriesName, file.queriesImplType(implementationPackage))
            .addModifiers(OVERRIDE)
            .initializer("%T(this, $DRIVER_NAME)", file.queriesImplType(implementationPackage))
            .build()
        )
      }

    if (!fileIndex.deriveSchemaFromMigrations) {
      // Derive the schema from queries files.
      sourceFolders.flatMap { it.findChildrenOfType<SqlDelightQueriesFile>() }
        .forInitializationStatements { sqlText ->
          createFunction.addStatement("$DRIVER_NAME.execute(null, %L, 0)", sqlText.toCodeLiteral())
        }
    } else {
      val orderedMigrations = sourceFolders.flatMap { it.findChildrenOfType<MigrationFile>() }
        .sortedBy { it.order }

      // Derive the schema from migration files.
      orderedMigrations.flatMap { it.sqliteStatements() }
        .filter { it.isSchema() }
        .forEach {
          createFunction.addStatement("$DRIVER_NAME.execute(null, %L, 0)", it.rawSqlText().toCodeLiteral())
        }
    }

    var maxVersion = 1

    sourceFolders.flatMap { it.findChildrenOfType<MigrationFile>() }
      .sortedBy { it.version }
      .forEach { migrationFile ->
        try {
          maxVersion = maxOf(maxVersion, migrationFile.version + 1)
        } catch (e: Throwable) {
          throw SqlDelightException("Migration files can only have versioned names (1.sqm, 2.sqm, etc)")
        }
        migrateFunction.beginControlFlow(
          "if (%N <= ${migrationFile.version} && %N > ${migrationFile.version})",
          oldVersion, newVersion
        )
        migrationFile.sqliteStatements().forEach {
          migrateFunction.addStatement("$DRIVER_NAME.execute(null, %S, 0)", it.rawSqlText())
        }
        migrateFunction.endControlFlow()
      }

    return typeSpec
      .addType(
        TypeSpec.objectBuilder(DATABASE_SCHEMA_TYPE.simpleName)
          .addSuperinterface(DATABASE_SCHEMA_TYPE)
          .addFunction(createFunction.build())
          .addFunction(migrateFunction.build())
          .addProperty(
            PropertySpec.builder("version", INT, OVERRIDE)
              .getter(FunSpec.getterBuilder().addStatement("return $maxVersion").build())
              .build()
          )
          .build()
      )
      .addSuperinterface(ClassName(fileIndex.packageName, fileIndex.className))
      .primaryConstructor(constructor.build())
      .build()
  }

  private fun SqlStmt.isSchema() = when {
    createIndexStmt != null -> true
    createTableStmt != null -> true
    createTriggerStmt != null -> true
    createViewStmt != null -> true
    createVirtualTableStmt != null -> true
    alterTableStmt != null -> true
    dropIndexStmt != null -> true
    dropTableStmt != null -> true
    dropTriggerStmt != null -> true
    dropViewStmt != null -> true
    else -> false
  }
}

package com.squareup.sqldelight.core.compiler

import com.alecstrong.sql.psi.core.psi.SqlDeleteStmtLimited
import com.alecstrong.sql.psi.core.psi.SqlInsertStmt
import com.alecstrong.sql.psi.core.psi.SqlUpdateStmtLimited
import com.intellij.psi.util.PsiTreeUtil
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.KModifier
import com.squareup.kotlinpoet.ParameterSpec
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.kotlinpoet.joinToCode
import com.squareup.sqldelight.core.compiler.model.NamedExecute
import com.squareup.sqldelight.core.compiler.model.NamedMutator
import com.squareup.sqldelight.core.compiler.model.NamedQuery
import com.squareup.sqldelight.core.lang.psi.StmtIdentifierMixin
import com.squareup.sqldelight.core.psi.SqlDelightStmtClojureStmtList

open class ExecuteQueryGenerator(private val query: NamedExecute) : QueryGenerator(query) {
  internal open fun queriesUpdated(): List<NamedQuery> {
    if (query.statement is SqlDelightStmtClojureStmtList) {
      return PsiTreeUtil.findChildrenOfAnyType(
        query.statement,
        SqlUpdateStmtLimited::class.java,
        SqlDeleteStmtLimited::class.java,
        SqlInsertStmt::class.java
      ).flatMap {
        MutatorQueryGenerator(
          when (it) {
            is SqlUpdateStmtLimited -> NamedMutator.Update(it, query.identifier as StmtIdentifierMixin)
            is SqlDeleteStmtLimited -> NamedMutator.Delete(it, query.identifier as StmtIdentifierMixin)
            is SqlInsertStmt -> NamedMutator.Insert(it, query.identifier as StmtIdentifierMixin)
            else -> throw IllegalArgumentException("Unexpected statement $it")
          }
        ).queriesUpdated()
      }.distinct()
    }
    return emptyList()
  }

  private fun FunSpec.Builder.notifyQueries(): FunSpec.Builder {
    val resultSetsUpdated = queriesUpdated()

    if (resultSetsUpdated.isEmpty()) return this

    // The list of effected queries:
    // (queryWrapper.dataQueries.selectForId + queryWrapper.otherQueries.selectForId)
    // TODO: Only notify queries that were dirtied (check using dirtied method).
    addStatement(
      "notifyQueries(%L, {%L})",
      query.id,
      resultSetsUpdated.sortedBy { it.id }.map { it.queryProperty }.joinToCode(separator = " + ")
    )

    return this
  }

  /**
   * The public api to execute [query]
   */
  fun function(): FunSpec {
    return interfaceFunction()
      .addModifiers(KModifier.OVERRIDE)
      .addCode(executeBlock())
      .notifyQueries()
      .build()
  }

  fun interfaceFunction(): FunSpec.Builder {
    return FunSpec.builder(query.name)
      .also(this::addJavadoc)
      .addParameters(
        query.parameters.map {
          ParameterSpec.builder(it.name, it.argumentType()).build()
        }
      )
  }

  fun value(): PropertySpec {
    return PropertySpec.builder(query.name, ClassName("", query.name.capitalize()), KModifier.PRIVATE)
      .initializer("${query.name.capitalize()}()")
      .build()
  }
}

package com.squareup.sqldelight.core.compiler

import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.MemberName
import com.squareup.kotlinpoet.buildCodeBlock

/**
 * The maximum string length that we can emit as a literal. This value is 2^15 instead of the real
 * value, 2^16, to account for the indentation that KotlinPoet will add by using `trimMargin`.
 */
private const val MAX_STRING_LENGTH = 32_768

private val BUILD_STRING = MemberName("kotlin.text", "buildString")

/**
 * Returns a [CodeBlock] which contains a representation of [this] as a literal. If longer than
 * the maximum allowed length for a string, it will be broken up into multiple literals with code
 * to re-assemble at runtime.
 */
internal fun String.toCodeLiteral(): CodeBlock {
  if (length < MAX_STRING_LENGTH) {
    return CodeBlock.of("%S", this)
  }
  return buildCodeBlock {
    add("%M(%L) {\n", BUILD_STRING, length)
    for (i in 0 until length step MAX_STRING_LENGTH) {
      add("append(%S)\n", substring(i, (i + MAX_STRING_LENGTH).coerceAtMost(length)))
    }
    add("}")
  }
}

package com.squareup.sqldelight.core.compiler

import com.alecstrong.sql.psi.core.psi.SqlForeignKeyClause
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.squareup.sqldelight.core.compiler.model.NamedMutator
import com.squareup.sqldelight.core.compiler.model.NamedQuery
import com.squareup.sqldelight.core.lang.util.childOfType
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.lang.util.referencedTables

class MutatorQueryGenerator(
  private val query: NamedMutator
) : ExecuteQueryGenerator(query) {
  override fun queriesUpdated(): List<NamedQuery> {
    val resultSetsUpdated = mutableListOf<NamedQuery>()
    val foreignKeyCascadeCheck = when (query) {
      is NamedMutator.Delete -> SqlTypes.DELETE
      is NamedMutator.Update -> SqlTypes.UPDATE
      else -> null
    }

    query.containingFile.iterateSqlFiles { psiFile ->
      val tablesAffected = mutableListOf(query.tableEffected)

      if (foreignKeyCascadeCheck != null) {
        psiFile.sqlStmtList?.stmtList?.mapNotNull { it.createTableStmt }?.forEach { table ->
          val effected = table.findChildrenOfType<SqlForeignKeyClause>().any {
            (it.foreignTable.name == query.tableEffected.name) && it.node.findChildByType(foreignKeyCascadeCheck) != null
          }
          if (effected) tablesAffected.add(table.tableName)
        }
      }

      psiFile.triggers.forEach { trigger ->
        if (trigger.tableName?.name == query.tableEffected.name) {
          val triggered = when (query) {
            is NamedMutator.Delete -> trigger.childOfType(SqlTypes.DELETE) != null
            is NamedMutator.Insert -> {
              trigger.childOfType(SqlTypes.INSERT) != null ||
                (query.hasUpsertClause && trigger.childOfType(SqlTypes.UPDATE) != null)
            }
            is NamedMutator.Update -> {
              val columns = trigger.columnNameList.map { it.name }
              val updateColumns = query.update.updateStmtSubsequentSetterList.map { it.columnName?.name } +
                query.update.columnName?.name
              trigger.childOfType(SqlTypes.UPDATE) != null && (
                columns.isEmpty() ||
                  updateColumns.any { it in columns }
                )
            }
          }

          if (triggered) {
            // Also need to notify for the trigger.
            tablesAffected += trigger.insertStmtList.map {
              it.tableName.referencedTables().single()
            }
            tablesAffected += trigger.deleteStmtList.map {
              it.qualifiedTableName!!.tableName.referencedTables().single()
            }
            tablesAffected += trigger.updateStmtList.map {
              it.qualifiedTableName.tableName.referencedTables().single()
            }
          }
        }
      }

      resultSetsUpdated.addAll(
        psiFile.namedQueries
          .filter { query -> query.tablesObserved.any { it in tablesAffected } }
      )
    }

    return resultSetsUpdated
  }
}

package com.squareup.sqldelight.core.compiler

import com.intellij.openapi.module.Module
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.KModifier.ABSTRACT
import com.squareup.kotlinpoet.KModifier.PRIVATE
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.kotlinpoet.TypeSpec
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.compiler.model.NamedExecute
import com.squareup.sqldelight.core.compiler.model.NamedMutator
import com.squareup.sqldelight.core.lang.CUSTOM_DATABASE_NAME
import com.squareup.sqldelight.core.lang.DRIVER_NAME
import com.squareup.sqldelight.core.lang.DRIVER_TYPE
import com.squareup.sqldelight.core.lang.SqlDelightQueriesFile
import com.squareup.sqldelight.core.lang.TRANSACTER_IMPL_TYPE
import com.squareup.sqldelight.core.lang.TRANSACTER_TYPE
import com.squareup.sqldelight.core.lang.queriesImplType
import com.squareup.sqldelight.core.lang.queriesType

class QueriesTypeGenerator(
  private val module: Module,
  private val file: SqlDelightQueriesFile
) {

  fun interfaceType(): TypeSpec {
    val type = TypeSpec.interfaceBuilder(file.queriesType.simpleName)
      .addSuperinterface(TRANSACTER_TYPE)

    file.namedQueries.forEach { query ->
      tryWithElement(query.select) {
        val generator = SelectQueryGenerator(query)

        type.addFunction(
          generator.customResultTypeFunctionInterface()
            .addModifiers(ABSTRACT)
            .build()
        )

        if (query.needsWrapper()) {
          type.addFunction(
            generator.defaultResultTypeFunctionInterface()
              .addModifiers(ABSTRACT)
              .build()
          )
        }
      }
    }

    file.namedMutators.forEach { mutator ->
      type.addExecute(mutator, true)
    }

    file.namedExecutes.forEach { execute ->
      type.addExecute(execute, true)
    }

    return type.build()
  }

  /**
   * Generate the full queries object - done once per file, containing all labeled select and
   * mutator queries.
   *
   * eg: class DataQueries(
   *       private val queryWrapper: QueryWrapper,
   *       private val driver: SqlDriver,
   *       transactions: ThreadLocal<Transacter.Transaction>
   *     ) : TransacterImpl(driver, transactions)
   */
  fun generateType(packageName: String): TypeSpec {
    val type = TypeSpec.classBuilder(file.queriesImplType(packageName).simpleName)
      .addModifiers(PRIVATE)
      .superclass(TRANSACTER_IMPL_TYPE)
      .addSuperinterface(file.queriesType)

    val constructor = FunSpec.constructorBuilder()

    // Add the query wrapper as a constructor property:
    // private val queryWrapper: QueryWrapper
    val databaseType = ClassName(packageName, "${SqlDelightFileIndex.getInstance(module).className}Impl")

    type.addProperty(
      PropertySpec.builder(CUSTOM_DATABASE_NAME, databaseType, PRIVATE)
        .initializer(CUSTOM_DATABASE_NAME)
        .build()
    )
    constructor.addParameter(CUSTOM_DATABASE_NAME, databaseType)

    // Add the database as a constructor property and superclass parameter:
    // private val driver: SqlDriver
    type.addProperty(
      PropertySpec.builder(DRIVER_NAME, DRIVER_TYPE, PRIVATE)
        .initializer(DRIVER_NAME)
        .build()
    )
    constructor.addParameter(DRIVER_NAME, DRIVER_TYPE)
    type.addSuperclassConstructorParameter(DRIVER_NAME)

    file.namedQueries.forEach { query ->
      tryWithElement(query.select) {
        val generator = SelectQueryGenerator(query)

        type.addProperty(generator.queryCollectionProperty())
        type.addFunction(generator.customResultTypeFunction())

        if (query.needsWrapper()) {
          type.addFunction(generator.defaultResultTypeFunction())
        }

        if (query.arguments.isNotEmpty()) {
          type.addType(generator.querySubtype())
        }
      }
    }

    file.namedMutators.forEach { mutator ->
      type.addExecute(mutator, false)
    }

    file.namedExecutes.forEach { execute ->
      type.addExecute(execute, false)
    }

    return type.primaryConstructor(constructor.build())
      .build()
  }

  private fun TypeSpec.Builder.addExecute(execute: NamedExecute, forInterface: Boolean) {
    tryWithElement(execute.statement) {
      val generator = if (execute is NamedMutator) {
        MutatorQueryGenerator(execute)
      } else {
        ExecuteQueryGenerator(execute)
      }

      addFunction(
        if (forInterface) generator.interfaceFunction().addModifiers(ABSTRACT).build()
        else generator.function()
      )
    }
  }
}

package com.squareup.sqldelight.core.compiler

import com.alecstrong.sql.psi.core.psi.SqlBinaryEqualityExpr
import com.alecstrong.sql.psi.core.psi.SqlBindExpr
import com.alecstrong.sql.psi.core.psi.SqlStmt
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiWhiteSpace
import com.intellij.psi.util.PsiTreeUtil
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.NameAllocator
import com.squareup.sqldelight.core.compiler.integration.javadocText
import com.squareup.sqldelight.core.compiler.model.BindableQuery
import com.squareup.sqldelight.core.compiler.model.NamedExecute
import com.squareup.sqldelight.core.compiler.model.NamedQuery
import com.squareup.sqldelight.core.lang.DRIVER_NAME
import com.squareup.sqldelight.core.lang.IntermediateType
import com.squareup.sqldelight.core.lang.util.childOfType
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.lang.util.isArrayParameter
import com.squareup.sqldelight.core.lang.util.range
import com.squareup.sqldelight.core.lang.util.rawSqlText
import com.squareup.sqldelight.core.psi.SqlDelightStmtClojureStmtList

abstract class QueryGenerator(private val query: BindableQuery) {
  /**
   * Creates the block of code that prepares [query] as a prepared statement and binds the
   * arguments to it. This code block does not make any use of class fields, and only populates a
   * single variable [STATEMENT_NAME]
   *
   * val numberIndexes = createArguments(count = number.size)
   * val statement = database.prepareStatement("""
   *     |SELECT *
   *     |FROM player
   *     |WHERE number IN $numberIndexes
   *     """.trimMargin(), SqlPreparedStatement.Type.SELECT, 1 + (number.size - 1))
   * number.forEachIndexed { index, number ->
   *     statement.bindLong(index + 2, number)
   *     }
   */
  protected fun executeBlock(): CodeBlock {
    val result = CodeBlock.builder()

    if (query is NamedExecute && query.statement is SqlDelightStmtClojureStmtList) {
      query.statement.findChildrenOfType<SqlStmt>().forEachIndexed { index, statement ->
        result.add(executeBlock(statement, query.idForIndex(index)))
      }
    } else {
      result.add(executeBlock(query.statement, query.id))
    }

    return result.build()
  }

  private fun executeBlock(
    statement: PsiElement,
    id: Int
  ): CodeBlock {
    val result = CodeBlock.builder()

    val positionToArgument = mutableListOf<Triple<Int, BindableQuery.Argument, SqlBindExpr?>>()
    val seenArgs = mutableSetOf<BindableQuery.Argument>()
    val duplicateTypes = mutableSetOf<IntermediateType>()
    query.arguments.forEach { argument ->
      if (argument.bindArgs.isNotEmpty()) {
        argument.bindArgs
          .filter { PsiTreeUtil.isAncestor(statement, it, true) }
          .forEach { bindArg ->
            if (!seenArgs.add(argument)) {
              duplicateTypes.add(argument.type)
            }
            positionToArgument.add(Triple(bindArg.node.textRange.startOffset, argument, bindArg))
          }
      } else {
        positionToArgument.add(Triple(0, argument, null))
      }
    }

    val bindStatements = CodeBlock.builder()
    val replacements = mutableListOf<Pair<IntRange, String>>()
    val argumentCounts = mutableListOf<String>()

    var needsFreshStatement = false

    val seenArrayArguments = mutableSetOf<BindableQuery.Argument>()

    val argumentNameAllocator = NameAllocator().apply {
      query.arguments.forEach { newName(it.type.name) }
    }

    // A list of [SqlBindExpr] in order of appearance in the query.
    val orderedBindArgs = positionToArgument.sortedBy { it.first }

    // The number of non-array bindArg's we've encountered so far.
    var nonArrayBindArgsCount = 0

    // A list of arrays we've encountered so far.
    val precedingArrays = mutableListOf<String>()

    val extractedVariables = mutableMapOf<IntermediateType, String>()
    // extract the variable for duplicate types, so we don't encode twice
    for (type in duplicateTypes) {
      val encodedJavaType = type.encodedJavaType() ?: continue
      val variableName = argumentNameAllocator.newName(type.name)
      extractedVariables[type] = variableName
      bindStatements.addStatement("val %N = $encodedJavaType", variableName)
    }
    // For each argument in the sql
    orderedBindArgs.forEach { (_, argument, bindArg) ->
      val type = argument.type
      // Need to replace the single argument with a group of indexed arguments, calculated at
      // runtime from the list parameter:
      // val idIndexes = id.mapIndexed { index, _ -> "?${1 + previousArray.size + index}" }.joinToString(prefix = "(", postfix = ")")
      val offset = (precedingArrays.map { "$it.size" } + "${nonArrayBindArgsCount + 1}")
        .joinToString(separator = " + ")
      if (bindArg?.isArrayParameter() == true) {
        needsFreshStatement = true

        if (seenArrayArguments.add(argument)) {
          result.addStatement(
            """
            |val ${type.name}Indexes = createArguments(count = ${type.name}.size)
          """.trimMargin()
          )
        }

        // Replace the single bind argument with the array of bind arguments:
        // WHERE id IN ${idIndexes}
        replacements.add(bindArg.range to "\$${type.name}Indexes")

        // Perform the necessary binds:
        // id.forEachIndex { index, parameter ->
        //   statement.bindLong(1 + previousArray.size + index, parameter)
        // }
        val indexCalculator = "index + $offset"
        val elementName = argumentNameAllocator.newName(type.name)
        bindStatements.addStatement(
          """
          |${type.name}.forEachIndexed { index, $elementName ->
          |%L}
        """.trimMargin(),
          type.copy(name = elementName).preparedStatementBinder(indexCalculator)
        )

        precedingArrays.add(type.name)
        argumentCounts.add("${type.name}.size")
      } else {
        nonArrayBindArgsCount += 1
        if (type.javaType.isNullable) {
          val parent = bindArg?.parent
          if (parent is SqlBinaryEqualityExpr) {
            needsFreshStatement = true

            var symbol = parent.childOfType(SqlTypes.EQ) ?: parent.childOfType(SqlTypes.EQ2)
            val nullableEquality: String
            if (symbol != null) {
              nullableEquality = "${symbol.leftWhitspace()}IS${symbol.rightWhitespace()}"
            } else {
              symbol = parent.childOfType(SqlTypes.NEQ) ?: parent.childOfType(SqlTypes.NEQ2)!!
              nullableEquality = "${symbol.leftWhitspace()}IS NOT${symbol.rightWhitespace()}"
            }

            val block = CodeBlock.of("if (${type.name} == null) \"$nullableEquality\" else \"${symbol.text}\"")
            replacements.add(symbol.range to "\${ $block }")
          }
        }
        // Binds each parameter to the statement:
        // statement.bindLong(1, id)
        bindStatements.add(type.preparedStatementBinder(offset, extractedVariables[type]))

        // Replace the named argument with a non named/indexed argument.
        // This allows us to use the same algorithm for non Sqlite dialects
        // :name becomes ?
        if (bindArg != null) {
          replacements.add(bindArg.range to "?")
        }
      }
    }

    // Adds the actual SqlPreparedStatement:
    // statement = database.prepareStatement("SELECT * FROM test")
    val executeMethod = if (query is NamedQuery) {
      "return $DRIVER_NAME.executeQuery"
    } else {
      "$DRIVER_NAME.execute"
    }
    if (nonArrayBindArgsCount != 0) {
      argumentCounts.add(0, nonArrayBindArgsCount.toString())
    }
    val arguments = mutableListOf<Any>(
      statement.rawSqlText(replacements),
      argumentCounts.ifEmpty { listOf(0) }.joinToString(" + ")
    )
    val binder: String

    if (argumentCounts.isEmpty()) {
      binder = ""
    } else {
      arguments.add(
        CodeBlock.builder()
          .addStatement(" {")
          .indent()
          .add(bindStatements.build())
          .unindent()
          .add("}")
          .build()
      )
      binder = "%L"
    }
    result.add(
      "$executeMethod(" +
        "${if (needsFreshStatement) "null" else "$id"}," +
        " %P," +
        " %L" +
        ")$binder\n",
      *arguments.toTypedArray()
    )

    return result.build()
  }

  private fun PsiElement.leftWhitspace(): String {
    return if (prevSibling is PsiWhiteSpace) "" else " "
  }

  private fun PsiElement.rightWhitespace(): String {
    return if (nextSibling is PsiWhiteSpace) "" else " "
  }

  protected fun addJavadoc(builder: FunSpec.Builder) {
    if (query.javadoc != null) javadocText(query.javadoc)?.let { builder.addKdoc(it) }
  }
}

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.core.compiler

import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.KModifier.DATA
import com.squareup.kotlinpoet.KModifier.OVERRIDE
import com.squareup.kotlinpoet.MemberName
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.kotlinpoet.TypeSpec
import com.squareup.kotlinpoet.asClassName
import com.squareup.kotlinpoet.joinToCode
import com.squareup.sqldelight.core.compiler.model.NamedQuery
import com.squareup.sqldelight.core.lang.psi.ColumnTypeMixin.Companion.isArrayType

class QueryInterfaceGenerator(val query: NamedQuery) {
  fun kotlinImplementationSpec(): TypeSpec {
    val typeSpec = TypeSpec.classBuilder(query.name.capitalize())
      .addModifiers(DATA)

    val propertyPrints = mutableListOf<CodeBlock>()
    val contentToString = MemberName("kotlin.collections", "contentToString")

    val constructor = FunSpec.constructorBuilder()

    query.resultColumns.forEach {
      typeSpec.addProperty(
        PropertySpec.builder(it.name, it.javaType)
          .initializer(it.name)
          .build()
      )
      constructor.addParameter(it.name, it.javaType)

      propertyPrints += if (it.javaType.isArrayType) {
        CodeBlock.of("${it.name}: \${${it.name}.%M()}", contentToString)
      } else {
        CodeBlock.of("${it.name}: \$${it.name}")
      }
    }

    typeSpec.addFunction(
      FunSpec.builder("toString")
        .returns(String::class.asClassName())
        .addModifiers(OVERRIDE)
        .addStatement(
          "return %L",
          propertyPrints.joinToCode(
            separator = "\n|  ",
            prefix = "\"\"\"\n|${query.name.capitalize()} [\n|  ",
            suffix = "\n|]\n\"\"\".trimMargin()"
          )
        )
        .build()
    )

    return typeSpec
      .primaryConstructor(constructor.build())
      .build()
  }
}

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.core.compiler

import com.alecstrong.sql.psi.core.psi.SqlColumnDef
import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.squareup.kotlinpoet.ANY
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.KModifier.INNER
import com.squareup.kotlinpoet.KModifier.INTERNAL
import com.squareup.kotlinpoet.KModifier.OUT
import com.squareup.kotlinpoet.KModifier.OVERRIDE
import com.squareup.kotlinpoet.KModifier.PRIVATE
import com.squareup.kotlinpoet.LambdaTypeName
import com.squareup.kotlinpoet.MemberName
import com.squareup.kotlinpoet.NameAllocator
import com.squareup.kotlinpoet.ParameterSpec
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.kotlinpoet.TypeSpec
import com.squareup.kotlinpoet.TypeVariableName
import com.squareup.kotlinpoet.joinToCode
import com.squareup.sqldelight.core.compiler.SqlDelightCompiler.allocateName
import com.squareup.sqldelight.core.compiler.model.NamedQuery
import com.squareup.sqldelight.core.lang.ADAPTER_NAME
import com.squareup.sqldelight.core.lang.CURSOR_NAME
import com.squareup.sqldelight.core.lang.CURSOR_TYPE
import com.squareup.sqldelight.core.lang.CUSTOM_DATABASE_NAME
import com.squareup.sqldelight.core.lang.DRIVER_NAME
import com.squareup.sqldelight.core.lang.EXECUTE_METHOD
import com.squareup.sqldelight.core.lang.MAPPER_NAME
import com.squareup.sqldelight.core.lang.QUERY_LIST_TYPE
import com.squareup.sqldelight.core.lang.QUERY_TYPE
import com.squareup.sqldelight.core.lang.psi.ColumnTypeMixin
import com.squareup.sqldelight.core.lang.util.rawSqlText

class SelectQueryGenerator(private val query: NamedQuery) : QueryGenerator(query) {
  /**
   * The exposed query method which returns the default data class implementation.
   *
   * `fun selectForId(id: Int): Query<Data>`
   */
  fun defaultResultTypeFunction(): FunSpec {
    val function = defaultResultTypeFunctionInterface()
      .addModifiers(OVERRIDE)
    val argNameAllocator = NameAllocator()
    val params =
      query
        .arguments
        .asSequence()
        .sortedBy { it.index }
        .onEach { (_, argument) ->
          argNameAllocator.newName(argument.name, argument)
        }
        .map { (_, argument) ->
          CodeBlock.of(argNameAllocator[argument])
        }
        .toList()

    val columnArgs = query.resultColumns.map { argument ->
      argNameAllocator.newName(argument.name, argument)
    }

    val lamdaParams = columnArgs.joinToString(separator = ", ")
    val ctorParams = columnArgs.joinToString(separator = ",\n", postfix = "\n")

    val trailingLambda = CodeBlock.builder()
      .add(CodeBlock.of("·{ $lamdaParams ->\n"))
      .indent()
      .add("%T(\n", query.interfaceType)
      .indent()
      .add(ctorParams)
      .unindent()
      .add(")\n")
      .unindent()
      .add("}")
      .build()

    return function
      .addStatement(
        "return %L",
        CodeBlock
          .builder()
          .add(
            if (params.isEmpty()) {
              CodeBlock.of(query.name)
            } else {
              params.joinToCode(", ", "${query.name}(", ")")
            }
          )
          .add(trailingLambda)
          .build()
      )
      .build()
  }

  fun defaultResultTypeFunctionInterface(): FunSpec.Builder {
    val function = FunSpec.builder(query.name)
      .also(this::addJavadoc)
    query.arguments.sortedBy { it.index }.forEach { (_, argument) ->
      function.addParameter(argument.name, argument.argumentType())
    }
    return function
      .returns(QUERY_TYPE.parameterizedBy(query.interfaceType))
  }

  fun customResultTypeFunctionInterface(): FunSpec.Builder {
    val function = FunSpec.builder(query.name)
    val params = mutableListOf<CodeBlock>()

    query.arguments.sortedBy { it.index }.forEach { (_, argument) ->
      // Adds each sqlite parameter to the argument list:
      // fun <T> selectForId(<<id>>, <<other_param>>, ...)
      function.addParameter(argument.name, argument.argumentType())
      params.add(CodeBlock.of(argument.name))
    }

    if (query.needsWrapper()) {
      // Function takes a custom mapper.

      // Add the type variable to the signature.
      val typeVariable = TypeVariableName("T", ANY)
      function.addTypeVariable(typeVariable)

      // Add the custom mapper to the signature:
      // mapper: (id: kotlin.Long, value: kotlin.String) -> T
      function.addParameter(
        ParameterSpec.builder(
          MAPPER_NAME,
          LambdaTypeName.get(
            parameters = query.resultColumns.map {
              ParameterSpec.builder(it.name, it.javaType.copy(annotations = emptyList()))
                .build()
            },
            returnType = typeVariable
          )
        ).build()
      )

      // Specify the return type for the mapper:
      // Query<T>
      function.returns(QUERY_TYPE.parameterizedBy(typeVariable))
    } else {
      // No custom type possible, just returns the single column:
      // fun selectSomeText(_id): Query<String>
      function.returns(QUERY_TYPE.parameterizedBy(query.resultColumns.single().javaType))
    }

    return function
  }

  /**
   * The exposed query method which returns a provided custom type.
   *
   * `fun <T> selectForId(id, mapper: (column1: String) -> T): Query<T>`
   */
  fun customResultTypeFunction(): FunSpec {
    val function = customResultTypeFunctionInterface()
      .addModifiers(OVERRIDE)

    query.resultColumns.forEach { resultColumn ->
      (listOf(resultColumn) + resultColumn.assumedCompatibleTypes)
        .takeIf { it.size > 1 }
        ?.map { assumedCompatibleType ->
          (assumedCompatibleType.column?.columnType as ColumnTypeMixin?)?.let { columnTypeMixin ->
            val tableAdapterName = "${(assumedCompatibleType.column!!.parent as SqlCreateTableStmt).name()}$ADAPTER_NAME"
            val columnAdapterName = "${allocateName((columnTypeMixin.parent as SqlColumnDef).columnName)}$ADAPTER_NAME"
            "$CUSTOM_DATABASE_NAME.$tableAdapterName.$columnAdapterName"
          }
        }
        ?.let { adapterNames ->
          function.addStatement(
            """%M(%M(%L).size == 1) { "Adapter·types·are·expected·to·be·identical." }""",
            MemberName("kotlin", "check"),
            MemberName("kotlin.collections", "setOf"),
            adapterNames.joinToString()
          )
        }
    }

    // Assemble the actual mapper lambda:
    // { resultSet ->
    //   mapper(
    //       resultSet.getLong(0),
    //       queryWrapper.tableAdapter.columnAdapter.decode(resultSet.getString(0))
    //   )
    // }
    val mapperLambda = CodeBlock.builder().addStatement("·{ $CURSOR_NAME ->").indent()

    if (query.needsWrapper()) {
      mapperLambda.add("$MAPPER_NAME(\n")

      // Add the call of mapper with the deserialized columns:
      // mapper(
      //     resultSet.getLong(0),
      //     queryWrapper.tableAdapter.columnAdapter.decode(resultSet.getString(0))
      // )
      mapperLambda
        .indent()
        .apply {
          val decoders = query.resultColumns.mapIndexed { index, column -> column.cursorGetter(index) }
          add(decoders.joinToCode(separator = ",\n", suffix = "\n"))
        }
        .unindent()
        .add(")\n")
    } else {
      mapperLambda.add(query.resultColumns.single().cursorGetter(0)).add("\n")
    }
    mapperLambda.unindent().add("}\n")

    if (query.arguments.isEmpty()) {
      // No need for a custom query type, return an instance of Query:
      // return Query(statement, selectForId) { resultSet -> ... }
      function.addCode(
        "return %T(${query.id}, ${query.name}, $DRIVER_NAME, %S, %S, %S)%L",
        QUERY_TYPE, query.statement.containingFile.name, query.name, query.statement.rawSqlText(),
        mapperLambda.build()
      )
    } else {
      // Custom type is needed to handle dirtying events, return an instance of custom type:
      // return SelectForId(id) { resultSet -> ... }
      function.addCode(
        "return %N(%L)%L", query.customQuerySubtype,
        query.arguments.joinToString { (_, parameter) -> parameter.name }, mapperLambda.build()
      )
    }

    return function.build()
  }

  /**
   * The private property used to delegate query result updates.
   *
   * `private val selectForId: MutableList<Query<*>> = mutableListOf()`
   */
  fun queryCollectionProperty(): PropertySpec {
    return PropertySpec.builder(query.name, QUERY_LIST_TYPE, INTERNAL)
      .initializer("%M()", MemberName("com.squareup.sqldelight.internal", "copyOnWriteList"))
      .build()
  }

  /**
   * The private query subtype for this specific query.
   *
   * ```
   * private class SelectForIdQuery<out T>(
   *   private val _id: Int,
   *   mapper: (SqlResultSet) -> T
   * ) : Query<T>(statement, selectForId, mapper) {
   * private inner class SelectForIdQuery<out T>(
   *   private val _id: Int, mapper: (Cursor) -> T
   * ): Query<T>(database.helper, selectForId, mapper)
   * ```
   */
  fun querySubtype(): TypeSpec {
    val queryType = TypeSpec.classBuilder(query.customQuerySubtype)
      .addModifiers(PRIVATE, INNER)

    val constructor = FunSpec.constructorBuilder()

    // The custom return type variable:
    // <out T>
    val returnType = TypeVariableName("T", bounds = arrayOf(ANY), variance = OUT)
    queryType.addTypeVariable(returnType)

    // The superclass:
    // Query<T>
    queryType.superclass(QUERY_TYPE.parameterizedBy(returnType))

    val createStatementFunction = FunSpec.builder(EXECUTE_METHOD)
      .addModifiers(OVERRIDE)
      .returns(CURSOR_TYPE)
      .addCode(executeBlock())

    // For each bind argument the query has.
    query.arguments.sortedBy { it.index }.forEach { (_, parameter) ->
      // Add the argument as a constructor property. (Used later to figure out if query dirtied)
      // val id: Int
      queryType.addProperty(
        PropertySpec.builder(parameter.name, parameter.argumentType())
          .initializer(parameter.name)
          .build()
      )
      constructor.addParameter(parameter.name, parameter.argumentType())
    }

    // Add the query property to the super constructor
    queryType.addSuperclassConstructorParameter(query.name)

    // Add the mapper constructor parameter and pass to the super constructor
    constructor.addParameter(
      MAPPER_NAME,
      LambdaTypeName.get(
        parameters = arrayOf(CURSOR_TYPE),
        returnType = returnType
      )
    )
    queryType.addSuperclassConstructorParameter(MAPPER_NAME)

    return queryType
      .primaryConstructor(constructor.build())
      .addFunction(createStatementFunction.build())
      .addFunction(
        FunSpec.builder("toString")
          .addModifiers(OVERRIDE)
          .returns(String::class)
          .addStatement("return %S", "${query.statement.containingFile.name}:${query.name}")
          .build()
      )
      .build()
  }
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.core.compiler

import com.alecstrong.sql.psi.core.psi.NamedElement
import com.alecstrong.sql.psi.core.psi.SqlCreateViewStmt
import com.intellij.openapi.module.Module
import com.intellij.psi.PsiElement
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.FileSpec
import com.squareup.kotlinpoet.NameAllocator
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.compiler.model.NamedQuery
import com.squareup.sqldelight.core.lang.MigrationFile
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.SqlDelightQueriesFile
import com.squareup.sqldelight.core.lang.queriesName
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.lang.util.sqFile
import java.io.Closeable

private typealias FileAppender = (fileName: String) -> Appendable

object SqlDelightCompiler {
  fun writeInterfaces(
    module: Module,
    file: SqlDelightQueriesFile,
    output: FileAppender
  ) {
    writeTableInterfaces(file, output)
    writeQueryInterfaces(file, output)
    writeQueriesInterface(module, file, output)
  }

  fun writeInterfaces(
    file: MigrationFile,
    output: FileAppender,
    includeAll: Boolean = false
  ) {
    writeTableInterfaces(file, output, includeAll)
  }

  fun writeDatabaseInterface(
    module: Module,
    file: SqlDelightFile,
    implementationFolder: String,
    output: FileAppender
  ) {
    writeQueryWrapperInterface(module, file, implementationFolder, output)
  }

  fun writeImplementations(
    module: Module,
    sourceFile: SqlDelightQueriesFile,
    implementationFolder: String,
    output: FileAppender
  ) {
    val fileIndex = SqlDelightFileIndex.getInstance(module)
    val packageName = "${fileIndex.packageName}.$implementationFolder"
    val databaseImplementationType = DatabaseGenerator(module, sourceFile).type(packageName)
    val exposer = DatabaseExposerGenerator(databaseImplementationType, fileIndex)

    val fileSpec = FileSpec.builder(packageName, databaseImplementationType.name!!)
      .addProperty(exposer.exposedSchema())
      .addFunction(exposer.exposedConstructor())
      .addType(databaseImplementationType)
      .apply {
        fileIndex.sourceFolders(sourceFile, includeDependencies = true)
          .flatMap { it.findChildrenOfType<SqlDelightQueriesFile>() }
          .forEach { file ->
            val queriesGenerator = QueriesTypeGenerator(module, file)
            addType(queriesGenerator.generateType(packageName))
          }
      }
      .build()

    fileIndex.outputDirectory(sourceFile).forEach { outputDirectory ->
      val packageDirectory = "$outputDirectory/${packageName.replace(".", "/")}"
      fileSpec.writeToAndClose(output("$packageDirectory/${databaseImplementationType.name}.kt"))
    }
  }

  private fun writeQueryWrapperInterface(
    module: Module,
    sourceFile: SqlDelightFile,
    implementationFolder: String,
    output: FileAppender
  ) {
    val fileIndex = SqlDelightFileIndex.getInstance(module)
    val packageName = fileIndex.packageName
    val queryWrapperType = DatabaseGenerator(module, sourceFile).interfaceType()
    val fileSpec = FileSpec.builder(packageName, queryWrapperType.name!!)
      // TODO: Remove these when kotlinpoet supports top level types.
      .addImport("$packageName.$implementationFolder", "newInstance", "schema")
      .apply {
        var index = 0
        fileIndex.dependencies.forEach {
          addAliasedImport(ClassName(it.packageName, it.className), "${it.className}${index++}")
        }
      }
      .addType(queryWrapperType)
      .build()

    fileIndex.outputDirectory(sourceFile).forEach { outputDirectory ->
      val packageDirectory = "$outputDirectory/${packageName.replace(".", "/")}"
      fileSpec.writeToAndClose(output("$packageDirectory/${queryWrapperType.name}.kt"))
    }
  }

  internal fun writeTableInterfaces(
    file: SqlDelightFile,
    output: FileAppender,
    includeAll: Boolean = false
  ) {
    val packageName = file.packageName ?: return
    file.tables(includeAll).forEach { query ->
      val statement = query.tableName.parent

      if (statement is SqlCreateViewStmt && statement.compoundSelectStmt != null) {
        listOf(NamedQuery(allocateName(statement.viewName), statement.compoundSelectStmt!!, statement.viewName))
          .writeQueryInterfaces(file, output)
        return@forEach
      }

      val fileSpec = FileSpec.builder(packageName, allocateName(query.tableName))
        .apply {
          tryWithElement(statement) {
            val generator = TableInterfaceGenerator(query)
            addType(generator.kotlinImplementationSpec())
          }
        }
        .build()

      statement.sqFile().generatedDirectories?.forEach { directory ->
        fileSpec.writeToAndClose(output("$directory/${allocateName(query.tableName).capitalize()}.kt"))
      }
    }
  }

  internal fun writeQueryInterfaces(
    file: SqlDelightQueriesFile,
    output: FileAppender
  ) {
    file.namedQueries.writeQueryInterfaces(file, output)
  }

  internal fun writeQueriesInterface(
    module: Module,
    file: SqlDelightQueriesFile,
    output: FileAppender
  ) {
    val packageName = file.packageName ?: return
    val queriesType = QueriesTypeGenerator(module, file).interfaceType()
    val fileSpec = FileSpec.builder(packageName, file.queriesName.capitalize())
      .addType(queriesType)
      .build()

    file.generatedDirectories?.forEach { directory ->
      fileSpec.writeToAndClose(output("$directory/${queriesType.name}.kt"))
    }
  }

  internal fun allocateName(namedElement: NamedElement): String {
    return NameAllocator().newName(namedElement.normalizedName)
  }

  private fun List<NamedQuery>.writeQueryInterfaces(file: SqlDelightFile, output: FileAppender) {
    return filter { tryWithElement(it.select) { it.needsInterface() } }
      .forEach { namedQuery ->
        val fileSpec = FileSpec.builder(namedQuery.interfaceType.packageName, namedQuery.name)
          .apply {
            tryWithElement(namedQuery.select) {
              val generator = QueryInterfaceGenerator(namedQuery)
              addType(generator.kotlinImplementationSpec())
            }
          }
          .build()

        file.generatedDirectories(namedQuery.interfaceType.packageName)?.forEach { directory ->
          fileSpec.writeToAndClose(output("$directory/${namedQuery.name.capitalize()}.kt"))
        }
      }
  }

  private val NamedElement.normalizedName: String
    get() {
      val f = name[0]
      val l = name[name.lastIndex]
      return if ((f in "\"'`" && f == l) || (f == '[' && l == ']')) {
        name.substring(1, name.length - 1)
      } else {
        name
      }
    }

  private fun FileSpec.writeToAndClose(appendable: Appendable) {
    writeTo(appendable)
    if (appendable is Closeable) appendable.close()
  }
}

internal fun <T> tryWithElement(
  element: PsiElement,
  block: () -> T
): T {
  try {
    return block()
  } catch (e: Throwable) {
    val exception = IllegalStateException("Failed to compile $element :\n${element.text}")
    exception.initCause(e)
    throw exception
  }
}

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.core.compiler

import com.alecstrong.sql.psi.core.psi.LazyQuery
import com.alecstrong.sql.psi.core.psi.SqlColumnDef
import com.alecstrong.sql.psi.core.psi.SqlStmt
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.intellij.psi.util.PsiTreeUtil
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.KModifier.DATA
import com.squareup.kotlinpoet.KModifier.OVERRIDE
import com.squareup.kotlinpoet.MemberName
import com.squareup.kotlinpoet.ParameterSpec
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.kotlinpoet.TypeSpec
import com.squareup.kotlinpoet.asClassName
import com.squareup.kotlinpoet.joinToCode
import com.squareup.sqldelight.core.compiler.SqlDelightCompiler.allocateName
import com.squareup.sqldelight.core.compiler.integration.javadocText
import com.squareup.sqldelight.core.lang.ADAPTER_NAME
import com.squareup.sqldelight.core.lang.psi.ColumnTypeMixin
import com.squareup.sqldelight.core.lang.psi.ColumnTypeMixin.Companion.isArrayType
import com.squareup.sqldelight.core.lang.util.childOfType
import com.squareup.sqldelight.core.lang.util.parentOfType
import com.squareup.sqldelight.core.psi.SqlDelightStmtIdentifier

internal class TableInterfaceGenerator(private val table: LazyQuery) {
  private val typeName = allocateName(table.tableName).capitalize()

  fun kotlinImplementationSpec(): TypeSpec {
    val typeSpec = TypeSpec.classBuilder(typeName)
      .addModifiers(DATA)

    val identifier = PsiTreeUtil.getPrevSiblingOfType(
      PsiTreeUtil.getParentOfType(table.tableName, SqlStmt::class.java),
      SqlDelightStmtIdentifier::class.java
    )
    identifier?.childOfType(SqlTypes.JAVADOC)?.let { javadoc ->
      javadocText(javadoc)?.let { typeSpec.addKdoc(it) }
    }

    val propertyPrints = mutableListOf<CodeBlock>()
    val contentToString = MemberName("kotlin.collections", "contentToString")

    val constructor = FunSpec.constructorBuilder()

    table.columns.forEach { column ->
      val columnName = allocateName(column.columnName)
      val columnType = column.columnType as ColumnTypeMixin
      typeSpec.addProperty(
        PropertySpec.builder(columnName, columnType.type().javaType)
          .initializer(columnName)
          .build()
      )
      val param = ParameterSpec.builder(columnName, columnType.type().javaType)
      column.javadoc?.let(::javadocText)?.let { param.addKdoc(it) }
      constructor.addParameter(param.build())

      propertyPrints += if (columnType.type().javaType.isArrayType) {
        CodeBlock.of("$columnName: \${$columnName.%M()}", contentToString)
      } else {
        CodeBlock.of("$columnName: \$$columnName")
      }
    }

    typeSpec.addFunction(
      FunSpec.builder("toString")
        .returns(String::class.asClassName())
        .addModifiers(OVERRIDE)
        .addStatement(
          "return %L",
          propertyPrints.joinToCode(
            separator = "\n|  ",
            prefix = "\"\"\"\n|$typeName [\n|  ",
            suffix = "\n|]\n\"\"\".trimMargin()"
          )
        )
        .build()
    )

    val adapters = table.columns.mapNotNull { (it.columnType as ColumnTypeMixin).adapter() }

    if (adapters.isNotEmpty()) {
      typeSpec.addType(
        TypeSpec.classBuilder(ADAPTER_NAME)
          .primaryConstructor(
            FunSpec.constructorBuilder()
              .addParameters(
                adapters.map {
                  ParameterSpec.builder(it.name, it.type, *it.modifiers.toTypedArray()).build()
                }
              )
              .build()
          )
          .addProperties(
            adapters.map {
              PropertySpec.builder(it.name, it.type, *it.modifiers.toTypedArray())
                .initializer(it.name)
                .build()
            }
          )
          .build()
      )
    }

    return typeSpec
      .primaryConstructor(constructor.build())
      .build()
  }

  private val LazyQuery.columns: Collection<SqlColumnDef>
    get() = query.columns.map { it.element.parentOfType() }
}

package com.squareup.sqldelight.core.lang

import com.intellij.openapi.vfs.VirtualFile
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import com.squareup.kotlinpoet.STAR

internal val CURSOR_TYPE = ClassName("com.squareup.sqldelight.db", "SqlCursor")
internal const val CURSOR_NAME = "cursor"

internal val DRIVER_TYPE = ClassName("com.squareup.sqldelight.db", "SqlDriver")
internal const val DRIVER_NAME = "driver"
internal val DATABASE_SCHEMA_TYPE = DRIVER_TYPE.nestedClass("Schema")

internal const val CUSTOM_DATABASE_NAME = "database"

internal const val ADAPTER_NAME = "Adapter"

internal val QUERY_TYPE = ClassName("com.squareup.sqldelight", "Query")

internal val QUERY_LIST_TYPE = ClassName("kotlin.collections", "MutableList")
  .parameterizedBy(QUERY_TYPE.parameterizedBy(STAR))

internal const val MAPPER_NAME = "mapper"

internal const val EXECUTE_METHOD = "execute"

const val QUERIES_SUFFIX_NAME = "Queries"

val VirtualFile.queriesName
  get() = "${nameWithoutExtension.capitalize()}$QUERIES_SUFFIX_NAME"

internal val SqlDelightFile.queriesName
  get() = "${virtualFile!!.nameWithoutExtension.decapitalize()}$QUERIES_SUFFIX_NAME"
internal val SqlDelightFile.queriesType
  get() = ClassName(packageName!!, "${virtualFile!!.nameWithoutExtension.capitalize()}$QUERIES_SUFFIX_NAME")
internal fun SqlDelightFile.queriesImplType(implementationPackage: String) =
  ClassName(implementationPackage, "${virtualFile!!.nameWithoutExtension.capitalize()}QueriesImpl")

internal val TRANSACTER_TYPE = ClassName("com.squareup.sqldelight", "Transacter")
internal val TRANSACTER_IMPL_TYPE = ClassName("com.squareup.sqldelight", "TransacterImpl")

package com.squareup.sqldelight.core.lang

import com.intellij.icons.AllIcons
import com.intellij.openapi.fileTypes.FileType
import com.intellij.openapi.vfs.VirtualFile

object DatabaseFileType : FileType {
  const val EXTENSION = "db"

  override fun getName() = "SqlDelight Database"
  override fun getDescription() = "SqlDelight Database"
  override fun getDefaultExtension() = EXTENSION
  override fun getIcon() = AllIcons.Nodes.DataSchema
  override fun isBinary() = true
  override fun isReadOnly() = true
  override fun getCharset(vFile: VirtualFile, bytes: ByteArray) = null
}

package com.squareup.sqldelight.core.lang

import com.alecstrong.sql.psi.core.DialectPreset
import com.intellij.lang.Language
import com.intellij.openapi.roots.ProjectFileIndex
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.FileViewProvider
import com.intellij.psi.FileViewProviderFactory
import com.intellij.psi.PsiManager
import com.intellij.psi.SingleRootFileViewProvider
import com.intellij.testFramework.LightVirtualFile
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.SqlDelightProjectService
import java.sql.Connection
import java.sql.DriverManager
import java.sql.SQLException

class DatabaseFileViewProviderFactory : FileViewProviderFactory {
  override fun createFileViewProvider(
    vFile: VirtualFile,
    language: Language?,
    psiManager: PsiManager,
    eventSystemEnabled: Boolean
  ): FileViewProvider {
    fun default() = SingleRootFileViewProvider(psiManager, vFile, eventSystemEnabled)

    if (vFile.fileType != DatabaseFileType) return default()

    val projectService = SqlDelightProjectService.getInstance(psiManager.project)
    if (!projectService.isSupportedSqlite()) return default()

    val module = projectService.module(vFile) ?: return default()
    val index = SqlDelightFileIndex.getInstance(module)
    if (index.deriveSchemaFromMigrations) return default()

    // If this database is unversioned, continue.
    val version = vFile.version ?: return default()

    val lowestVersion = ProjectFileIndex.getInstance(psiManager.project).iterateContent {
      if (it.fileType != DatabaseFileType && it.fileType != MigrationFileType) return@iterateContent true
      if (module != projectService.module(vFile)) return@iterateContent true

      val otherVersion = it.version
      if (otherVersion != null && otherVersion < version) {
        return@iterateContent false
      }
      return@iterateContent true
    }

    // Only parse a db if it is the base.
    if (!lowestVersion) return default()

    return DatabaseFileViewProvider(psiManager, vFile, eventSystemEnabled)
  }
}

internal class DatabaseFileViewProvider(
  manager: PsiManager,
  virtualFile: VirtualFile,
  eventSystemEnabled: Boolean
) : SingleRootFileViewProvider(manager, virtualFile, eventSystemEnabled, DatabaseFileType) {
  private var schemaFile: VirtualFile? = null

  fun getSchemaFile(): VirtualFile? {
    schemaFile?.let { return it }

    val virtualFile = super.getVirtualFile()
    try {
      val statements = createConnection(virtualFile.path).use {
        it.prepareStatement("SELECT sql FROM sqlite_master WHERE sql IS NOT NULL;").use {
          it.executeQuery().use {
            mutableListOf<String>().apply {
              while (it.next()) {
                add("${it.getString(1)};")
              }
            }
          }
        }
      }
      return LightVirtualFile(
        "${virtualFile.version!! - 1}.db",
        MigrationFileType,
        statements.joinToString(separator = "\n")
      ).also { schemaFile = it }
    } catch (e: Exception) {
      return null
    }
  }

  private fun createConnection(path: String): Connection {
    return try {
      DriverManager.getConnection("jdbc:sqlite:$path")
    } catch (e: SQLException) {
      DriverManager.getConnection("jdbc:sqlite:$path")
    }
  }
}

private val VirtualFile.version
  get() = nameWithoutExtension.filter { it in '0'..'9' }.toIntOrNull()

private fun SqlDelightProjectService.isSupportedSqlite(): Boolean = when (dialectPreset) {
  DialectPreset.SQLITE_3_18, DialectPreset.SQLITE_3_24, DialectPreset.SQLITE_3_25 -> true
  DialectPreset.MYSQL, DialectPreset.POSTGRESQL, DialectPreset.HSQL -> false
}

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.core.lang

import com.alecstrong.sql.psi.core.psi.Queryable
import com.alecstrong.sql.psi.core.psi.SqlBindExpr
import com.alecstrong.sql.psi.core.psi.SqlColumnDef
import com.intellij.psi.util.PsiTreeUtil
import com.squareup.kotlinpoet.BOOLEAN
import com.squareup.kotlinpoet.BYTE
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.FLOAT
import com.squareup.kotlinpoet.INT
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import com.squareup.kotlinpoet.SHORT
import com.squareup.kotlinpoet.TypeName
import com.squareup.kotlinpoet.asClassName
import com.squareup.sqldelight.core.compiler.integration.adapterName
import com.squareup.sqldelight.core.dialect.api.DialectType
import com.squareup.sqldelight.core.lang.psi.ColumnTypeMixin
import com.squareup.sqldelight.core.lang.util.isArrayParameter

/**
 * Internal representation for a column type, which has SQLite data affinity as well as JVM class
 * type.
 */
internal data class IntermediateType(
  val dialectType: DialectType,
  val javaType: TypeName = dialectType.javaType,
  /**
   * The column definition this type is sourced from, or null if there is none.
   */
  val column: SqlColumnDef? = null,
  /**
   * The name of this intermediate type as exposed in the generated api.
   */
  val name: String = "value",
  /**
   * The original bind argument expression this intermediate type comes from.
   */
  val bindArg: SqlBindExpr? = null,
  /**
   * Whether or not this argument is extracted from a different type
   */
  val extracted: Boolean = false,
  /**
   * The types assumed to be compatible with this type. Validated at runtime.
   */
  val assumedCompatibleTypes: List<IntermediateType> = emptyList(),
) {
  fun asNullable() = copy(javaType = javaType.copy(nullable = true))

  fun asNonNullable() = copy(javaType = javaType.copy(nullable = false))

  fun nullableIf(predicate: Boolean): IntermediateType {
    return if (predicate) asNullable() else asNonNullable()
  }

  fun argumentType() = if (bindArg?.isArrayParameter() == true) {
    Collection::class.asClassName().parameterizedBy(javaType)
  } else {
    javaType
  }

  /**
   * @return A [CodeBlock] which binds this type to [columnIndex] on [STATEMENT_NAME].
   *
   * eg: statement.bindBytes(0, queryWrapper.tableNameAdapter.columnNameAdapter.encode(column))
   */
  fun preparedStatementBinder(
    columnIndex: String,
    extractedVariable: String? = null
  ): CodeBlock {
    val codeBlock = extractedVariable?.let(CodeBlock::of) ?: encodedJavaType()
    if (codeBlock != null) {
      return dialectType.prepareStatementBinder(columnIndex, codeBlock)
    }

    val name = if (javaType.isNullable) "it" else this.name
    val value = when (javaType.copy(nullable = false)) {
      FLOAT -> CodeBlock.of("$name.toDouble()")
      BYTE -> CodeBlock.of("$name.toLong()")
      SHORT -> CodeBlock.of("$name.toLong()")
      INT -> CodeBlock.of("$name.toLong()")
      BOOLEAN -> CodeBlock.of("if ($name) 1L else 0L")
      else -> {
        return dialectType.prepareStatementBinder(
          columnIndex,
          dialectType.encode(CodeBlock.of(this.name))
        )
      }
    }

    return if (javaType.isNullable) {
      dialectType.prepareStatementBinder(columnIndex, value.wrapInLet())
    } else {
      dialectType.prepareStatementBinder(columnIndex, value)
    }
  }

  fun encodedJavaType(): CodeBlock? {
    val name = if (javaType.isNullable) "it" else this.name
    return (column?.columnType as ColumnTypeMixin?)?.adapter()?.let { adapter ->
      val parent = PsiTreeUtil.getParentOfType(column, Queryable::class.java)
      val adapterName = parent!!.tableExposed().adapterName
      val value = dialectType.encode(
        CodeBlock.of("$CUSTOM_DATABASE_NAME.$adapterName.%N.encode($name)", adapter)
      )
      if (javaType.isNullable) {
        value.wrapInLet()
      } else {
        value
      }
    }
  }

  private fun CodeBlock.wrapInLet(): CodeBlock {
    return CodeBlock.builder()
      .add("${this@IntermediateType.name}?.let { ")
      .add(this)
      .add(" }")
      .build()
  }

  fun cursorGetter(columnIndex: Int): CodeBlock {
    var cursorGetter = dialectType.cursorGetter(columnIndex)

    if (!javaType.isNullable) {
      cursorGetter = CodeBlock.of("$cursorGetter!!")
    }

    return (column?.columnType as ColumnTypeMixin?)?.adapter()?.let { adapter ->
      val adapterName = PsiTreeUtil.getParentOfType(column, Queryable::class.java)!!.tableExposed().adapterName
      if (javaType.isNullable) {
        CodeBlock.of("%L?.let { $CUSTOM_DATABASE_NAME.$adapterName.%N.decode(%L) }", cursorGetter, adapter, dialectType.decode(CodeBlock.of("it")))
      } else {
        CodeBlock.of("$CUSTOM_DATABASE_NAME.$adapterName.%N.decode(%L)", adapter, dialectType.decode(cursorGetter))
      }
    } ?: when (javaType) {
      FLOAT -> CodeBlock.of("$cursorGetter.toFloat()")
      FLOAT.copy(nullable = true) -> CodeBlock.of("$cursorGetter?.toFloat()")
      BYTE -> CodeBlock.of("$cursorGetter.toByte()")
      BYTE.copy(nullable = true) -> CodeBlock.of("$cursorGetter?.toByte()")
      SHORT -> CodeBlock.of("$cursorGetter.toShort()")
      SHORT.copy(nullable = true) -> CodeBlock.of("$cursorGetter?.toShort()")
      INT -> CodeBlock.of("$cursorGetter.toInt()")
      INT.copy(nullable = true) -> CodeBlock.of("$cursorGetter?.toInt()")
      BOOLEAN -> CodeBlock.of("$cursorGetter == 1L")
      BOOLEAN.copy(nullable = true) -> CodeBlock.of("$cursorGetter?.let { it == 1L }")
      else -> cursorGetter
    }
  }
}

package com.squareup.sqldelight.core.lang

import com.alecstrong.sql.psi.core.SqlFileBase
import com.intellij.openapi.vfs.VfsUtilCore
import com.intellij.psi.FileViewProvider
import com.intellij.psi.PsiManager
import com.squareup.sqldelight.core.SqlDelightFileIndex

class MigrationFile(
  viewProvider: FileViewProvider
) : SqlDelightFile(viewProvider, MigrationLanguage) {
  val version: Int by lazy {
    name.substringBeforeLast(".${fileType.EXTENSION}")
      .filter { it in '0'..'9' }.toIntOrNull() ?: 0
  }

  internal fun sqliteStatements() = sqlStmtList!!.stmtList

  override val packageName
    get() = module?.let { module -> SqlDelightFileIndex.getInstance(module).packageName }

  override val order
    get() = version

  override fun getFileType() = MigrationFileType

  override fun baseContributorFile(): SqlFileBase? {
    val module = module
    if (module == null || SqlDelightFileIndex.getInstance(
        module
      ).deriveSchemaFromMigrations
    ) return null

    val manager = PsiManager.getInstance(project)
    var result: SqlFileBase? = null
    val folders = SqlDelightFileIndex.getInstance(module).sourceFolders(virtualFile ?: return null)
    folders.forEach { dir ->
      VfsUtilCore.iterateChildrenRecursively(
        dir, { it.isDirectory || it.fileType == DatabaseFileType },
        { file ->
          if (file.isDirectory) return@iterateChildrenRecursively true

          val vFile = (manager.findViewProvider(file) as? DatabaseFileViewProvider)?.getSchemaFile()
            ?: return@iterateChildrenRecursively true

          manager.findFile(vFile)?.let { psiFile ->
            if (psiFile is SqlFileBase) {
              result = psiFile
              return@iterateChildrenRecursively false
            }
          }

          return@iterateChildrenRecursively true
        }
      )
    }
    return result
  }
}

package com.squareup.sqldelight.core.lang

import com.intellij.icons.AllIcons
import com.intellij.openapi.fileTypes.LanguageFileType

object MigrationFileType : LanguageFileType(MigrationLanguage) {
  private val ICON = AllIcons.Providers.Sqlite

  const val EXTENSION = "sqm"

  override fun getName() = "SqlDelight Migration"
  override fun getDescription() = "SqlDelight Migration"
  override fun getDefaultExtension() = EXTENSION
  override fun getIcon() = ICON
}

package com.squareup.sqldelight.core.lang

import com.intellij.lang.Language

object MigrationLanguage : Language("SqlDelightMigration")

package com.squareup.sqldelight.core.lang

import com.alecstrong.sql.psi.core.SqlFileBase
import com.alecstrong.sql.psi.core.SqlParser
import com.alecstrong.sql.psi.core.SqlParserDefinition
import com.intellij.openapi.project.Project
import com.intellij.psi.FileViewProvider
import com.intellij.psi.stubs.PsiFileStub
import com.intellij.psi.tree.ILightStubFileElementType

class MigrationParserDefinition : SqlParserDefinition() {
  private val parserUtil = ParserUtil()

  override fun createFile(viewProvider: FileViewProvider) = MigrationFile(viewProvider)
  override fun getFileNodeType() = FILE
  override fun getLanguage() = SqlDelightLanguage

  override fun createParser(project: Project): SqlParser {
    parserUtil.initializeDialect(project)
    return super.createParser(project)
  }

  companion object {
    private val FILE = ILightStubFileElementType<PsiFileStub<SqlFileBase>>(MigrationLanguage)
  }
}

package com.squareup.sqldelight.core.lang

import com.alecstrong.sql.psi.core.DialectPreset
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.intellij.openapi.project.Project
import com.squareup.sqldelight.core.SqlDelightProjectService
import com.squareup.sqldelight.core.SqldelightParserUtil
import com.squareup.sqldelight.core.lang.psi.FunctionExprMixin

internal class ParserUtil {
  private var dialect: DialectPreset? = null

  fun initializeDialect(project: Project) {
    val newDialect = SqlDelightProjectService.getInstance(project).dialectPreset
    if (newDialect != dialect) {
      newDialect.setup()
      SqldelightParserUtil.overrideSqlParser()
      dialect = newDialect

      val currentElementCreation = SqldelightParserUtil.createElement
      SqldelightParserUtil.createElement = {
        when (it.elementType) {
          SqlTypes.FUNCTION_EXPR -> FunctionExprMixin(it)
          else -> currentElementCreation(it)
        }
      }
    }
  }
}

package com.squareup.sqldelight.core.lang

import com.alecstrong.sql.psi.core.SqlFileBase
import com.intellij.lang.Language
import com.intellij.openapi.module.Module
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.FileViewProvider
import com.intellij.psi.search.GlobalSearchScope
import com.intellij.psi.search.GlobalSearchScopesCore
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.SqlDelightProjectService

abstract class SqlDelightFile(
  viewProvider: FileViewProvider,
  language: Language
) : SqlFileBase(viewProvider, language) {
  protected val module: Module?
    get() = SqlDelightProjectService.getInstance(project).module(requireNotNull(virtualFile, { "Null virtualFile" }))

  val generatedDirectories by lazy {
    val packageName = packageName ?: return@lazy null
    generatedDirectories(packageName)
  }

  internal fun generatedDirectories(packageName: String): List<String>? {
    val module = module ?: return null
    return SqlDelightFileIndex.getInstance(module).outputDirectory(this).map { outputDirectory ->
      "$outputDirectory/${packageName.replace('.', '/')}"
    }
  }

  internal val dialect
    get() = SqlDelightProjectService.getInstance(project).dialectPreset

  internal abstract val packageName: String?

  override fun getVirtualFile(): VirtualFile? {
    if (myOriginalFile != null) return myOriginalFile.virtualFile
    return super.getVirtualFile()
  }

  override fun searchScope(): GlobalSearchScope {
    val default = GlobalSearchScope.fileScope(this)

    val module = module ?: return default
    val index = SqlDelightFileIndex.getInstance(module)
    val sourceFolders = index.sourceFolders(virtualFile ?: return default)
    if (sourceFolders.isEmpty()) return default

    // TODO Deal with database files?

    return sourceFolders
      .map { GlobalSearchScopesCore.directoryScope(project, it, true) }
      .reduce { totalScope, directoryScope -> totalScope.union(directoryScope) }
  }
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.core.lang

import com.intellij.icons.AllIcons
import com.intellij.openapi.fileTypes.LanguageFileType

object SqlDelightFileType : LanguageFileType(SqlDelightLanguage) {
  private val ICON = AllIcons.Providers.Sqlite

  const val EXTENSION = "sq"
  const val FOLDER_NAME = "sqldelight"

  override fun getName() = "SqlDelight"
  override fun getDescription() = "SqlDelight"
  override fun getDefaultExtension() = EXTENSION
  override fun getIcon() = ICON
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.core.lang

import com.intellij.lang.Language

object SqlDelightLanguage : Language("SqlDelight")

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.core.lang

import com.alecstrong.sql.psi.core.SqlFileBase
import com.alecstrong.sql.psi.core.SqlParser
import com.alecstrong.sql.psi.core.SqlParserDefinition
import com.intellij.openapi.project.Project
import com.intellij.psi.FileViewProvider
import com.intellij.psi.stubs.PsiFileStub
import com.intellij.psi.tree.ILightStubFileElementType

class SqlDelightParserDefinition : SqlParserDefinition() {
  private val parserUtil = ParserUtil()

  override fun createFile(viewProvider: FileViewProvider) = SqlDelightQueriesFile(viewProvider)
  override fun getFileNodeType() = FILE
  override fun getLanguage() = SqlDelightLanguage

  override fun createParser(project: Project): SqlParser {
    parserUtil.initializeDialect(project)
    return super.createParser(project)
  }

  companion object {
    private val FILE = ILightStubFileElementType<PsiFileStub<SqlFileBase>>(SqlDelightLanguage)
  }
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.core.lang

import com.alecstrong.sql.psi.core.SqlAnnotationHolder
import com.alecstrong.sql.psi.core.psi.SqlAnnotatedElement
import com.alecstrong.sql.psi.core.psi.SqlStmt
import com.intellij.psi.FileViewProvider
import com.intellij.psi.search.GlobalSearchScope
import com.intellij.psi.util.PsiTreeUtil
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.compiler.model.NamedExecute
import com.squareup.sqldelight.core.compiler.model.NamedMutator.Delete
import com.squareup.sqldelight.core.compiler.model.NamedMutator.Insert
import com.squareup.sqldelight.core.compiler.model.NamedMutator.Update
import com.squareup.sqldelight.core.compiler.model.NamedQuery
import com.squareup.sqldelight.core.lang.psi.StmtIdentifierMixin
import com.squareup.sqldelight.core.psi.SqlDelightStmtList

class SqlDelightQueriesFile(
  viewProvider: FileViewProvider
) : SqlDelightFile(viewProvider, SqlDelightLanguage),
  SqlAnnotatedElement {
  override val packageName by lazy {
    module?.let { module ->
      SqlDelightFileIndex.getInstance(module).packageName(this)
    }
  }

  internal val namedQueries by lazy {
    sqliteStatements()
      .filter { it.statement.compoundSelectStmt != null && it.identifier.name != null }
      .map { NamedQuery(it.identifier.name!!, it.statement.compoundSelectStmt!!, it.identifier) }
  }

  internal val namedMutators by lazy {
    sqliteStatements().filter { it.identifier.name != null }
      .mapNotNull {
        when {
          it.statement.deleteStmtLimited != null -> Delete(it.statement.deleteStmtLimited!!, it.identifier)
          it.statement.insertStmt != null -> Insert(it.statement.insertStmt!!, it.identifier)
          it.statement.updateStmtLimited != null -> Update(it.statement.updateStmtLimited!!, it.identifier)
          else -> null
        }
      }
  }

  internal val namedExecutes by lazy {
    val sqlStmtList = PsiTreeUtil.getChildOfType(this, SqlDelightStmtList::class.java)!!

    val transactions = sqlStmtList.stmtClojureList.map {
      NamedExecute(
        identifier = it.stmtIdentifierClojure as StmtIdentifierMixin,
        statement = it.stmtClojureStmtList!!
      )
    }

    val statements = sqliteStatements()
      .filter {
        it.identifier.name != null &&
          it.statement.deleteStmtLimited == null &&
          it.statement.insertStmt == null &&
          it.statement.updateStmtLimited == null &&
          it.statement.compoundSelectStmt == null
      }
      .map { NamedExecute(it.identifier, it.statement) }

    return@lazy transactions + statements
  }

  internal val triggers by lazy { triggers(this) }

  override val order = null

  override fun getFileType() = SqlDelightFileType

  internal fun sqliteStatements(): Collection<LabeledStatement> {
    val sqlStmtList = PsiTreeUtil.getChildOfType(this, SqlDelightStmtList::class.java)!!
    return sqlStmtList.stmtIdentifierList.zip(sqlStmtList.stmtList) { id, stmt ->
      return@zip LabeledStatement(id as StmtIdentifierMixin, stmt)
    }
  }

  fun iterateSqlFiles(block: (SqlDelightQueriesFile) -> Unit) {
    val module = module ?: return

    SqlDelightFileIndex.getInstance(module).sourceFolders(this).forEach { dir ->
      PsiTreeUtil.findChildrenOfAnyType(dir, SqlDelightQueriesFile::class.java).forEach(block)
    }
  }

  override fun annotate(annotationHolder: SqlAnnotationHolder) {
    if (packageName.isNullOrEmpty()) {
      annotationHolder.createErrorAnnotation(this, "SqlDelight files must be placed in a package directory.")
    }
  }

  override fun searchScope(): GlobalSearchScope {
    val module = module
    if (module != null && !SqlDelightFileIndex.getInstance(module).deriveSchemaFromMigrations) {
      return GlobalSearchScope.getScopeRestrictedByFileTypes(super.searchScope(), SqlDelightFileType)
    }
    return super.searchScope()
  }

  data class LabeledStatement(val identifier: StmtIdentifierMixin, val statement: SqlStmt)
}

package com.squareup.sqldelight.core.lang

import com.alecstrong.sql.psi.core.psi.SqlInsertStmt
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.squareup.sqldelight.core.lang.util.childOfType

fun SqlInsertStmt.acceptsTableInterface(): Boolean {
  return insertStmtValues?.childOfType(SqlTypes.BIND_EXPR) != null
}

package com.squareup.sqldelight.gradle.android

import com.android.build.gradle.BaseExtension
import com.android.ide.common.symbols.getPackageNameFromManifest
import org.gradle.api.Project

internal fun Project.packageName(): String {
  val androidExtension = extensions.getByType(BaseExtension::class.java)
  androidExtension.sourceSets
    .map { it.manifest.srcFile }
    .filter { it.exists() }
    .forEach {
      return getPackageNameFromManifest(it)
    }
  throw IllegalStateException("No source sets available")
}

internal fun Project.sqliteVersion(): String? {
  val androidExtension = extensions.getByType(BaseExtension::class.java)
  val minSdk = androidExtension.defaultConfig.minSdk ?: return null
  if (minSdk >= 30) return "sqlite:3.25"
  return "sqlite:3.18"
}

package com.squareup.sqldelight.gradle.kotlin

import org.gradle.api.Project
import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinNativeCompilation

fun Project.linkSqlite() {
  val extension = project.extensions.findByType(KotlinMultiplatformExtension ::class.java) ?: return
  extension.targets
    .flatMap { it.compilations }
    .filterIsInstance<KotlinNativeCompilation>()
    .forEach { compilationUnit ->
      compilationUnit.kotlinOptions.freeCompilerArgs += arrayOf("-linker-options", "-lsqlite3")
    }
}

package com.squareup.sqldelight.gradle.kotlin

import com.android.build.gradle.AppExtension
import com.android.build.gradle.BaseExtension
import com.android.build.gradle.LibraryExtension
import com.android.build.gradle.api.BaseVariant
import com.android.build.gradle.internal.tasks.factory.dependsOn
import com.squareup.sqldelight.gradle.SqlDelightDatabase
import com.squareup.sqldelight.gradle.SqlDelightTask
import org.gradle.api.DomainObjectSet
import org.gradle.api.Project
import org.gradle.api.Task
import org.gradle.api.file.SourceDirectorySet
import org.gradle.api.tasks.SourceSetContainer
import org.gradle.api.tasks.TaskContainer
import org.gradle.api.tasks.TaskProvider
import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType
import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinJvmAndroidCompilation
import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinMetadataTarget
import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinNativeTarget

/**
 * @return A list of source roots and their dependencies.
 *
 * Examples:
 *   Multiplatform Environment. Ios target labeled "ios".
 *     -> iosMain deps [commonMain]
 *
 *   Android environment. internal, production, release, debug variants.
 *     -> internalDebug deps [internal, debug, main]
 *     -> internalRelease deps [internal, release, main]
 *     -> productionDebug deps [production, debug, main]
 *     -> productionRelease deps [production, release, main]
 *
 *    Multiplatform environment with android target (oh boy)
 */
internal fun SqlDelightDatabase.sources(): List<Source> {
  // Multiplatform project.
  project.extensions.findByType(KotlinMultiplatformExtension::class.java)?.let {
    return it.sources(project)
  }

  // Android project.
  project.extensions.findByName("android")?.let {
    return (it as BaseExtension).sources(project)
  }

  // Kotlin project.
  val sourceSets = project.extensions.getByName("sourceSets") as SourceSetContainer
  return listOf(
    Source(
      type = KotlinPlatformType.jvm,
      name = "main",
      sourceSets = listOf("main"),
      sourceDirectorySet = sourceSets.getByName("main").kotlin!!,
      registerTaskDependency = { task ->
        project.tasks.named("compileKotlin").configure { it.dependsOn(task) }
      }
    )
  )
}

private fun KotlinMultiplatformExtension.sources(project: Project): List<Source> {
  // For multiplatform we only support SQLDelight in commonMain - to support other source sets
  // we would need to generate expect/actual SQLDelight code which at least right now doesn't
  // seem like there is a use case for. However this code is capable of running on any Target type.
  val target = targets.single { it is KotlinMetadataTarget }
  return target.compilations.mapNotNull { compilation ->
    if (compilation.name.endsWith(suffix = "Test", ignoreCase = true)) {
      // TODO: If we can include these compilations as sqldelight compilation units, we solve
      //  the testing problem. However there's no api to get the main compilation for a test
      //  compilation, except for native where KotlinNativeCompilation has a
      //  "friendCompilationName" which is the main compilation unit. There looks to be
      //  nothing for the other compilation units, but we should revisit later to see if
      //  theres a way to accomplish this.
      return@mapNotNull null
    }
    val targetName = if (target is KotlinMetadataTarget) "common" else target.name
    Source(
      type = target.platformType,
      nativePresetName = (target as? KotlinNativeTarget)?.preset?.name,
      name = "$targetName${compilation.name.capitalize()}",
      variantName = (compilation as? KotlinJvmAndroidCompilation)?.name,
      sourceDirectorySet = compilation.defaultSourceSet.kotlin,
      sourceSets = compilation.allKotlinSourceSets.map { it.name },
      registerTaskDependency = { task ->
        targets.forEach { target ->
          target.compilations.forEach { compilation ->
            (target as? KotlinNativeTarget)?.binaries?.forEach {
              it.linkTask.dependsOn(task)
            }
            compilation.compileKotlinTask.dependsOn(task)
          }
        }
      }
    )
  }
}

private fun BaseExtension.sources(project: Project): List<Source> {
  val variants: DomainObjectSet<out BaseVariant> = when (this) {
    is AppExtension -> applicationVariants
    is LibraryExtension -> libraryVariants
    else -> throw IllegalStateException("Unknown Android plugin $this")
  }
  val sourceSets = sourceSets
    .associate { sourceSet ->
      sourceSet.name to sourceSet.kotlin
    }

  return variants.map { variant ->
    Source(
      type = KotlinPlatformType.androidJvm,
      name = variant.name,
      variantName = variant.name,
      sourceDirectorySet = sourceSets[variant.name]
        ?: throw IllegalStateException("Couldn't find ${variant.name} in $sourceSets"),
      sourceSets = variant.sourceSets.map { it.name },
      registerTaskDependency = { task ->
        // TODO: Lazy task configuration!!!
        variant.registerJavaGeneratingTask(task.get(), task.get().outputDirectory)
        // We have to explicitly add dependencies between kotlin tasks
        // and the generation task as the method registerJavaGeneratingTask above doesn't
        // fully support generation of kotlin code.
        project.tasks.named("compile${variant.name.capitalize()}Kotlin").dependsOn(task)
        project.tasks
          .namedOrNull("kaptGenerateStubs${variant.name.capitalize()}Kotlin")
          ?.dependsOn(task)
      }
    )
  }
}

private fun TaskContainer.namedOrNull(
  taskName: String
): TaskProvider<Task>? {
  return try {
    named(taskName)
  } catch (_: Exception) {
    null
  }
}

internal data class Source(
  val type: KotlinPlatformType,
  val nativePresetName: String? = null,
  val sourceDirectorySet: SourceDirectorySet,
  val name: String,
  val variantName: String? = null,
  val sourceSets: List<String>,
  val registerTaskDependency: (TaskProvider<SqlDelightTask>) -> Unit
) {
  fun closestMatch(sources: Collection<Source>): Source? {
    var matches = sources.filter {
      type == it.type || (type == KotlinPlatformType.androidJvm && it.type == KotlinPlatformType.jvm) || it.type == KotlinPlatformType.common
    }
    if (matches.size <= 1) return matches.singleOrNull()

    // Multiplatform native matched or android variants matched.
    matches = matches.filter {
      nativePresetName == it.nativePresetName && variantName == it.variantName
    }
    return matches.singleOrNull()
  }
}

package com.squareup.sqldelight.gradle.kotlin

import com.android.build.gradle.api.AndroidSourceSet
import org.gradle.api.file.SourceDirectorySet
import org.gradle.api.internal.HasConvention
import org.gradle.api.tasks.SourceSet
import org.jetbrains.kotlin.gradle.plugin.KOTLIN_DSL_NAME
import org.jetbrains.kotlin.gradle.plugin.KOTLIN_JS_DSL_NAME
import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet

internal val AndroidSourceSet.kotlin: SourceDirectorySet?
  get() = kotlinSourceSet

internal val SourceSet.kotlin: SourceDirectorySet?
  get() = kotlinSourceSet

// Copied from kotlin plugin
private val Any.kotlinSourceSet: SourceDirectorySet?
  get() {
    val convention = (getConvention(KOTLIN_DSL_NAME) ?: getConvention(KOTLIN_JS_DSL_NAME)) ?: return null
    val kotlinSourceSetIface =
      convention.javaClass.interfaces.find { it.name == KotlinSourceSet::class.qualifiedName }
    val getKotlin = kotlinSourceSetIface?.methods?.find { it.name == "getKotlin" } ?: return null
    return getKotlin(convention) as? SourceDirectorySet
  }

private fun Any.getConvention(name: String): Any? =
  (this as HasConvention).convention.plugins[name]

package com.squareup.sqldelight.intellij.actions

import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.intellij.codeInsight.CodeInsightActionHandler
import com.intellij.codeInsight.actions.CodeInsightAction
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiFile
import com.intellij.psi.util.parentOfType
import com.squareup.sqldelight.core.lang.SqlDelightFile

internal abstract class BaseGenerateAction(private val handler: CodeInsightActionHandler) : CodeInsightAction() {

  override fun isValidForFile(project: Project, editor: Editor, file: PsiFile): Boolean {
    val offset = editor.caretModel.offset
    return file is SqlDelightFile && file.findElementAt(offset)
      ?.parentOfType<SqlCreateTableStmt>(true) != null
  }

  override fun getHandler(): CodeInsightActionHandler = handler
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij.actions

import com.alecstrong.sql.psi.core.psi.SqlStmt
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.actionSystem.LangDataKeys
import com.intellij.openapi.actionSystem.PlatformDataKeys
import com.intellij.openapi.ide.CopyPasteManager
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.SqlDelightFileType
import com.squareup.sqldelight.core.lang.util.rawSqlText
import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
import java.awt.datatransfer.StringSelection

class CopyAsSqliteAction : AnAction() {
  override fun actionPerformed(e: AnActionEvent) {
    val element = e.sqlElementAtCaret() ?: return
    CopyPasteManager.getInstance().setContents(StringSelection(element.rawSqlText()))
  }

  override fun update(e: AnActionEvent) {
    e.presentation.isVisible =
      e.getData(PlatformDataKeys.VIRTUAL_FILE)?.extension == SqlDelightFileType.defaultExtension &&
      e.sqlElementAtCaret() != null
  }

  private fun AnActionEvent.sqlElementAtCaret(): SqlStmt? {
    val caret = getData(LangDataKeys.CARET)
      ?: return getData(LangDataKeys.PSI_ELEMENT)?.getStrictParentOfType()
    val file = (getData(LangDataKeys.PSI_FILE) as? SqlDelightFile)
    return file?.findElementAt(caret.offset)?.getStrictParentOfType()
  }
}

package com.squareup.sqldelight.intellij.actions

import com.intellij.openapi.editor.Editor
import com.intellij.openapi.editor.ScrollType
import com.intellij.psi.PsiElement
import org.jetbrains.kotlin.psi.psiUtil.endOffset

internal fun insertNewLineAndCleanup(editor: Editor, anchor: PsiElement) {
  val document = editor.document
  val s = "\n\n"
  document.insertString(anchor.endOffset, s)
  val startOffset = anchor.endOffset + s.length
  document.replaceString(startOffset, document.textLength, "")
  editor.caretModel.moveToOffset(startOffset)
  editor.scrollingModel.scrollToCaret(ScrollType.RELATIVE)
  editor.selectionModel.removeSelection()
}

package com.squareup.sqldelight.intellij.actions

internal class GenerateDeleteByPrimaryKeyQueryAction : BaseGenerateAction(QueryByPrimaryKeyHandler("del"))

package com.squareup.sqldelight.intellij.actions

import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.intellij.codeInsight.CodeInsightActionHandler
import com.intellij.codeInsight.template.TemplateManager
import com.intellij.codeInsight.template.impl.TemplateManagerImpl
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiFile
import com.intellij.psi.util.parentOfType
import com.squareup.sqldelight.core.lang.util.findChildOfType
import com.squareup.sqldelight.core.psi.SqlDelightStmtList

internal class GenerateInsertIntoQueryAction : BaseGenerateAction(InsertIntoHandler()) {
  class InsertIntoHandler : CodeInsightActionHandler {
    override fun invoke(project: Project, editor: Editor, file: PsiFile) {
      val caretOffset = editor.caretModel.offset
      val createTableStmt = file.findElementAt(caretOffset)?.parentOfType<SqlCreateTableStmt>() ?: return

      val tableName = createTableStmt.tableName.name
      val template = TemplateManagerImpl.listApplicableTemplates(file, caretOffset, false)
        .first { it.key == "ins" }

      val stmtList = file.findChildOfType<SqlDelightStmtList>() ?: return
      insertNewLineAndCleanup(editor, stmtList)
      TemplateManager.getInstance(project).startTemplate(
        editor, template, false, mapOf("table" to tableName), null
      )
    }
  }
}

package com.squareup.sqldelight.intellij.actions

import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.intellij.codeInsight.CodeInsightActionHandler
import com.intellij.codeInsight.template.TemplateManager
import com.intellij.codeInsight.template.impl.TemplateManagerImpl
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiFile
import com.intellij.psi.util.parentOfType
import com.squareup.sqldelight.core.lang.util.findChildOfType
import com.squareup.sqldelight.core.psi.SqlDelightStmtList

internal class GenerateSelectAllQueryAction : BaseGenerateAction(SelectAllHandler()) {

  class SelectAllHandler : CodeInsightActionHandler {
    override fun invoke(project: Project, editor: Editor, file: PsiFile) {
      val caretOffset = editor.caretModel.offset
      val createTableStmt = file.findElementAt(caretOffset)?.parentOfType<SqlCreateTableStmt>() ?: return

      val tableName = createTableStmt.tableName.name
      val selectAllTemplate = TemplateManagerImpl.listApplicableTemplates(file, caretOffset, false)
        .first { it.key == "sel" }

      val stmtList = file.findChildOfType<SqlDelightStmtList>() ?: return
      insertNewLineAndCleanup(editor, stmtList)
      TemplateManager.getInstance(project).startTemplate(
        editor, selectAllTemplate, false, mapOf("table" to tableName), null
      )
    }
  }
}

package com.squareup.sqldelight.intellij.actions

internal class GenerateSelectByPrimaryKeyQueryAction : BaseGenerateAction(QueryByPrimaryKeyHandler("selw"))

package com.squareup.sqldelight.intellij.actions

internal class GenerateUpdateByPrimaryKeyQueryAction : BaseGenerateAction(QueryByPrimaryKeyHandler("upd"))

package com.squareup.sqldelight.intellij.actions

import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.intellij.codeInsight.CodeInsightActionHandler
import com.intellij.codeInsight.template.TemplateManager
import com.intellij.codeInsight.template.impl.TemplateManagerImpl
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiFile
import com.intellij.psi.util.parentOfType
import com.squareup.sqldelight.core.lang.util.findChildOfType
import com.squareup.sqldelight.core.psi.SqlDelightStmtList

class QueryByPrimaryKeyHandler(private val templateName: String) : CodeInsightActionHandler {
  override fun invoke(project: Project, editor: Editor, file: PsiFile) {
    val caretOffset = editor.caretModel.offset
    val createTableStmt = file.findElementAt(caretOffset)?.parentOfType<SqlCreateTableStmt>() ?: return

    val tableName = createTableStmt.tableName.name
    val pk = createTableStmt.columnDefList.firstOrNull { columnDef ->
      columnDef.columnConstraintList.any { columnConstraint ->
        columnConstraint.textMatches("PRIMARY KEY")
      }
    }?.columnName

    val template = TemplateManagerImpl.listApplicableTemplates(file, caretOffset, false)
      .first { it.key == templateName }

    val args = mutableMapOf("table" to tableName)
    if (pk != null) {
      args += "pkey" to pk.name
    }

    val stmtList = file.findChildOfType<SqlDelightStmtList>() ?: return
    insertNewLineAndCleanup(editor, stmtList)
    TemplateManager.getInstance(project).startTemplate(
      editor, template, false, args, null
    )
  }
}

package com.squareup.sqldelight.intellij.gradle

import com.squareup.sqldelight.core.SqlDelightPropertiesFile
import org.gradle.tooling.BuildAction
import org.gradle.tooling.BuildController
import org.gradle.tooling.UnknownModelException

object FetchProjectModelsBuildAction : BuildAction<Map<String, SqlDelightPropertiesFile?>> {
  override fun execute(controller: BuildController): Map<String, SqlDelightPropertiesFile?> {
    val models = mutableMapOf<String, SqlDelightPropertiesFile?>()
    for (project in controller.buildModel.projects) {
      // Make sure the keys of our map always use '/' as path separator
      val path = project.projectDirectory.absoluteFile.invariantSeparatorsPath
      try {
        models[path] = controller.getModel(project, SqlDelightPropertiesFile::class.java)
      } catch (e: UnknownModelException) {
        models[path] = null
      }
    }
    return models
  }
}

package com.squareup.sqldelight.intellij.gradle

import com.intellij.notification.NotificationDisplayType
import com.intellij.notification.NotificationGroup
import com.intellij.notification.NotificationType
import com.intellij.openapi.externalSystem.model.ExternalSystemException
import com.intellij.openapi.externalSystem.service.execution.ExternalSystemJdkUtil
import com.intellij.openapi.externalSystem.util.ExternalSystemApiUtil
import com.intellij.openapi.module.Module
import com.intellij.openapi.progress.ProgressIndicator
import com.intellij.openapi.progress.ProgressManager
import com.intellij.openapi.progress.Task
import com.intellij.openapi.project.Project
import com.squareup.sqldelight.core.GradleCompatibility
import com.squareup.sqldelight.core.GradleCompatibility.CompatibilityReport.Incompatible
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.SqlDelightProjectService
import com.squareup.sqldelight.core.dialectPreset
import com.squareup.sqldelight.intellij.FileIndex
import com.squareup.sqldelight.intellij.SqlDelightFileIndexImpl
import org.jetbrains.plugins.gradle.service.execution.GradleExecutionHelper
import org.jetbrains.plugins.gradle.settings.DistributionType
import org.jetbrains.plugins.gradle.settings.GradleExecutionSettings
import timber.log.Timber
import java.io.File

internal class FileIndexMap {
  private val fileIndices = mutableMapOf<String, SqlDelightFileIndex>()
  private val incompatibilityNotifier = IncompatibilityNotifier()

  private var initialized = false
  private var retries = 0

  operator fun get(module: Module): SqlDelightFileIndex {
    val projectPath = ExternalSystemApiUtil.getExternalProjectPath(module) ?: return defaultIndex
    val result = fileIndices[projectPath]
    if (result != null) return result
    synchronized(this) {
      if (!initialized) {
        initialized = true
        if (!module.isDisposed && !module.project.isDisposed)
          try {
            ProgressManager.getInstance().run(FetchModuleModels(module, projectPath))
          } catch (e: Throwable) {
            // IntelliJ can fail to start the fetch command, reinitialize later in this case.
            if (retries++ < 3) {
              initialized = false
            }
          }
      }
    }
    return fileIndices[projectPath] ?: defaultIndex
  }

  private inner class FetchModuleModels(
    private val module: Module,
    private val projectPath: String
  ) : Task.Backgroundable(
    /* project = */ module.project,
    /* title = */ "Importing ${module.name} SQLDelight"
  ) {
    override fun run(indicator: ProgressIndicator) {
      val executionSettings = GradleExecutionSettings(
        /* gradleHome = */ null,
        /* serviceDirectory = */ null,
        /* distributionType = */ DistributionType.DEFAULT_WRAPPED,
        /* isOfflineWork = */ false
      )
      try {
        fileIndices.putAll(
          GradleExecutionHelper().execute(projectPath, executionSettings) { connection ->
            Timber.i("Fetching SQLDelight models")
            val javaHome =
              ExternalSystemJdkUtil.getJdk(project, ExternalSystemJdkUtil.USE_PROJECT_JDK)
                ?.homeDirectory?.path?.let { File(it) }
            val properties =
              connection.action(FetchProjectModelsBuildAction).setJavaHome(javaHome).run()

            Timber.i("Assembling file index")
            return@execute properties.mapValues { (_, value) ->
              if (value == null) return@mapValues defaultIndex

              val compatibility = GradleCompatibility.validate(value)
              if (compatibility is Incompatible) {
                incompatibilityNotifier.notifyError(project, compatibility.reason)
                return@mapValues defaultIndex
              }

              val database = value.databases.first()
              SqlDelightProjectService.getInstance(module.project).dialectPreset =
                database.dialectPreset
              return@mapValues FileIndex(database)
            }
          }
        )
      } catch (externalException: ExternalSystemException) {
        // It's a gradle error, ignore and let the user fix when they try and build the project
      }

      Timber.i("Initialized file index")
      initialized = false
    }
  }

  companion object {
    internal var defaultIndex: SqlDelightFileIndex = SqlDelightFileIndexImpl()
  }
}

private class IncompatibilityNotifier {
  private var showedIncompatibilityBalloon = false

  fun notifyError(project: Project, content: String) {
    if (showedIncompatibilityBalloon) return

    // We only want to show this balloon once per project open.
    showedIncompatibilityBalloon = true
    NOTIFICATION_GROUP.createNotification(content, NotificationType.ERROR).notify(project)
  }

  companion object {
    private val NOTIFICATION_GROUP =
      NotificationGroup("SQLDelight Notification Group", NotificationDisplayType.BALLOON, true)
  }
}

package com.squareup.sqldelight.intellij.inspections

import com.alecstrong.sql.psi.core.psi.SqlBinaryEqualityExpr
import com.alecstrong.sql.psi.core.psi.SqlColumnExpr
import com.alecstrong.sql.psi.core.psi.SqlExpr
import com.alecstrong.sql.psi.core.psi.SqlJoinConstraint
import com.alecstrong.sql.psi.core.psi.SqlParenExpr
import com.intellij.codeInspection.InspectionManager
import com.intellij.codeInspection.LocalInspectionTool
import com.intellij.codeInspection.ProblemDescriptor
import com.intellij.codeInspection.ProblemHighlightType
import com.intellij.psi.PsiFile
import com.squareup.sqldelight.core.lang.psi.isColumnSameAs
import com.squareup.sqldelight.core.lang.psi.isTypeSameAs
import com.squareup.sqldelight.core.lang.util.findChildrenOfType

class MismatchJoinColumnInspection : LocalInspectionTool() {
  override fun checkFile(
    file: PsiFile,
    manager: InspectionManager,
    isOnTheFly: Boolean
  ): Array<ProblemDescriptor> {
    val joinConstraints = file.findChildrenOfType<SqlJoinConstraint>()
      .mapNotNull { it.expr?.binaryEqualityExpr() }

    return joinConstraints.mapNotNull { joinEquality ->
      val (expr1, expr2) = joinEquality.getExprList()
      if (expr1 !is SqlColumnExpr || expr2 !is SqlColumnExpr) return@mapNotNull null

      val column1 = expr1.columnName
      val column2 = expr2.columnName

      if (column1.isColumnSameAs(column2)) {
        return@mapNotNull manager.createProblemDescriptor(
          joinEquality, "Join condition always evaluates to true", isOnTheFly,
          emptyArray(), ProblemHighlightType.WARNING
        )
      }

      if (!column1.isTypeSameAs(column2)) {
        return@mapNotNull manager.createProblemDescriptor(
          joinEquality, "Join compares two columns of different types", isOnTheFly,
          emptyArray(), ProblemHighlightType.WARNING
        )
      }
      return@mapNotNull null
    }.toTypedArray()
  }

  private fun SqlExpr.binaryEqualityExpr(): SqlBinaryEqualityExpr? = when (this) {
    is SqlParenExpr -> expr?.binaryEqualityExpr()
    is SqlBinaryEqualityExpr -> this
    else -> null
  }
}

package com.squareup.sqldelight.intellij.inspections

import com.alecstrong.sql.psi.core.psi.SqlColumnDef
import com.alecstrong.sql.psi.core.psi.SqlColumnExpr
import com.alecstrong.sql.psi.core.psi.SqlIsExpr
import com.alecstrong.sql.psi.core.psi.SqlSelectStmt
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.alecstrong.sql.psi.core.psi.SqlVisitor
import com.intellij.codeInspection.LocalInspectionTool
import com.intellij.codeInspection.LocalInspectionToolSession
import com.intellij.codeInspection.ProblemHighlightType
import com.intellij.codeInspection.ProblemsHolder
import com.intellij.psi.PsiElementVisitor
import com.intellij.psi.util.PsiTreeUtil
import com.intellij.psi.util.elementType
import com.intellij.psi.util.parentOfType

internal class RedundantNullCheckInspection : LocalInspectionTool() {

  override fun buildVisitor(
    holder: ProblemsHolder,
    isOnTheFly: Boolean,
    session: LocalInspectionToolSession
  ): PsiElementVisitor = object : SqlVisitor() {

    override fun visitIsExpr(o: SqlIsExpr) {
      if (o.context !is SqlSelectStmt) return
      if (PsiTreeUtil.prevVisibleLeaf(o).elementType != SqlTypes.WHERE) return

      val firstChild = o.firstChild
      if (firstChild !is SqlColumnExpr) return
      val clauseText = o.text
      if (!clauseText.endsWith("IS NOT NULL") && !clauseText.endsWith("IS NULL")) return

      val columnName = firstChild.columnName
      val sqlColumnDef = columnName.reference?.resolve()?.parentOfType<SqlColumnDef>() ?: return
      val hasNotNullConstraint = sqlColumnDef.columnConstraintList.any { constraint ->
        constraint.textMatches("NOT NULL")
      }
      if (!hasNotNullConstraint) {
        return
      }

      holder.registerProblem(o, "Column ${columnName.name} defined as NOT NULL", ProblemHighlightType.WARNING)
    }
  }
}

package com.squareup.sqldelight.intellij.inspections

import com.alecstrong.sql.psi.core.psi.AliasElement
import com.alecstrong.sql.psi.core.psi.SqlColumnDef
import com.alecstrong.sql.psi.core.psi.SqlColumnName
import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.alecstrong.sql.psi.core.psi.SqlVisitor
import com.intellij.codeInspection.LocalInspectionTool
import com.intellij.codeInspection.LocalInspectionToolSession
import com.intellij.codeInspection.LocalQuickFixOnPsiElement
import com.intellij.codeInspection.ProblemHighlightType
import com.intellij.codeInspection.ProblemsHolder
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.editor.ReadOnlyModificationException
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiElementVisitor
import com.intellij.psi.PsiFile
import com.intellij.psi.PsiFileFactory
import com.intellij.psi.PsiManager
import com.intellij.psi.SmartPointerManager
import com.intellij.psi.search.FileTypeIndex
import com.intellij.psi.search.GlobalSearchScope
import com.intellij.psi.search.searches.ReferencesSearch
import com.intellij.psi.util.PsiTreeUtil
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.SqlDelightFileType
import com.squareup.sqldelight.core.lang.util.findChildOfType
import com.squareup.sqldelight.core.lang.util.findChildrenOfType

class UnusedColumnInspection : LocalInspectionTool() {

  override fun buildVisitor(
    holder: ProblemsHolder,
    isOnTheFly: Boolean,
    session: LocalInspectionToolSession
  ): PsiElementVisitor = object : SqlVisitor() {
    override fun visitCreateTableStmt(o: SqlCreateTableStmt) {
      val candidates = o.findChildrenOfType<SqlColumnDef>()
        .filter { columnDef ->
          val columnName = columnDef.columnName
          ReferencesSearch.search(columnName, columnName.useScope)
            .allMatch { reference -> reference.element.parent is SqlColumnDef }
        }
        .toMutableList()

      if (candidates.isEmpty()) {
        return
      }

      val project = o.project
      val psiManager = PsiManager.getInstance(project)

      fun PsiElement.columnDef(): SqlColumnDef? {
        if (this is SqlColumnName) return parent.columnDef()
        if (this is SqlColumnDef) return this
        if (this is AliasElement) return source().columnDef()
        return null
      }

      FileTypeIndex.getFiles(SqlDelightFileType, GlobalSearchScope.allScope(project))
        .mapNotNull { vFile -> psiManager.findFile(vFile) as SqlDelightFile? }
        .asSequence()
        .flatMap { file -> file.sqlStmtList?.stmtList.orEmpty() }
        .flatMap { stmt -> stmt.compoundSelectStmt?.queryExposed().orEmpty() }
        .flatMap { queryResult -> queryResult.columns }
        .forEach { column ->
          column.element.columnDef()?.let(candidates::remove)
        }

      candidates.forEach { columnDef ->
        holder.registerProblem(
          columnDef.columnName, "Unused symbol", ProblemHighlightType.LIKE_UNUSED_SYMBOL,
          SafeDeleteQuickFix(o, columnDef)
        )
      }
    }
  }

  class SafeDeleteQuickFix(createTableStmt: SqlCreateTableStmt, columnDef: SqlColumnDef) :
    LocalQuickFixOnPsiElement(createTableStmt) {

    private val createTableRef = SmartPointerManager.getInstance(createTableStmt.project)
      .createSmartPsiElementPointer(createTableStmt, createTableStmt.containingFile)

    private val columnDefRef = SmartPointerManager.getInstance(columnDef.project)
      .createSmartPsiElementPointer(columnDef, columnDef.containingFile)

    override fun getFamilyName(): String = name

    override fun getText(): String = "Safe delete ${columnDefRef.element?.columnName?.text.orEmpty()}"

    override fun invoke(
      project: Project,
      file: PsiFile,
      startElement: PsiElement,
      endElement: PsiElement
    ) {
      WriteCommandAction.writeCommandAction(project).run<ReadOnlyModificationException> {
        val createTable = createTableRef.element ?: return@run
        val columnDef = columnDefRef.element ?: return@run

        val columnDefString = createTable.columnDefList.filter { it !== columnDef }
          .joinToString(",\n  ") { it.text.trim() }

        val startTable = PsiTreeUtil.findSiblingForward(createTable.firstChild, SqlTypes.LP, false, null) ?: return@run
        val endTable = PsiTreeUtil.findSiblingBackward(createTable.lastChild, SqlTypes.RP, false, null) ?: return@run
        createTable.deleteChildRange(startTable.nextSibling, endTable.prevSibling)

        if (columnDefString.isEmpty()) {
          return@run
        }

        val factory = PsiFileFactory.getInstance(file.project)
        val createTableStmt = """
          |CREATE TABLE ${createTable.tableName.id?.text}(
          |  $columnDefString
          |)
        """.trimMargin()
        val dummyFile = factory.createFileFromText(
          "_Dummy_.${SqlDelightFileType.EXTENSION}",
          SqlDelightFileType, createTableStmt
        )

        val newCreateTable = dummyFile.findChildOfType<SqlCreateTableStmt>() ?: return@run
        val newStart = PsiTreeUtil.findSiblingForward(newCreateTable.firstChild, SqlTypes.LP, false, null) ?: return@run
        val newEnd = PsiTreeUtil.findSiblingBackward(newCreateTable.lastChild, SqlTypes.RP, false, null) ?: return@run

        createTable.addRangeAfter(newStart.nextSibling, newEnd.prevSibling, startTable)
      }
    }
  }
}

package com.squareup.sqldelight.intellij.inspections

import com.intellij.codeInsight.actions.OptimizeImportsProcessor
import com.intellij.codeInspection.InspectionManager
import com.intellij.codeInspection.LocalInspectionTool
import com.intellij.codeInspection.LocalQuickFixOnPsiElement
import com.intellij.codeInspection.ProblemDescriptor
import com.intellij.codeInspection.ProblemHighlightType
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import com.squareup.sqldelight.core.lang.psi.ImportStmtMixin
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.psi.SqlDelightColumnType
import com.squareup.sqldelight.core.psi.SqlDelightJavaType
import com.squareup.sqldelight.core.psi.SqlDelightJavaTypeName

class UnusedImportInspection : LocalInspectionTool() {

  override fun runForWholeFile(): Boolean = true

  override fun checkFile(
    file: PsiFile,
    manager: InspectionManager,
    isOnTheFly: Boolean
  ): Array<ProblemDescriptor> {
    val javaTypes = file.columnJavaTypes()

    return file.findChildrenOfType<ImportStmtMixin>()
      .filter { importStmtMixin ->
        importStmtMixin.javaType.text.substringAfterLast(".").removeSuffix(";") !in javaTypes
      }
      .map { importStmtMixin ->
        manager.createProblemDescriptor(
          importStmtMixin, "Unused import", isOnTheFly,
          arrayOf(RemoveUnusedImportQuickFix(file)), ProblemHighlightType.LIKE_UNUSED_SYMBOL
        )
      }
      .toTypedArray()
  }

  class RemoveUnusedImportQuickFix(file: PsiFile) : LocalQuickFixOnPsiElement(file) {

    override fun getFamilyName(): String = name

    override fun getText(): String = "Optimize imports"

    override fun invoke(
      project: Project,
      file: PsiFile,
      startElement: PsiElement,
      endElement: PsiElement
    ) {
      OptimizeImportsProcessor(project, file).run()
    }
  }
}

fun PsiFile.columnJavaTypes(): Set<String> =
  findChildrenOfType<SqlDelightColumnType>()
    .asSequence()
    .flatMap { columnType ->
      columnType.findChildrenOfType<SqlDelightJavaType>() + columnType.findChildrenOfType<SqlDelightJavaTypeName>()
    }
    .mapNotNull { it.text.substringBefore(".") }
    .toSet()

package com.squareup.sqldelight.intellij.inspections

import com.alecstrong.sql.psi.core.psi.SqlStmtList
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.intellij.codeInspection.LocalInspectionTool
import com.intellij.codeInspection.LocalInspectionToolSession
import com.intellij.codeInspection.LocalQuickFixOnPsiElement
import com.intellij.codeInspection.ProblemHighlightType
import com.intellij.codeInspection.ProblemsHolder
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.editor.ReadOnlyModificationException
import com.intellij.openapi.module.ModuleUtil
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiElementVisitor
import com.intellij.psi.PsiFile
import com.intellij.psi.SmartPointerManager
import com.intellij.psi.search.FilenameIndex
import com.intellij.psi.search.GlobalSearchScope
import com.intellij.psi.search.searches.ReferencesSearch
import com.intellij.psi.util.PsiTreeUtil
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.psi.StmtIdentifierMixin
import com.squareup.sqldelight.core.lang.queriesName
import com.squareup.sqldelight.core.lang.util.findChildOfType
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.psi.SqlDelightStmtIdentifier
import com.squareup.sqldelight.core.psi.SqlDelightVisitor
import org.jetbrains.kotlin.asJava.toLightMethods
import org.jetbrains.kotlin.psi.KtNamedFunction

class UnusedQueryInspection : LocalInspectionTool() {

  override fun buildVisitor(
    holder: ProblemsHolder,
    isOnTheFly: Boolean,
    session: LocalInspectionToolSession
  ): PsiElementVisitor {
    val sqlDelightFile = session.file as SqlDelightFile
    val module =
      ModuleUtil.findModuleForPsiElement(sqlDelightFile) ?: return PsiElementVisitor.EMPTY_VISITOR
    val fileName = "${sqlDelightFile.virtualFile?.queriesName}.kt"
    val generatedFile = FilenameIndex.getFilesByName(
      sqlDelightFile.project, fileName, GlobalSearchScope.moduleScope(module)
    ).firstOrNull() ?: return PsiElementVisitor.EMPTY_VISITOR
    val allMethods = generatedFile.findChildrenOfType<KtNamedFunction>()
    return object : SqlDelightVisitor() {
      override fun visitStmtIdentifier(o: SqlDelightStmtIdentifier) {
        if (o !is StmtIdentifierMixin || o.identifier() == null) {
          return
        }
        val generatedMethods = allMethods.filter { namedFunction ->
          namedFunction.name == o.identifier()?.text
        }
        for (generatedMethod in generatedMethods) {
          val lightMethods = generatedMethod.toLightMethods()
          for (psiMethod in lightMethods) {
            val reference = ReferencesSearch.search(psiMethod, psiMethod.useScope).findFirst()
            if (reference != null) {
              return
            }
          }
        }
        holder.registerProblem(
          o, "Unused symbol", ProblemHighlightType.LIKE_UNUSED_SYMBOL, SafeDeleteQuickFix(o)
        )
      }
    }
  }

  class SafeDeleteQuickFix(element: PsiElement) : LocalQuickFixOnPsiElement(element) {
    private val ref = SmartPointerManager.getInstance(element.project)
      .createSmartPsiElementPointer(element, element.containingFile)

    override fun getFamilyName(): String = name

    override fun getText(): String = "Safe delete ${ref.element?.text?.removeSuffix(":").orEmpty()}"

    override fun invoke(
      project: Project,
      file: PsiFile,
      startElement: PsiElement,
      endElement: PsiElement
    ) {
      WriteCommandAction.writeCommandAction(project).run<ReadOnlyModificationException> {
        val element = ref.element ?: return@run
        val semicolon = PsiTreeUtil.findSiblingForward(element, SqlTypes.SEMI, false, null)
        file.findChildOfType<SqlStmtList>()?.deleteChildRange(element, semicolon)
      }
    }
  }
}

package com.squareup.sqldelight.intellij.intentions

import com.intellij.codeInsight.daemon.QuickFixBundle
import com.intellij.codeInsight.intention.BaseElementAtCaretIntentionAction
import com.intellij.codeInsight.navigation.NavigationUtil
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.editor.Document
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.module.ModuleUtil
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.popup.PopupStep
import com.intellij.openapi.ui.popup.util.BaseListPopupStep
import com.intellij.openapi.util.Iconable
import com.intellij.psi.PsiClass
import com.intellij.psi.PsiDocumentManager
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import com.intellij.psi.search.GlobalSearchScope
import com.intellij.ui.popup.list.ListPopupImpl
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.psi.SqlDelightImportStmt
import com.squareup.sqldelight.intellij.util.PsiClassSearchHelper
import javax.swing.Icon

class AddImportIntention(private val key: String) : BaseElementAtCaretIntentionAction() {
  override fun getFamilyName(): String = INTENTIONS_FAMILY_NAME_IMPORTS

  override fun getText(): String = "Add import for $key"

  override fun isAvailable(project: Project, editor: Editor, element: PsiElement): Boolean {
    return true
  }

  override fun invoke(project: Project, editor: Editor, element: PsiElement) {
    val module = ModuleUtil.findModuleForPsiElement(element) ?: return
    val scope = GlobalSearchScope.moduleWithDependenciesAndLibrariesScope(module)
    val classes = PsiClassSearchHelper.getClassesByShortName(key, project, scope)
    val document = editor.document
    val file = element.containingFile
    if (classes.size == 1) {
      document.addImport(file, "import ${classes.first().qualifiedName};")
    } else {
      showImportPopup(project, editor, file, classes.sortedBy { it.qualifiedName })
    }
  }

  private fun showImportPopup(
    project: Project,
    editor: Editor,
    psiFile: PsiFile,
    classes: List<PsiClass>
  ) {
    val document = editor.document
    val step = object : BaseListPopupStep<PsiClass>(
      QuickFixBundle.message("class.to.import.chooser.title"), classes
    ) {
      override fun isAutoSelectionEnabled(): Boolean {
        return false
      }

      override fun getTextFor(value: PsiClass): String {
        return requireNotNull(value.qualifiedName)
      }

      override fun getIconFor(value: PsiClass): Icon? {
        return value.getIcon(Iconable.ICON_FLAG_VISIBILITY)
      }

      override fun onChosen(selectedValue: PsiClass?, finalChoice: Boolean): PopupStep<*>? {
        if (selectedValue == null) {
          return FINAL_CHOICE
        }
        if (finalChoice) {
          return doFinalStep {
            PsiDocumentManager.getInstance(project).commitAllDocuments()
            WriteCommandAction.runWriteCommandAction(
              project, QuickFixBundle.message("add.import"),
              null,
              {
                document.addImport(psiFile, "import ${selectedValue.qualifiedName};")
              }
            )
          }
        }
        return super.onChosen(selectedValue, finalChoice)
      }
    }
    val popup = ListPopupImpl(project, step)
    NavigationUtil.hidePopupIfDumbModeStarts(popup, project)
    popup.showInBestPositionFor(editor)
  }

  private fun Document.addImport(file: PsiFile, import: String) {
    val imports = file.findChildrenOfType<SqlDelightImportStmt>()
    if (imports.isEmpty()) {
      insertString(0, "$import\n\n")
    } else {
      val newImports = mutableListOf(import)
      var endOffset = 0
      for (imp in imports) {
        newImports.add(imp.text)
        endOffset = maxOf(endOffset, imp.textOffset + imp.textLength)
      }
      replaceString(0, endOffset, newImports.sorted().joinToString("\n"))
    }
  }
}

package com.squareup.sqldelight.intellij.intentions

import com.alecstrong.sql.psi.core.psi.SqlResultColumn
import com.alecstrong.sql.psi.core.psi.SqlSelectStmt
import com.intellij.codeInsight.intention.impl.BaseIntentionAction
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.editor.ReadOnlyModificationException
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiFile
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.intellij.util.findElementOfTypeAtOffset
import org.jetbrains.kotlin.psi.psiUtil.endOffset
import org.jetbrains.kotlin.psi.psiUtil.startOffset

class ExpandColumnNamesWildcardQuickFix : BaseIntentionAction() {

  override fun getFamilyName() = INTENTIONS_FAMILY_NAME_REFACTORINGS

  override fun getText() = INTENTION_EXPAND_COLUMN_NAMES_TEXT

  override fun isAvailable(project: Project, editor: Editor, file: PsiFile): Boolean {
    val caret = editor.caretModel.offset
    return selectStatementAtCaretWithColumnNamesWildcard(
      file = file as? SqlDelightFile ?: return false,
      caret = caret
    ) != null
  }

  private fun selectStatementAtCaretWithColumnNamesWildcard(
    file: SqlDelightFile,
    caret: Int
  ): SqlSelectStmt? {
    val selectStatement = file.findElementOfTypeAtOffset<SqlSelectStmt>(caret) ?: return null
    val resultColumns = selectStatement.resultColumnList
    return if (resultColumns.any { it.textMatches("*") }) selectStatement else null
  }

  override fun invoke(project: Project, editor: Editor, file: PsiFile) {
    WriteCommandAction.writeCommandAction(project).run<ReadOnlyModificationException> {
      val caret = editor.caretModel.offset
      selectStatementAtCaretWithColumnNamesWildcard(file as SqlDelightFile, caret)?.run {
        val allColumns = queryExposed()
          .flatMap { it.columns }
          .distinctBy { it.element.text }
          .joinToString(separator = ", ") { it.element.text }
        val (start, end) = resultColumnList.startEndOffset()
        editor.document.replaceString(start, end, allColumns)
      }
    }
  }

  private fun List<SqlResultColumn>.startEndOffset(): Pair<Int, Int> {
    var start = firstOrNull()?.startOffset ?: 0
    var end = firstOrNull()?.endOffset ?: 0
    for (i in 1..lastIndex) {
      start = minOf(start, get(i).startOffset)
      end = maxOf(end, get(i).endOffset)
    }
    return start to end
  }
}

package com.squareup.sqldelight.intellij.intentions

internal const val INTENTIONS_FAMILY_NAME_REFACTORINGS = "Refactorings"

internal const val INTENTIONS_FAMILY_NAME_IMPORTS = "Imports"

internal const val INTENTION_EXPAND_COLUMN_NAMES_TEXT = "Expand * into column names"

package com.squareup.sqldelight.intellij.intentions

import com.alecstrong.sql.psi.core.psi.SqlExpr
import com.alecstrong.sql.psi.core.psi.SqlSelectStmt
import com.alecstrong.sql.psi.core.psi.SqlTableAlias
import com.alecstrong.sql.psi.core.psi.SqlTableName
import com.intellij.codeInsight.intention.BaseElementAtCaretIntentionAction
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.popup.JBPopupFactory
import com.intellij.psi.PsiElement
import com.intellij.psi.util.PsiTreeUtil
import com.intellij.psi.util.parentOfType
import com.squareup.sqldelight.core.lang.util.findChildrenOfType

internal class IntroduceTableAliasIntention : BaseElementAtCaretIntentionAction() {

  private val regex = "(?=\\p{Upper})|(_)".toRegex()

  override fun getFamilyName(): String {
    return INTENTIONS_FAMILY_NAME_REFACTORINGS
  }

  override fun getText(): String {
    return "Introduce table alias"
  }

  override fun isAvailable(project: Project, editor: Editor, element: PsiElement): Boolean {
    val tableName = element.parentOfType<SqlTableName>(true) ?: return false
    if (tableName.parent is SqlExpr) {
      return false
    }
    val columnAlias = PsiTreeUtil.getNextSiblingOfType(tableName, SqlTableAlias::class.java)
    if (columnAlias != null) {
      return false
    }
    val sqlSelectStmt = tableName.parentOfType<SqlSelectStmt>() ?: return false
    return sqlSelectStmt.resultColumnList.none { it.textMatches("*") }
  }

  override fun invoke(project: Project, editor: Editor, element: PsiElement) {
    val sqlTableName = element.parentOfType<SqlTableName>(true) ?: return
    val sqlSelectStmt = sqlTableName.parentOfType<SqlSelectStmt>() ?: return

    val tableName = sqlTableName.name
    val aliasVariants = listOf(
      tableName.first().toString(),
      tableName.split(regex)
        .filter(String::isNotBlank)
        .joinToString("") { it.first().toLowerCase().toString() }
    )
      .distinct()

    val document = editor.document
    val tableNameTextRange = sqlTableName.textRange
    val tableNameRangeMarker = document.createRangeMarker(tableNameTextRange)
    val tableReferenceMarkers = sqlSelectStmt.findChildrenOfType<SqlTableName>()
      .filter {
        it.textMatches(tableName) && !it.textRange.equals(tableNameTextRange)
      }
      .map { document.createRangeMarker(it.textRange) }

    val callback = { alias: String ->
      WriteCommandAction.runWriteCommandAction(project) {
        document.replaceString(
          tableNameRangeMarker.startOffset, tableNameRangeMarker.endOffset, "$tableName $alias"
        )
        tableReferenceMarkers.forEach { rangeMarker ->
          document.replaceString(rangeMarker.startOffset, rangeMarker.endOffset, alias)
        }
      }
    }
    if (aliasVariants.size == 1) {
      callback(aliasVariants.first())
    } else {
      JBPopupFactory.getInstance()
        .createPopupChooserBuilder(aliasVariants)
        .setMovable(true)
        .setResizable(true)
        .setItemChosenCallback(callback)
        .createPopup()
        .showInBestPositionFor(editor)
    }
  }
}

package com.squareup.sqldelight.intellij.intentions

import com.alecstrong.sql.psi.core.psi.SqlColumnExpr
import com.alecstrong.sql.psi.core.psi.SqlColumnName
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.intellij.codeInsight.intention.BaseElementAtCaretIntentionAction
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.popup.JBPopupFactory
import com.intellij.psi.PsiElement
import com.intellij.psi.util.elementType
import com.intellij.psi.util.parentOfType

class QualifyColumnNameIntention : BaseElementAtCaretIntentionAction() {

  override fun getFamilyName(): String {
    return INTENTIONS_FAMILY_NAME_REFACTORINGS
  }

  override fun getText(): String {
    return "Qualify identifier"
  }

  override fun isAvailable(project: Project, editor: Editor?, element: PsiElement): Boolean {
    val columnName = element.parentOfType<SqlColumnName>(true) ?: return false
    return columnName.parent is SqlColumnExpr && columnName.prevSibling.elementType != SqlTypes.DOT
  }

  override fun invoke(project: Project, editor: Editor, element: PsiElement) {
    val columnName = element.parentOfType<SqlColumnName>(true) ?: return
    val columnText = columnName.text
    val tableNamesOrAliases = columnName.queryAvailable(columnName)
      .asSequence()
      .filter { result -> result.table != null }
      .flatMap { result ->
        result.columns.filter { column -> column.element.textMatches(columnText) }
          .map { column -> column.element to result.table!! }
      }
      .groupBy({ it.first }, { it.second })
      .values
      .flatten()
      .map { it.text }
      .distinct()

    if (tableNamesOrAliases.isEmpty()) {
      return
    }

    val document = editor.document
    val columnRange = document.createRangeMarker(columnName.textRange)

    val callback = { qualifier: String ->
      WriteCommandAction.runWriteCommandAction(project) {
        document.replaceString(
          columnRange.startOffset, columnRange.endOffset, "$qualifier.$columnText"
        )
      }
    }
    if (tableNamesOrAliases.size == 1) {
      callback(tableNamesOrAliases.first())
    } else {
      JBPopupFactory.getInstance()
        .createPopupChooserBuilder(tableNamesOrAliases)
        .setMovable(true)
        .setResizable(true)
        .setItemChosenCallback(callback)
        .createPopup()
        .showInBestPositionFor(editor)
    }
  }
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij.lang

import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.intellij.lang.CodeDocumentationAwareCommenter
import com.intellij.psi.PsiComment
import com.intellij.psi.tree.IElementType

class SqlDelightCommenter : CodeDocumentationAwareCommenter {

  override fun getLineCommentTokenType(): IElementType = SqlTypes.COMMENT
  override fun getLineCommentPrefix() = "-- "

  override fun getBlockCommentTokenType(): IElementType? = null
  override fun getBlockCommentPrefix(): String? = null
  override fun getBlockCommentSuffix(): String? = null

  override fun getDocumentationCommentTokenType(): IElementType = SqlTypes.JAVADOC
  override fun isDocumentationComment(psiComment: PsiComment?) =
    psiComment?.tokenType == documentationCommentTokenType

  override fun getDocumentationCommentPrefix() = "/**"
  override fun getDocumentationCommentLinePrefix() = "*"
  override fun getDocumentationCommentSuffix() = "*/"

  override fun getCommentedBlockCommentPrefix(): String? = null
  override fun getCommentedBlockCommentSuffix(): String? = null
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij.lang

import com.alecstrong.sql.psi.core.SqlAnnotationHolder
import com.alecstrong.sql.psi.core.psi.SqlAnnotatedElement
import com.intellij.lang.Language
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.module.Module
import com.intellij.openapi.util.Condition
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.FileViewProvider
import com.intellij.psi.FileViewProviderFactory
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiErrorElement
import com.intellij.psi.PsiManager
import com.intellij.psi.SingleRootFileViewProvider
import com.intellij.psi.util.PsiTreeUtil
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.SqlDelightProjectService
import com.squareup.sqldelight.core.compiler.SqlDelightCompiler
import com.squareup.sqldelight.core.lang.MigrationFile
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.SqlDelightQueriesFile
import com.squareup.sqldelight.intellij.util.GeneratedVirtualFile
import java.io.PrintStream
import java.util.concurrent.Executors
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicBoolean

class SqlDelightFileViewProviderFactory : FileViewProviderFactory {
  override fun createFileViewProvider(
    file: VirtualFile,
    language: Language,
    manager: PsiManager,
    eventSystemEnabled: Boolean
  ): FileViewProvider {
    val module = SqlDelightProjectService.getInstance(manager.project).module(file)
      ?: return SingleRootFileViewProvider(manager, file, eventSystemEnabled)
    return SqlDelightFileViewProvider(manager, file, eventSystemEnabled, language, module)
  }
}

private class SqlDelightFileViewProvider(
  manager: PsiManager,
  virtualFile: VirtualFile,
  eventSystemEnabled: Boolean,
  private val language: Language,
  private val module: Module
) : SingleRootFileViewProvider(manager, virtualFile, eventSystemEnabled, language) {
  private val threadPool = Executors.newScheduledThreadPool(1)

  private val file: SqlDelightFile
    get() = getPsiInner(language) as SqlDelightFile

  private var filesGenerated = emptyList<VirtualFile>()
    set(value) {
      (field - value).forEach { it.delete(this) }
      field = value
    }

  private var condition = WriteCondition()

  override fun contentsSynchronized() {
    super.contentsSynchronized()

    if (!SqlDelightFileIndex.getInstance(module).isConfigured ||
      SqlDelightFileIndex.getInstance(module).sourceFolders(file).isEmpty()
    ) {
      return
    }

    condition.invalidated.set(true)

    if (ApplicationManager.getApplication().isUnitTestMode) {
      generateSqlDelightCode()
      return
    }

    val thisCondition = WriteCondition()
    condition = thisCondition
    threadPool.schedule(
      {
        ApplicationManager.getApplication().invokeLater(
          {
            try {
              generateSqlDelightCode()
            } catch (e: Throwable) {
              // IDE generating code should be best effort - source of truth is always the gradle
              // build, and its better to ignore the error and try again than crash and require
              // the IDE restarts.
              e.printStackTrace()
            }
          },
          thisCondition
        )
      },
      1, TimeUnit.SECONDS
    )
  }

  /**
   * Attempt to generate the SQLDelight code for the file represented by the view provider.
   */
  private fun generateSqlDelightCode() {
    var shouldGenerate = true
    val annotationHolder = object : SqlAnnotationHolder {
      override fun createErrorAnnotation(element: PsiElement, s: String) {
        shouldGenerate = false
      }
    }

    // File is mutable so create a copy that wont be mutated.
    val file = file.copy() as SqlDelightFile

    shouldGenerate = try {
      PsiTreeUtil.processElements(file) { element ->
        when (element) {
          is PsiErrorElement -> return@processElements false
          is SqlAnnotatedElement -> element.annotate(annotationHolder)
        }
        return@processElements shouldGenerate
      }
    } catch (e: Throwable) {
      // If we encountered an exception while looking for errors, assume it was an error.
      false
    }

    if (shouldGenerate && !ApplicationManager.getApplication().isUnitTestMode) ApplicationManager.getApplication().runWriteAction {
      val files = mutableListOf<VirtualFile>()
      val fileAppender = { filePath: String ->
        val vFile: VirtualFile by GeneratedVirtualFile(filePath, module)
        files.add(vFile)
        PrintStream(vFile.getOutputStream(this))
      }
      if (file is SqlDelightQueriesFile) {
        SqlDelightCompiler.writeInterfaces(module, file, fileAppender)
      } else if (file is MigrationFile) {
        SqlDelightCompiler.writeInterfaces(file, fileAppender)
      }
      this.filesGenerated = files
    }
  }

  private class WriteCondition : Condition<Any> {
    var invalidated = AtomicBoolean(false)

    override fun value(t: Any?) = invalidated.get()
  }
}

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij.lang

import com.alecstrong.sql.psi.core.psi.SqlCreateIndexStmt
import com.alecstrong.sql.psi.core.psi.SqlCreateTriggerStmt
import com.alecstrong.sql.psi.core.psi.SqlCreateViewStmt
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.intellij.lang.ASTNode
import com.intellij.lang.folding.FoldingBuilder
import com.intellij.lang.folding.FoldingDescriptor
import com.intellij.openapi.editor.Document
import com.intellij.openapi.project.DumbAware
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiWhiteSpace
import com.squareup.sqldelight.core.psi.SqlDelightStmtIdentifier
import com.squareup.sqldelight.core.psi.SqldelightTypes
import com.squareup.sqldelight.intellij.util.prevSiblingOfType
import org.jetbrains.kotlin.psi.psiUtil.endOffset
import org.jetbrains.kotlin.psi.psiUtil.getChildOfType

class SqlDelightFoldingBuilder : FoldingBuilder, DumbAware {

  override fun buildFoldRegions(root: ASTNode, document: Document) =
    root.createFoldingDescriptors()

  private fun ASTNode.createFoldingDescriptors(): Array<FoldingDescriptor> {
    return getChildren(null)
      .filter { it.elementType == SqldelightTypes.STMT_LIST }
      .flatMap {
        val descriptors = mutableListOf<FoldingDescriptor>()
        val statements = it.getChildren(null).toList()
        for (statement in statements) {
          when (statement.elementType) {
            SqldelightTypes.IMPORT_STMT_LIST ->
              statement.psi.toImportListDescriptor()?.let(descriptors::add)
            SqlTypes.STMT -> {
              val psi = statement.psi
              val sqlStatement = statement.firstChildNode
              when (sqlStatement?.elementType) {
                SqlTypes.CREATE_TABLE_STMT ->
                  psi.toCreateTableDescriptor(sqlStatement?.psi)?.let(descriptors::add)
                SqlTypes.CREATE_VIEW_STMT ->
                  psi.toCreateViewDescriptor(sqlStatement?.psi)?.let(descriptors::add)
                SqlTypes.CREATE_TRIGGER_STMT ->
                  psi.toCreateTriggerDescriptor(sqlStatement?.psi)?.let(descriptors::add)
                SqlTypes.CREATE_INDEX_STMT ->
                  psi.toCreateIndexDescriptor(sqlStatement?.psi)?.let(descriptors::add)
              }
              val stmtIdentifier = psi.prevSiblingOfType<SqlDelightStmtIdentifier>()
              if (stmtIdentifier?.identifier() != null) {
                psi.toStatementDescriptor(stmtIdentifier)?.let(descriptors::add)
              }
            }
          }
        }
        return@flatMap descriptors
      }
      .toTypedArray()
  }

  private fun PsiElement.toCreateTableDescriptor(createTableStmt: PsiElement?): FoldingDescriptor? {
    val openingBraceElement = createTableStmt?.node?.findChildByType(SqlTypes.LP) ?: return null
    val start = openingBraceElement.startOffset
    val end = nextSibling.endOffset
    if (start >= end) return null
    return FoldingDescriptor(this, start, end, null, "(...);")
  }

  private fun PsiElement.toCreateViewDescriptor(createViewStmt: PsiElement?): FoldingDescriptor? {
    val viewNameElement = (createViewStmt as? SqlCreateViewStmt)?.viewName ?: return null
    return toStatementDescriptor(viewNameElement)
  }

  private fun PsiElement.toCreateTriggerDescriptor(
    createTriggerStmt: PsiElement?
  ): FoldingDescriptor? {
    val triggerNameElement =
      (createTriggerStmt as? SqlCreateTriggerStmt)?.triggerName ?: return null
    return toStatementDescriptor(triggerNameElement)
  }

  private fun PsiElement.toCreateIndexDescriptor(createIndexStmt: PsiElement?): FoldingDescriptor? {
    val indexNameElement = (createIndexStmt as? SqlCreateIndexStmt)?.indexName ?: return null
    return toStatementDescriptor(indexNameElement)
  }

  private fun PsiElement.toStatementDescriptor(stmtIdentifier: PsiElement?): FoldingDescriptor? {
    if (stmtIdentifier == null) return null
    if (nextSibling?.node?.elementType != SqlTypes.SEMI) return null
    val start = stmtIdentifier.endOffset
    val end = nextSibling.endOffset
    if (start >= end) return null
    return FoldingDescriptor(this, start, end, null, "...")
  }

  private fun PsiElement.toImportListDescriptor(): FoldingDescriptor? {
    if (children.size < 2) return null
    val whitespaceElement = firstChild.getChildOfType<PsiWhiteSpace>() ?: return null
    val start = whitespaceElement.endOffset
    val end = lastChild.endOffset
    if (start >= end) return null
    return FoldingDescriptor(this, start, end, null, "...")
  }

  override fun getPlaceholderText(node: ASTNode) = "..."

  override fun isCollapsedByDefault(node: ASTNode) = with(node) {
    elementType == SqldelightTypes.IMPORT_STMT_LIST && getChildren(null).size > 1
  }
}

package com.squareup.sqldelight.intellij.lang

import com.alecstrong.sql.psi.core.psi.AliasElement
import com.alecstrong.sql.psi.core.psi.SqlColumnName
import com.alecstrong.sql.psi.core.psi.SqlFunctionName
import com.alecstrong.sql.psi.core.psi.SqlIdentifier
import com.alecstrong.sql.psi.core.psi.SqlTableName
import com.alecstrong.sql.psi.core.psi.SqlTypeName
import com.alecstrong.sql.psi.core.psi.SqlVisitor
import com.intellij.codeInsight.daemon.impl.HighlightInfo
import com.intellij.codeInsight.daemon.impl.HighlightInfoType
import com.intellij.codeInsight.daemon.impl.HighlightVisitor
import com.intellij.codeInsight.daemon.impl.analysis.HighlightInfoHolder
import com.intellij.openapi.editor.DefaultLanguageHighlighterColors
import com.intellij.openapi.editor.colors.TextAttributesKey
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.psi.SqlDelightStmtIdentifier

class SqlDelightHighlightVisitor : SqlVisitor(), HighlightVisitor {

  private var myHolder: HighlightInfoHolder? = null

  override fun suitableForFile(file: PsiFile): Boolean {
    return file is SqlDelightFile
  }

  override fun visit(element: PsiElement) {
    when (element) {
      is AliasElement -> visitAliasElement(element)
      is SqlColumnName -> visitColumnName(element)
      is SqlFunctionName -> visitFunctionName(element)
      is SqlIdentifier -> visitIdentifier(element)
      is SqlTableName -> visitTableName(element)
      is SqlTypeName -> visitTypeName(element)
    }
  }

  override fun analyze(
    file: PsiFile,
    updateWholeFile: Boolean,
    holder: HighlightInfoHolder,
    action: Runnable
  ): Boolean {
    myHolder = holder
    try {
      action.run()
    } finally {
      myHolder = null
    }
    return true
  }

  override fun clone(): HighlightVisitor {
    return SqlDelightHighlightVisitor()
  }

  override fun visitColumnName(o: SqlColumnName) {
    val info = HighlightInfo.newHighlightInfo(createSymbolTypeInfo(SQL_COLUMN_NAME))
      .range(o)
      .create()
    myHolder?.add(info)
  }

  override fun visitFunctionName(o: SqlFunctionName) {
    val info = HighlightInfo.newHighlightInfo(createSymbolTypeInfo(SQL_FUNCTION_NAME))
      .range(o)
      .create()
    myHolder?.add(info)
  }

  override fun visitIdentifier(o: SqlIdentifier) {
    if (o.parent is SqlDelightStmtIdentifier) {
      val info = HighlightInfo.newHighlightInfo(createSymbolTypeInfo(SQL_STMT_IDENTIFIER_NAME))
        .range(o)
        .create()
      myHolder?.add(info)
    }
  }

  override fun visitTypeName(o: SqlTypeName) {
    val info = HighlightInfo.newHighlightInfo(createSymbolTypeInfo(SQL_TYPE_NAME))
      .range(o)
      .create()
    myHolder?.add(info)
  }

  override fun visitTableName(o: SqlTableName) {
    val info = HighlightInfo.newHighlightInfo(createSymbolTypeInfo(SQL_TABLE_NAME))
      .range(o)
      .create()
    myHolder?.add(info)
  }

  override fun visitAliasElement(o: AliasElement) {
    val info = HighlightInfo.newHighlightInfo(createSymbolTypeInfo(SQL_ALIAS))
      .range(o)
      .create()
    myHolder?.add(info)
  }

  private fun createSymbolTypeInfo(attributesKey: TextAttributesKey): HighlightInfoType {
    return HighlightInfoType.HighlightInfoTypeImpl(
      HighlightInfoType.SYMBOL_TYPE_SEVERITY, attributesKey, false
    )
  }

  companion object {
    val SQL_ALIAS = TextAttributesKey.createTextAttributesKey(
      "SQL.ALIAS",
      DefaultLanguageHighlighterColors.LOCAL_VARIABLE
    )
    val SQL_COLUMN_NAME = TextAttributesKey.createTextAttributesKey(
      "SQL.COLUMN_NAME",
      DefaultLanguageHighlighterColors.INSTANCE_FIELD
    )
    val SQL_FUNCTION_NAME = TextAttributesKey.createTextAttributesKey(
      "SQL.FUNCTION_NAME",
      DefaultLanguageHighlighterColors.STATIC_METHOD
    )
    val SQL_STMT_IDENTIFIER_NAME = TextAttributesKey.createTextAttributesKey(
      "SQL.STMT_IDENTIFIER_NAME",
      DefaultLanguageHighlighterColors.INSTANCE_METHOD
    )
    val SQL_TABLE_NAME = TextAttributesKey.createTextAttributesKey(
      "SQL.TABLE_NAME",
      DefaultLanguageHighlighterColors.CLASS_NAME
    )
    val SQL_TYPE_NAME = TextAttributesKey.createTextAttributesKey(
      "SQL.TYPE_NAME",
      DefaultLanguageHighlighterColors.KEYWORD
    )
  }
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij.lang

import com.alecstrong.sql.psi.core.SqlLexerAdapter
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.intellij.openapi.editor.DefaultLanguageHighlighterColors
import com.intellij.openapi.editor.colors.TextAttributesKey
import com.intellij.openapi.fileTypes.SyntaxHighlighterBase
import com.intellij.psi.tree.IElementType

class SqlDelightHighlighter : SyntaxHighlighterBase() {
  override fun getHighlightingLexer() = SqlLexerAdapter()

  override fun getTokenHighlights(tokenType: IElementType) =
    SyntaxHighlighterBase.pack(TEXT_ATTRIBUTES_MAP[tokenType.index])

  companion object {
    val SQLITE_KEYWORD = TextAttributesKey.createTextAttributesKey(
      "SQLITE.KEYWORD", DefaultLanguageHighlighterColors.KEYWORD
    )
    val SQLITE_NUMBER = TextAttributesKey.createTextAttributesKey(
      "SQLITE.NUMBER", DefaultLanguageHighlighterColors.NUMBER
    )
    val SQLITE_STRING = TextAttributesKey.createTextAttributesKey(
      "SQLITE.STRING", DefaultLanguageHighlighterColors.STRING
    )
    val SQLITE_OPERATOR = TextAttributesKey.createTextAttributesKey(
      "SQLITE.OPERATOR", DefaultLanguageHighlighterColors.OPERATION_SIGN
    )
    val SQLITE_PAREN = TextAttributesKey.createTextAttributesKey(
      "SQLITE.PAREN", DefaultLanguageHighlighterColors.PARENTHESES
    )
    val SQLITE_DOT = TextAttributesKey.createTextAttributesKey(
      "SQLITE.DOT", DefaultLanguageHighlighterColors.DOT
    )
    val SQLITE_SEMICOLON = TextAttributesKey.createTextAttributesKey(
      "SQLITE.SEMICOLON", DefaultLanguageHighlighterColors.SEMICOLON
    )
    val SQLITE_COMMA = TextAttributesKey.createTextAttributesKey(
      "SQLITE.COMMA", DefaultLanguageHighlighterColors.COMMA
    )
    val SQLITE_LINE_COMMENT = TextAttributesKey.createTextAttributesKey(
      "SQLITE.LINE_COMMENT", DefaultLanguageHighlighterColors.LINE_COMMENT
    )
    val SQLITE_DOC = TextAttributesKey.createTextAttributesKey(
      "SQLITE.DOC", DefaultLanguageHighlighterColors.DOC_COMMENT
    )

    private val TEXT_ATTRIBUTES_MAP = mapOf(

      // Keywords

      SqlTypes.ABORT.index to SQLITE_KEYWORD,
      SqlTypes.ACTION.index to SQLITE_KEYWORD,
      SqlTypes.ADD.index to SQLITE_KEYWORD,
      SqlTypes.AFTER.index to SQLITE_KEYWORD,
      SqlTypes.ALL.index to SQLITE_KEYWORD,
      SqlTypes.ALTER.index to SQLITE_KEYWORD,
      SqlTypes.ANALYZE.index to SQLITE_KEYWORD,
      SqlTypes.AND.index to SQLITE_KEYWORD,
      SqlTypes.AS.index to SQLITE_KEYWORD,
      SqlTypes.ASC.index to SQLITE_KEYWORD,
      SqlTypes.ATTACH.index to SQLITE_KEYWORD,
      SqlTypes.AUTOINCREMENT.index to SQLITE_KEYWORD,
      SqlTypes.BEFORE.index to SQLITE_KEYWORD,
      SqlTypes.BEGIN.index to SQLITE_KEYWORD,
      SqlTypes.BETWEEN.index to SQLITE_KEYWORD,
      SqlTypes.BY.index to SQLITE_KEYWORD,
      SqlTypes.CASCADE.index to SQLITE_KEYWORD,
      SqlTypes.CASE.index to SQLITE_KEYWORD,
      SqlTypes.CAST.index to SQLITE_KEYWORD,
      SqlTypes.CHECK.index to SQLITE_KEYWORD,
      SqlTypes.COLLATE.index to SQLITE_KEYWORD,
      SqlTypes.COLUMN.index to SQLITE_KEYWORD,
      SqlTypes.COMMIT.index to SQLITE_KEYWORD,
      SqlTypes.CONFLICT.index to SQLITE_KEYWORD,
      SqlTypes.CONSTRAINT.index to SQLITE_KEYWORD,
      SqlTypes.CREATE.index to SQLITE_KEYWORD,
      SqlTypes.CROSS.index to SQLITE_KEYWORD,
      SqlTypes.CURRENT_DATE.index to SQLITE_KEYWORD,
      SqlTypes.CURRENT_TIME.index to SQLITE_KEYWORD,
      SqlTypes.CURRENT_TIMESTAMP.index to SQLITE_KEYWORD,
      SqlTypes.DATABASE.index to SQLITE_KEYWORD,
      SqlTypes.DEFAULT.index to SQLITE_KEYWORD,
      SqlTypes.DEFERRABLE.index to SQLITE_KEYWORD,
      SqlTypes.DEFERRED.index to SQLITE_KEYWORD,
      SqlTypes.DELETE.index to SQLITE_KEYWORD,
      SqlTypes.DESC.index to SQLITE_KEYWORD,
      SqlTypes.DETACH.index to SQLITE_KEYWORD,
      SqlTypes.DISTINCT.index to SQLITE_KEYWORD,
      SqlTypes.DROP.index to SQLITE_KEYWORD,
      SqlTypes.EACH.index to SQLITE_KEYWORD,
      SqlTypes.ELSE.index to SQLITE_KEYWORD,
      SqlTypes.END.index to SQLITE_KEYWORD,
      SqlTypes.ESCAPE.index to SQLITE_KEYWORD,
      SqlTypes.EXCEPT.index to SQLITE_KEYWORD,
      SqlTypes.EXCLUSIVE.index to SQLITE_KEYWORD,
      SqlTypes.EXISTS.index to SQLITE_KEYWORD,
      SqlTypes.EXPLAIN.index to SQLITE_KEYWORD,
      SqlTypes.FAIL.index to SQLITE_KEYWORD,
      SqlTypes.FOR.index to SQLITE_KEYWORD,
      SqlTypes.FOREIGN.index to SQLITE_KEYWORD,
      SqlTypes.FROM.index to SQLITE_KEYWORD,
      SqlTypes.GLOB.index to SQLITE_KEYWORD,
      SqlTypes.GROUP.index to SQLITE_KEYWORD,
      SqlTypes.HAVING.index to SQLITE_KEYWORD,
      SqlTypes.IF.index to SQLITE_KEYWORD,
      SqlTypes.IGNORE.index to SQLITE_KEYWORD,
      SqlTypes.IMMEDIATE.index to SQLITE_KEYWORD,
      SqlTypes.IN.index to SQLITE_KEYWORD,
      SqlTypes.INDEX.index to SQLITE_KEYWORD,
      SqlTypes.INDEXED.index to SQLITE_KEYWORD,
      SqlTypes.INITIALLY.index to SQLITE_KEYWORD,
      SqlTypes.INNER.index to SQLITE_KEYWORD,
      SqlTypes.INSERT.index to SQLITE_KEYWORD,
      SqlTypes.INSTEAD.index to SQLITE_KEYWORD,
      SqlTypes.INTERSECT.index to SQLITE_KEYWORD,
      SqlTypes.INTO.index to SQLITE_KEYWORD,
      SqlTypes.IS.index to SQLITE_KEYWORD,
      SqlTypes.ISNULL.index to SQLITE_KEYWORD,
      SqlTypes.JOIN.index to SQLITE_KEYWORD,
      SqlTypes.KEY.index to SQLITE_KEYWORD,
      SqlTypes.LEFT.index to SQLITE_KEYWORD,
      SqlTypes.LIKE.index to SQLITE_KEYWORD,
      SqlTypes.LIMIT.index to SQLITE_KEYWORD,
      SqlTypes.MATCH.index to SQLITE_KEYWORD,
      SqlTypes.NATURAL.index to SQLITE_KEYWORD,
      SqlTypes.NO.index to SQLITE_KEYWORD,
      SqlTypes.NOT.index to SQLITE_KEYWORD,
      SqlTypes.NOTNULL.index to SQLITE_KEYWORD,
      SqlTypes.NULL.index to SQLITE_KEYWORD,
      SqlTypes.OF.index to SQLITE_KEYWORD,
      SqlTypes.OFFSET.index to SQLITE_KEYWORD,
      SqlTypes.ON.index to SQLITE_KEYWORD,
      SqlTypes.OR.index to SQLITE_KEYWORD,
      SqlTypes.ORDER.index to SQLITE_KEYWORD,
      SqlTypes.OUTER.index to SQLITE_KEYWORD,
      SqlTypes.PLAN.index to SQLITE_KEYWORD,
      SqlTypes.PRAGMA.index to SQLITE_KEYWORD,
      SqlTypes.PRIMARY.index to SQLITE_KEYWORD,
      SqlTypes.QUERY.index to SQLITE_KEYWORD,
      SqlTypes.RAISE.index to SQLITE_KEYWORD,
      SqlTypes.RECURSIVE.index to SQLITE_KEYWORD,
      SqlTypes.REFERENCES.index to SQLITE_KEYWORD,
      SqlTypes.REGEXP.index to SQLITE_KEYWORD,
      SqlTypes.REINDEX.index to SQLITE_KEYWORD,
      SqlTypes.RELEASE.index to SQLITE_KEYWORD,
      SqlTypes.RENAME.index to SQLITE_KEYWORD,
      SqlTypes.REPLACE.index to SQLITE_KEYWORD,
      SqlTypes.RESTRICT.index to SQLITE_KEYWORD,
      SqlTypes.ROLLBACK.index to SQLITE_KEYWORD,
      SqlTypes.ROW.index to SQLITE_KEYWORD,
      SqlTypes.SAVEPOINT.index to SQLITE_KEYWORD,
      SqlTypes.SELECT.index to SQLITE_KEYWORD,
      SqlTypes.SET.index to SQLITE_KEYWORD,
      SqlTypes.TABLE.index to SQLITE_KEYWORD,
      SqlTypes.TEMP.index to SQLITE_KEYWORD,
      SqlTypes.TEMPORARY.index to SQLITE_KEYWORD,
      SqlTypes.THEN.index to SQLITE_KEYWORD,
      SqlTypes.TO.index to SQLITE_KEYWORD,
      SqlTypes.TRANSACTION.index to SQLITE_KEYWORD,
      SqlTypes.TRIGGER.index to SQLITE_KEYWORD,
      SqlTypes.UNION.index to SQLITE_KEYWORD,
      SqlTypes.UNIQUE.index to SQLITE_KEYWORD,
      SqlTypes.UPDATE.index to SQLITE_KEYWORD,
      SqlTypes.USING.index to SQLITE_KEYWORD,
      SqlTypes.VACUUM.index to SQLITE_KEYWORD,
      SqlTypes.VALUES.index to SQLITE_KEYWORD,
      SqlTypes.VIEW.index to SQLITE_KEYWORD,
      SqlTypes.VIRTUAL.index to SQLITE_KEYWORD,
      SqlTypes.WHEN.index to SQLITE_KEYWORD,
      SqlTypes.WHERE.index to SQLITE_KEYWORD,
      SqlTypes.WITH.index to SQLITE_KEYWORD,
      SqlTypes.WITHOUT.index to SQLITE_KEYWORD,

      // Numbers

      SqlTypes.NUMERIC_LITERAL.index to SQLITE_NUMBER,
      SqlTypes.SIGNED_NUMBER.index to SQLITE_NUMBER,
      SqlTypes.DIGIT.index to SQLITE_NUMBER,

      // Operators

      SqlTypes.COMMA.index to SQLITE_COMMA,
      SqlTypes.DOT.index to SQLITE_DOT,
      SqlTypes.EQ.index to SQLITE_OPERATOR,
      SqlTypes.LP.index to SQLITE_PAREN,
      SqlTypes.RP.index to SQLITE_PAREN,
      SqlTypes.SEMI.index to SQLITE_SEMICOLON,
      SqlTypes.PLUS.index to SQLITE_OPERATOR,
      SqlTypes.MINUS.index to SQLITE_OPERATOR,
      SqlTypes.SHIFT_RIGHT.index to SQLITE_OPERATOR,
      SqlTypes.SHIFT_LEFT.index to SQLITE_OPERATOR,
      SqlTypes.LT.index to SQLITE_OPERATOR,
      SqlTypes.GT.index to SQLITE_OPERATOR,
      SqlTypes.LTE.index to SQLITE_OPERATOR,
      SqlTypes.GTE.index to SQLITE_OPERATOR,
      SqlTypes.EQ2.index to SQLITE_OPERATOR,
      SqlTypes.NEQ.index to SQLITE_OPERATOR,
      SqlTypes.NEQ2.index to SQLITE_OPERATOR,
      SqlTypes.MULTIPLY.index to SQLITE_OPERATOR,
      SqlTypes.DIVIDE.index to SQLITE_OPERATOR,
      SqlTypes.MOD.index to SQLITE_OPERATOR,
      SqlTypes.BITWISE_AND.index to SQLITE_OPERATOR,
      SqlTypes.BITWISE_OR.index to SQLITE_OPERATOR,
      SqlTypes.CONCAT.index to SQLITE_OPERATOR,

      SqlTypes.STRING.index to SQLITE_STRING,

      // Comments

      SqlTypes.COMMENT.index to SQLITE_LINE_COMMENT,
      SqlTypes.JAVADOC.index to SQLITE_DOC
    )
  }
}

package com.squareup.sqldelight.intellij.lang

import com.intellij.lang.ImportOptimizer
import com.intellij.psi.PsiDocumentManager
import com.intellij.psi.PsiFile
import com.intellij.psi.PsiFileFactory
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.SqlDelightFileType
import com.squareup.sqldelight.core.lang.psi.ImportStmtMixin
import com.squareup.sqldelight.core.lang.util.findChildOfType
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.psi.SqlDelightImportStmtList
import com.squareup.sqldelight.intellij.inspections.columnJavaTypes

class SqlDelightImportOptimizer : ImportOptimizer {

  override fun supports(file: PsiFile): Boolean = file is SqlDelightFile

  override fun processFile(file: PsiFile): Runnable = Runnable {
    val manager = PsiDocumentManager.getInstance(file.project)
    val document = manager.getDocument(file) ?: return@Runnable
    manager.commitDocument(document)

    val javaTypes = file.columnJavaTypes()

    val remainingImports = file.findChildrenOfType<ImportStmtMixin>()
      .asSequence()
      .map { import -> import.text }
      .filter { import -> import.substringAfterLast(".").removeSuffix(";") in javaTypes }
      .sorted()
      .joinToString("\n")
    val factory = PsiFileFactory.getInstance(file.project)
    val dummyFile = factory.createFileFromText(
      "_Dummy_.${SqlDelightFileType.EXTENSION}",
      SqlDelightFileType, remainingImports
    )
    val newImportList = dummyFile.findChildOfType<SqlDelightImportStmtList>()
    if (newImportList != null) {
      file.findChildOfType<SqlDelightImportStmtList>()?.replace(newImportList)
    }
  }
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij.lang

import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.intellij.lang.BracePair
import com.intellij.lang.PairedBraceMatcher
import com.intellij.psi.PsiFile
import com.intellij.psi.tree.IElementType

class SqlDelightPairedBraceMatcher : PairedBraceMatcher {

  private val bracePairs = arrayOf(BracePair(SqlTypes.LP, SqlTypes.RP, false))

  override fun getPairs() = bracePairs

  override fun isPairedBracesAllowedBeforeType(
    lbraceType: IElementType,
    contextType: IElementType?
  ) = true

  override fun getCodeConstructStart(file: PsiFile?, openingBraceOffset: Int) = openingBraceOffset
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij.lang

import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.alecstrong.sql.psi.core.psi.SqlCreateVirtualTableStmt
import com.alecstrong.sql.psi.core.psi.SqlIdentifier
import com.alecstrong.sql.psi.core.psi.SqlIndexName
import com.alecstrong.sql.psi.core.psi.SqlStmt
import com.alecstrong.sql.psi.core.psi.SqlTableName
import com.alecstrong.sql.psi.core.psi.SqlTriggerName
import com.alecstrong.sql.psi.core.psi.SqlViewName
import com.intellij.ide.structureView.StructureViewModel
import com.intellij.ide.structureView.StructureViewModelBase
import com.intellij.ide.structureView.StructureViewTreeElement
import com.intellij.ide.structureView.TreeBasedStructureViewBuilder
import com.intellij.ide.util.treeView.smartTree.TreeElement
import com.intellij.lang.PsiStructureViewFactory
import com.intellij.navigation.ItemPresentation
import com.intellij.navigation.NavigationItem
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.util.Iconable
import com.intellij.pom.Navigatable
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.psi.SqlDelightStmtIdentifier
import javax.swing.Icon

class SqlDelightStructureViewFactory : PsiStructureViewFactory {
  override fun getStructureViewBuilder(psiFile: PsiFile) = object : TreeBasedStructureViewBuilder() {
    override fun createStructureViewModel(editor: Editor?) = SqlDelightStructureViewModel(psiFile)
    override fun isRootNodeShown() = false
  }
}

internal class SqlDelightStructureViewModel(
  psiFile: PsiFile
) : StructureViewModelBase(psiFile, SqlDelightStructureViewElement(psiFile)),
  StructureViewModel.ElementInfoProvider {
  override fun isAlwaysLeaf(element: StructureViewTreeElement) = false
  override fun isAlwaysShowsPlus(element: StructureViewTreeElement) = false
}

internal class SqlDelightStructureViewElement(
  private val element: PsiElement
) : StructureViewTreeElement,
  ItemPresentation,
  NavigationItem {
  override fun getValue() = element
  override fun getChildren() = when (element) {
    is SqlDelightFile -> element.childIdentifiers()
    else -> emptyArray()
  }

  override fun canNavigateToSource() = (element as Navigatable).canNavigateToSource()
  override fun canNavigate() = (element as Navigatable).canNavigate()
  override fun navigate(requestFocus: Boolean) = (element as Navigatable).navigate(requestFocus)
  override fun getPresentation() = this
  override fun getName(): String? = when (element) {
    is SqlDelightFile -> element.virtualFile?.name
    is SqlTableName -> when (element.parent) {
      is SqlCreateTableStmt -> "CREATE TABLE ${element.name}"
      is SqlCreateVirtualTableStmt -> "CREATE VIRTUAL TABLE ${element.name}"
      else -> throw IllegalStateException(
        "Unhandled table name element for parent ${element.parent}"
      )
    }
    is SqlIndexName -> "CREATE INDEX ${element.text}"
    is SqlTriggerName -> "CREATE TRIGGER ${element.text}"
    is SqlViewName -> "CREATE VIEW ${element.name}"
    is SqlIdentifier -> element.text
    else -> throw IllegalStateException("Unhandled element $element")
  }

  override fun getLocationString(): String? = null
  override fun getIcon(open: Boolean): Icon? = element.getIcon(Iconable.ICON_FLAG_READ_STATUS)

  override fun getPresentableText() = name

  companion object {
    private fun SqlDelightFile.childIdentifiers(): Array<out TreeElement> {
      return sqlStmtList?.children
        .orEmpty()
        .mapNotNull {
          val element = when (it) {
            is SqlDelightStmtIdentifier -> it.identifier()
            is SqlStmt -> with(it) {
              when {
                createTableStmt != null -> createTableStmt?.tableName
                createIndexStmt != null -> createIndexStmt?.indexName
                createTriggerStmt != null -> createTriggerStmt?.triggerName
                createViewStmt != null -> createViewStmt?.viewName
                createVirtualTableStmt != null -> createVirtualTableStmt?.tableName
                else -> null
              }
            }
            else -> null
          }
          return@mapNotNull element?.let(::SqlDelightStructureViewElement)
        }
        .toTypedArray()
    }
  }
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij.lang

import com.intellij.openapi.fileTypes.SyntaxHighlighterFactory
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.VirtualFile

class SqlDelightSyntaxHighlighterFactory : SyntaxHighlighterFactory() {
  override fun getSyntaxHighlighter(project: Project?, virtualFile: VirtualFile?) =
    SqlDelightHighlighter()
}

package com.squareup.sqldelight.intellij.macro

import com.alecstrong.sql.psi.core.SqlFileBase
import com.intellij.codeInsight.template.Expression
import com.intellij.codeInsight.template.ExpressionContext
import com.intellij.codeInsight.template.ListResult
import com.intellij.codeInsight.template.Result
import com.intellij.codeInsight.template.TemplateContextType
import com.intellij.codeInsight.template.TextResult
import com.intellij.codeInsight.template.macro.MacroBase
import com.intellij.psi.PsiDocumentManager
import com.squareup.sqldelight.intellij.SqlDelightLiveTemplateContextType

class DbColumnsMacro : MacroBase("dbColumns", "dbColumns(tableName)") {

  override fun isAcceptableInContext(context: TemplateContextType?): Boolean {
    return context is SqlDelightLiveTemplateContextType
  }

  override fun calculateResult(
    params: Array<out Expression>,
    context: ExpressionContext,
    quick: Boolean
  ): Result? {
    if (params.isEmpty()) {
      return null
    }
    val expression = params.first()
    val tableName = expression.calculateResult(context) ?: return null
    val document = context.editor?.document ?: return null
    val psiFile = PsiDocumentManager.getInstance(context.project)
      .getPsiFile(document) as? SqlFileBase ?: return null
    val tables = psiFile.tables(true)
    val table = tables.firstOrNull { it.tableName.name == tableName.toString() } ?: return null
    val columns = table.query.columns
      .map { TextResult(it.element.text) }
    return ListResult(columns)
  }
}

package com.squareup.sqldelight.intellij.macro

import com.alecstrong.sql.psi.core.SqlFileBase
import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.intellij.codeInsight.template.Expression
import com.intellij.codeInsight.template.ExpressionContext
import com.intellij.codeInsight.template.Result
import com.intellij.codeInsight.template.TextResult
import com.intellij.codeInsight.template.macro.MacroBase
import com.intellij.psi.PsiDocumentManager
import com.squareup.sqldelight.core.lang.util.findChildrenOfType

class PkMacro : MacroBase("pk", "pk(tableName)") {

  override fun calculateResult(
    params: Array<out Expression>,
    context: ExpressionContext,
    quick: Boolean
  ): Result? {
    if (params.isEmpty()) {
      return null
    }

    val expression = params.first()
    val tableName = expression.calculateResult(context) ?: return null
    val document = context.editor?.document ?: return null
    val psiFile = PsiDocumentManager.getInstance(context.project)
      .getPsiFile(document) as? SqlFileBase ?: return null

    val createTableStmt = psiFile.sqlStmtList?.findChildrenOfType<SqlCreateTableStmt>()
      .orEmpty()
      .firstOrNull { it.tableName.name == tableName.toString() } ?: return null

    val pk = createTableStmt.columnDefList.firstOrNull { columnDef ->
      columnDef.columnConstraintList.any { columnConstraint ->
        columnConstraint.textMatches("PRIMARY KEY")
      }
    }?.columnName
    return pk?.let { TextResult(it.name) }
  }
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij.util

import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.module.Module
import com.intellij.openapi.vfs.VirtualFile
import com.squareup.sqldelight.core.SqlDelightFileIndex
import kotlin.reflect.KProperty

class GeneratedVirtualFile(private val path: String, module: Module) {
  private val applicationManager = ApplicationManager.getApplication()
  private val index = SqlDelightFileIndex.getInstance(module)
  private var backingFile: VirtualFile? = null

  operator fun getValue(
    thisRef: Nothing?,
    property: KProperty<*>
  ): VirtualFile {
    applicationManager.assertWriteAccessAllowed()
    synchronized(this) {
      val backingFile = this.backingFile
      if (backingFile == null || !backingFile.exists()) {
        var file = index.contentRoot.findFileByRelativePath(path)
        if (file == null || !file.exists()) {
          file = getOrCreateFile(path)
        }
        check(file.exists()) { "VirtualFile $path still doesn't exist after creating." }
        this.backingFile = file
        return file
      } else {
        return backingFile
      }
    }
  }

  private fun getOrCreateFile(path: String): VirtualFile {
    val indexOfName = path.lastIndexOf('/')
    val parent = getOrCreateDirectory(path.substring(0, indexOfName))
    return parent.findOrCreateChildData(this, path.substring(indexOfName + 1, path.length))
  }

  private fun getOrCreateDirectory(path: String): VirtualFile {
    val indexOfName = path.lastIndexOf('/')
    if (indexOfName == -1) {
      return index.contentRoot.findChild(path) ?: index.contentRoot.createChildDirectory(this, path)
    }
    val parentPath = path.substring(0, indexOfName)
    var parent = index.contentRoot.findFileByRelativePath(parentPath)
    if (parent == null || !parent.exists()) {
      parent = getOrCreateDirectory(parentPath)
      if (!parent.exists() || !parent.isDirectory) throw AssertionError()
    }
    val name = path.substring(indexOfName + 1, path.length)
    return parent.findChild(name) ?: parent.createChildDirectory(this, name)
  }
}

package com.squareup.sqldelight.intellij.util

import com.intellij.openapi.project.Project
import com.intellij.psi.JavaPsiFacade
import com.intellij.psi.PsiClass
import com.intellij.psi.search.GlobalSearchScope
import com.intellij.psi.search.PsiShortNamesCache
import org.jetbrains.kotlin.idea.refactoring.memberInfo.qualifiedClassNameForRendering
import org.jetbrains.kotlin.idea.stubindex.KotlinClassShortNameIndex

object PsiClassSearchHelper {

  fun getClassesByShortName(shortName: String, project: Project, scope: GlobalSearchScope): List<PsiClass> {
    val javaPsiFacade = JavaPsiFacade.getInstance(project)

    val kotlinClasses = KotlinClassShortNameIndex.getInstance()
      .get(shortName, project, scope)
      .mapNotNull { javaPsiFacade.findClass(it.qualifiedClassNameForRendering(), scope) }

    val javaClasses = PsiShortNamesCache.getInstance(project)
      .getClassesByName(shortName, scope)
    return (kotlinClasses + javaClasses)
      .distinctBy { it.qualifiedName }
  }
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.intellij.util

import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import com.intellij.psi.util.PsiTreeUtil

internal inline fun <reified T : PsiElement> PsiElement.prevSiblingOfType(): T? =
  PsiTreeUtil.getPrevSiblingOfType(this, T::class.java)

internal inline fun <reified T : PsiElement> PsiFile.findElementOfTypeAtOffset(
  offset: Int
): T? {
  return PsiTreeUtil.findElementOfClassAtOffset(this, offset, T::class.java, false)
}

package com.squareup.sqldelight.intellij.util

import com.intellij.openapi.vfs.VirtualFile

internal fun VirtualFile.isAncestorOf(child: VirtualFile): Boolean {
  if (child in children) return true
  return isAncestorOf(child.parent ?: return false)
}

@file:JvmName("JdbcDrivers")
package com.squareup.sqldelight.sqlite.driver

import com.squareup.sqldelight.Transacter
import com.squareup.sqldelight.db.SqlCursor
import com.squareup.sqldelight.db.SqlDriver
import com.squareup.sqldelight.db.SqlPreparedStatement
import com.squareup.sqldelight.sqlite.driver.ConnectionManager.Transaction
import java.sql.Connection
import java.sql.PreparedStatement
import java.sql.ResultSet
import java.sql.Types
import javax.sql.DataSource

@JvmName("fromDataSource")
fun DataSource.asJdbcDriver() = object : JdbcDriver() {
  override fun getConnection(): Connection {
    return connection
  }

  override fun closeConnection(connection: Connection) {
    connection.close()
  }
}

interface ConnectionManager {
  fun close()

  fun getConnection(): Connection

  fun closeConnection(connection: Connection)

  fun Connection.beginTransaction()

  fun Connection.endTransaction()

  fun Connection.rollbackTransaction()

  var transaction: Transaction?

  class Transaction(
    override val enclosingTransaction: Transaction?,
    private val connectionManager: ConnectionManager,
    val connection: Connection
  ) : Transacter.Transaction() {
    override fun endTransaction(successful: Boolean) {
      if (enclosingTransaction == null) {
        if (successful) connectionManager.apply { connection.endTransaction() }
        else connectionManager.apply { connection.rollbackTransaction() }
      }
      connectionManager.transaction = enclosingTransaction
    }
  }
}

abstract class JdbcDriver : SqlDriver, ConnectionManager {
  override fun close() {
  }

  override fun Connection.endTransaction() {
    commit()
    autoCommit = true
    closeConnection(this)
  }

  override fun Connection.rollbackTransaction() {
    rollback()
    autoCommit = true
    closeConnection(this)
  }

  override fun Connection.beginTransaction() {
    check(autoCommit) {
      """
      Expected autoCommit to be true by default. For compatibility with SQLDelight make sure it is
      set to true when returning a connection from [JdbcDriver.getConnection()]
      """.trimIndent()
    }
    autoCommit = false
  }

  private val transactions = ThreadLocal<Transaction>()

  override var transaction: Transaction?
    get() = transactions.get()
    set(value) { transactions.set(value) }

  private fun connectionAndClose(): Pair<Connection, () -> Unit> {
    val enclosing = transaction
    return if (enclosing != null) {
      enclosing.connection to {}
    } else {
      val connection = getConnection()
      return connection to { closeConnection(connection) }
    }
  }

  override fun execute(
    identifier: Int?,
    sql: String,
    parameters: Int,
    binders: (SqlPreparedStatement.() -> Unit)?
  ) {
    val (connection, onClose) = connectionAndClose()
    try {
      connection.prepareStatement(sql).use { jdbcStatement ->
        SqliteJdbcPreparedStatement(jdbcStatement)
          .apply { if (binders != null) this.binders() }
          .execute()
      }
    } finally {
      onClose()
    }
  }

  override fun executeQuery(
    identifier: Int?,
    sql: String,
    parameters: Int,
    binders: (SqlPreparedStatement.() -> Unit)?
  ): SqlCursor {
    val (connection, onClose) = connectionAndClose()
    try {
      return SqliteJdbcPreparedStatement(connection.prepareStatement(sql))
        .apply { if (binders != null) this.binders() }
        .executeQuery(onClose)
    } catch (e: Exception) {
      onClose()
      throw e
    }
  }

  override fun newTransaction(): Transacter.Transaction {
    val enclosing = transaction
    val connection = enclosing?.connection ?: getConnection()
    val transaction = Transaction(enclosing, this, connection)
    this.transaction = transaction

    if (enclosing == null) {
      connection.beginTransaction()
    }

    return transaction
  }

  override fun currentTransaction(): Transacter.Transaction? = transaction
}

private class SqliteJdbcPreparedStatement(
  private val preparedStatement: PreparedStatement
) : SqlPreparedStatement {
  override fun bindBytes(index: Int, bytes: ByteArray?) {
    if (bytes == null) {
      preparedStatement.setNull(index, Types.BLOB)
    } else {
      preparedStatement.setBytes(index, bytes)
    }
  }

  override fun bindLong(index: Int, long: Long?) {
    if (long == null) {
      preparedStatement.setNull(index, Types.INTEGER)
    } else {
      preparedStatement.setLong(index, long)
    }
  }

  override fun bindDouble(index: Int, double: Double?) {
    if (double == null) {
      preparedStatement.setNull(index, Types.REAL)
    } else {
      preparedStatement.setDouble(index, double)
    }
  }

  override fun bindString(index: Int, string: String?) {
    if (string == null) {
      preparedStatement.setNull(index, Types.VARCHAR)
    } else {
      preparedStatement.setString(index, string)
    }
  }

  fun executeQuery(onClose: () -> Unit) =
    SqliteJdbcCursor(preparedStatement, preparedStatement.executeQuery(), onClose)

  fun execute() {
    preparedStatement.execute()
  }
}

private class SqliteJdbcCursor(
  private val preparedStatement: PreparedStatement,
  private val resultSet: ResultSet,
  private val onClose: () -> Unit
) : SqlCursor {
  override fun getString(index: Int) = resultSet.getString(index + 1)
  override fun getBytes(index: Int) = resultSet.getBytes(index + 1)
  override fun getLong(index: Int): Long? {
    return resultSet.getLong(index + 1).takeUnless { resultSet.wasNull() }
  }
  override fun getDouble(index: Int): Double? {
    return resultSet.getDouble(index + 1).takeUnless { resultSet.wasNull() }
  }
  override fun close() {
    resultSet.close()
    preparedStatement.close()
    onClose()
  }
  override fun next() = resultSet.next()
}

package com.squareup.sqldelight.drivers.native

internal interface Borrowed<T> {
  val value: T
  fun release()
}

package com.squareup.sqldelight.drivers.native

import co.touchlab.sqliter.DatabaseConfiguration
import co.touchlab.sqliter.DatabaseConnection
import co.touchlab.sqliter.DatabaseManager
import co.touchlab.sqliter.Statement
import co.touchlab.sqliter.createDatabaseManager
import co.touchlab.sqliter.withStatement
import co.touchlab.stately.concurrency.ThreadLocalRef
import co.touchlab.stately.concurrency.value
import com.squareup.sqldelight.Transacter
import com.squareup.sqldelight.db.Closeable
import com.squareup.sqldelight.db.SqlCursor
import com.squareup.sqldelight.db.SqlDriver
import com.squareup.sqldelight.db.SqlPreparedStatement
import com.squareup.sqldelight.drivers.native.util.nativeCache
import kotlin.native.concurrent.AtomicInt
import kotlin.native.concurrent.ensureNeverFrozen

sealed class ConnectionWrapper : SqlDriver {
  internal abstract fun <R> accessConnection(
    readOnly: Boolean,
    block: ThreadConnection.() -> R
  ): R

  final override fun execute(
    identifier: Int?,
    sql: String,
    parameters: Int,
    binders: (SqlPreparedStatement.() -> Unit)?
  ) {
    accessConnection(false) {
      val statement = useStatement(identifier, sql)
      if (binders != null) {
        try {
          SqliterStatement(statement).binders()
        } catch (t: Throwable) {
          statement.resetStatement()
          clearIfNeeded(identifier, statement)
          throw t
        }
      }

      statement.execute()
      statement.resetStatement()
      clearIfNeeded(identifier, statement)
    }
  }

  final override fun executeQuery(
    identifier: Int?,
    sql: String,
    parameters: Int,
    binders: (SqlPreparedStatement.() -> Unit)?
  ): SqlCursor {
    return accessConnection(true) {
      val statement = getStatement(identifier, sql)

      if (binders != null) {
        try {
          SqliterStatement(statement).binders()
        } catch (t: Throwable) {
          statement.resetStatement()
          safePut(identifier, statement)
          throw t
        }
      }

      val cursor = statement.query()

      SqliterSqlCursor(cursor) {
        statement.resetStatement()
        if (closed)
          statement.finalizeStatement()
        safePut(identifier, statement)
      }
    }
  }
}

/**
 * Native driver implementation.
 *
 * The driver creates two connection pools, which default to 1 connection maximum. There is a reader pool, which
 * handles all query requests outside of a transaction. The other pool is the transaction pool, which handles
 * all transactions and write requests outside of a transaction.
 *
 * When a transaction is started, that thread is aligned with a transaction pool connection. Attempting a write or
 * starting another transaction, if no connections are available, will cause the caller to wait.
 *
 * You can have multiple connections in the transaction pool, but this would only be useful for read transactions. Writing
 * from multiple connections in an overlapping manner can be problematic.
 *
 * Aligning a transaction to a thread means you cannot operate on a single transaction from multiple threads.
 * However, it would be difficult to find a use case where this would be desirable or safe. Currently, the native
 * implementation of kotlinx.coroutines does not use thread pooling. When that changes, we'll need a way to handle
 * transaction/connection alignment similar to what the Android/JVM driver implemented.
 *
 * https://medium.com/androiddevelopers/threading-models-in-coroutines-and-android-sqlite-api-6cab11f7eb90
 *
 * To use SqlDelight during create/upgrade processes, you can alternatively wrap a real connection
 * with wrapConnection.
 *
 * SqlPreparedStatement instances also do not point to real resources until either execute or
 * executeQuery is called. The SqlPreparedStatement structure also maintains a thread-aligned
 * instance which accumulates bind calls. Those are replayed on a real SQLite statement instance
 * when execute or executeQuery is called. This avoids race conditions with bind calls.
 */
class NativeSqliteDriver(
  private val databaseManager: DatabaseManager,
  maxReaderConnections: Int = 1,
) : ConnectionWrapper(), SqlDriver {
  constructor(
    configuration: DatabaseConfiguration,
    maxReaderConnections: Int = 1
  ) : this(
    databaseManager = createDatabaseManager(configuration),
    maxReaderConnections = maxReaderConnections
  )

  constructor(
    schema: SqlDriver.Schema,
    name: String,
    maxReaderConnections: Int = 1
  ) : this(
    configuration = DatabaseConfiguration(
      name = name,
      version = schema.version,
      create = { connection ->
        wrapConnection(connection) { schema.create(it) }
      },
      upgrade = { connection, oldVersion, newVersion ->
        wrapConnection(connection) { schema.migrate(it, oldVersion, newVersion) }
      }
    ),
    maxReaderConnections = maxReaderConnections
  )

  // A pool of reader connections used by all operations not in a transaction
  internal val transactionPool: Pool<ThreadConnection>
  internal val readerPool: Pool<ThreadConnection>

  // Once a transaction is started and connection borrowed, it will be here, but only for that
  // thread
  private val borrowedConnectionThread = ThreadLocalRef<Borrowed<ThreadConnection>>()

  init {
    // Single connection for transactions
    transactionPool = Pool(1) {
      ThreadConnection(databaseManager.createMultiThreadedConnection()) { conn ->
        borrowedConnectionThread.let {
          it.get()?.release()
          it.value = null
        }
      }
    }

    val maxReaderConnectionsForConfig: Int = when {
      databaseManager.configuration.inMemory -> 1 // Memory db's are single connection, generally. You can use named connections, but there are other issues that need to be designed for
      else -> maxReaderConnections
    }
    readerPool = Pool(maxReaderConnectionsForConfig) {
      val connection = databaseManager.createMultiThreadedConnection()
      connection.withStatement("PRAGMA query_only = 1") { execute() } // Ensure read only
      ThreadConnection(connection) {
        throw UnsupportedOperationException("Should never be in a transaction")
      }
    }
  }

  override fun currentTransaction(): Transacter.Transaction? {
    return borrowedConnectionThread.get()?.value?.transaction?.value
  }

  override fun newTransaction(): Transacter.Transaction {
    val alreadyBorrowed = borrowedConnectionThread.get()
    return if (alreadyBorrowed == null) {
      val borrowed = transactionPool.borrowEntry()

      try {
        val trans = borrowed.value.newTransaction()

        borrowedConnectionThread.value = borrowed
        trans
      } catch (e: Throwable) {
        // Unlock on failure.
        borrowed.release()
        throw e
      }
    } else {
      alreadyBorrowed.value.newTransaction()
    }
  }

  /**
   * If we're in a transaction, then I have a connection. Otherwise use shared.
   */
  override fun <R> accessConnection(
    readOnly: Boolean,
    block: ThreadConnection.() -> R
  ): R {
    val mine = borrowedConnectionThread.get()

    return if (readOnly) {
      // Code intends to read, which doesn't need to block
      if (mine != null) {
        mine.value.block()
      } else {
        readerPool.access(block)
      }
    } else {
      // Code intends to write, for which we're managing locks in code
      if (mine != null) {
        mine.value.block()
      } else {
        transactionPool.access(block)
      }
    }
  }

  override fun close() {
    transactionPool.close()
    readerPool.close()
  }
}

/**
 * Sqliter's DatabaseConfiguration takes lambda arguments for it's create and upgrade operations,
 * which each take a DatabaseConnection argument. Use wrapConnection to have SqlDelight access this
 * passed connection and avoid the pooling that the full SqlDriver instance performs.
 *
 * Note that queries created during this operation will be cleaned up. If holding onto a cursor from
 * a wrap call, it will no longer be viable.
 */
fun wrapConnection(
  connection: DatabaseConnection,
  block: (SqlDriver) -> Unit
) {
  val conn = SqliterWrappedConnection(ThreadConnection(connection) {})
  try {
    block(conn)
  } finally {
    conn.close()
  }
}

/**
 * SqlDriverConnection that wraps a Sqliter connection. Useful for migration tasks, or if you
 * don't want the polling.
 */
internal class SqliterWrappedConnection(
  private val threadConnection: ThreadConnection
) : ConnectionWrapper(),
  SqlDriver {
  override fun currentTransaction(): Transacter.Transaction? = threadConnection.transaction.value

  override fun newTransaction(): Transacter.Transaction = threadConnection.newTransaction()

  override fun <R> accessConnection(
    readOnly: Boolean,
    block: ThreadConnection.() -> R
  ): R = threadConnection.block()

  override fun close() {
    threadConnection.cleanUp()
  }
}

/**
 * Wraps and manages a "real" database connection.
 *
 * SQLite statements are specific to connections, and must be finalized explicitly. Cursors are
 * backed by a statement resource, so we keep links to open cursors to allow us to close them out
 * properly in cases where the user does not.
 */
internal class ThreadConnection(
  private val connection: DatabaseConnection,
  private val onEndTransaction: (ThreadConnection) -> Unit
) : Closeable {

  companion object {
    private val idCounter = AtomicInt(0)
  }

  internal val transaction = ThreadLocalRef<Transacter.Transaction?>()
  internal val closed: Boolean
    get() = connection.closed

  internal val statementCache = nativeCache<Statement>()
  internal val connectionId: Int = idCounter.addAndGet(1)

  fun safePut(identifier: Int?, statement: Statement) {
    val removed = if (identifier == null) {
      statement
    } else {
      statementCache.put(identifier.toString(), statement)
    }
    removed?.finalizeStatement()
  }

  fun getStatement(identifier: Int?, sql: String): Statement {
    val statement = removeCreateStatement(identifier, sql)
    return statement
  }

  fun useStatement(identifier: Int?, sql: String): Statement {
    return if (identifier != null) {
      statementCache.getOrCreate(identifier.toString()) {
        connection.createStatement(sql)
      }
    } else {
      connection.createStatement(sql)
    }
  }

  fun clearIfNeeded(identifier: Int?, statement: Statement) {
    if (identifier == null) {
      statement.finalizeStatement()
    }
  }

  /**
   * For cursors. Cursors are actually backed by SQLite statement instances, so they need to be
   * removed from the cache when in use. We're giving out a SQLite resource here, so extra care.
   */
  private fun removeCreateStatement(identifier: Int?, sql: String): Statement {
    if (identifier != null) {
      val cached = statementCache.remove(identifier.toString())
      if (cached != null)
        return cached
    }

    return connection.createStatement(sql)
  }

  fun newTransaction(): Transacter.Transaction {
    val enclosing = transaction.value

    // Create here, in case we bomb...
    if (enclosing == null) {
      connection.beginTransaction()
    }

    val trans = Transaction(enclosing)
    transaction.value = trans

    return trans
  }

  /**
   * This should only be called directly from wrapConnection. Clean resources without actually closing
   * the underlying connection.
   */
  internal fun cleanUp() {
    statementCache.cleanUp {
      it.finalizeStatement()
    }
  }

  override fun close() {
    cleanUp()
    connection.close()
  }

  private inner class Transaction(
    override val enclosingTransaction: Transacter.Transaction?
  ) : Transacter.Transaction() {
    init { ensureNeverFrozen() }

    override fun endTransaction(successful: Boolean) {
      transaction.value = enclosingTransaction

      if (enclosingTransaction == null) {
        try {
          if (successful) {
            connection.setTransactionSuccessful()
          }

          connection.endTransaction()
        } finally {
          // Release if we have
          onEndTransaction(this@ThreadConnection)
        }
      }
    }
  }
}

package com.squareup.sqldelight.drivers.native

import co.touchlab.stately.concurrency.AtomicBoolean
import com.squareup.sqldelight.db.Closeable
import com.squareup.sqldelight.drivers.native.util.PoolLock
import kotlin.native.concurrent.AtomicReference
import kotlin.native.concurrent.freeze

/**
 * A shared pool of connections. Borrowing is blocking when all connections are in use, and the pool has reached its
 * designated capacity.
 */
internal class Pool<T : Closeable>(internal val capacity: Int, private val producer: () -> T) {
  /**
   * Hold a list of active connections. If it is null, it means the MultiPool has been closed.
   */
  private val entriesRef = AtomicReference<List<Entry>?>(listOf<Entry>().freeze())
  private val poolLock = PoolLock()

  /**
   * For test purposes only
   */
  internal fun entryCount(): Int = poolLock.withLock {
    entriesRef.value?.size ?: 0
  }

  fun borrowEntry(): Borrowed<T> {
    val snapshot = entriesRef.value ?: throw ClosedMultiPoolException

    // Fastpath: Borrow the first available entry.
    val firstAvailable = snapshot.firstOrNull { it.tryToAcquire() }

    if (firstAvailable != null) {
      return firstAvailable.asBorrowed(poolLock)
    }

    // Slowpath: Create a new entry if capacity limit has not been reached, or wait for the next available entry.
    val nextAvailable = poolLock.withLock {
      // Reload the list since it could've been updated by other threads concurrently.
      val entries = entriesRef.value ?: throw ClosedMultiPoolException

      if (entries.count() < capacity) {
        // Capacity hasn't been reached — create a new entry to serve this call.
        val newEntry = Entry(producer())
        val done = newEntry.tryToAcquire()
        check(done)

        entriesRef.value = (entries + listOf(newEntry)).freeze()
        return@withLock newEntry
      } else {
        // Capacity is reached — wait for the next available entry.
        return@withLock loopForConditionalResult {
          // Reload the list, since the thread can be suspended here while the list of entries has been modified.
          val innerEntries = entriesRef.value ?: throw ClosedMultiPoolException
          innerEntries.firstOrNull { it.tryToAcquire() }
        }
      }
    }

    return nextAvailable.asBorrowed(poolLock)
  }

  fun <R> access(action: (T) -> R): R {
    val borrowed = borrowEntry()
    return try {
      action(borrowed.value)
    } finally {
      borrowed.release()
    }
  }

  fun close() {
    if (!poolLock.close())
      return

    val entries = entriesRef.value
    val done = entriesRef.compareAndSet(entries, null)
    check(done)

    entries?.forEach { it.value.close() }
  }

  inner class Entry(val value: T) {
    val isAvailable = AtomicBoolean(true)

    init { freeze() }

    fun tryToAcquire(): Boolean = isAvailable.compareAndSet(expected = true, new = false)

    fun asBorrowed(poolLock: PoolLock): Borrowed<T> = object : Borrowed<T> {
      override val value: T
        get() = this@Entry.value

      override fun release() {
        /**
         * Mark-as-available should be done before signalling blocked threads via [PoolLock.notifyConditionChanged],
         * since the happens-before relationship guarantees the woken thread to see the
         * available entry (if not having been taken by other threads during the wake-up lead time).
         */

        val done = isAvailable.compareAndSet(expected = false, new = true)
        check(done)

        // While signalling blocked threads does not require locking, doing so avoids a subtle race
        // condition in which:
        //
        // 1. a [loopForConditionalResult] iteration in [borrowEntry] slow path is happening concurrently;
        // 2. the iteration fails to see the atomic `isAvailable = true` above;
        // 3. we signal availability here but it is a no-op due to no waiting blocker; and finally
        // 4. the iteration entered an indefinite blocking wait, not being aware of us having signalled availability here.
        //
        // By acquiring the pool lock first, signalling cannot happen concurrently with the loop
        // iterations in [borrowEntry], thus eliminating the race condition.
        poolLock.withLock {
          poolLock.notifyConditionChanged()
        }
      }
    }
  }
}

private val ClosedMultiPoolException get() = IllegalStateException("Attempt to access a closed MultiPool.")

package com.squareup.sqldelight.drivers.native

import co.touchlab.sqliter.Statement
import co.touchlab.sqliter.bindBlob
import co.touchlab.sqliter.bindDouble
import co.touchlab.sqliter.bindLong
import co.touchlab.sqliter.bindString
import com.squareup.sqldelight.db.SqlPreparedStatement

/**
 * @param [recycle] A function which recycles any resources this statement is backed by.
 */
internal class SqliterStatement(
  private val statement: Statement
) : SqlPreparedStatement {
  override fun bindBytes(index: Int, bytes: ByteArray?) {
    statement.bindBlob(index, bytes)
  }

  override fun bindLong(index: Int, long: Long?) {
    statement.bindLong(index, long)
  }

  override fun bindDouble(index: Int, double: Double?) {
    statement.bindDouble(index, double)
  }

  override fun bindString(index: Int, string: String?) {
    statement.bindString(index, string)
  }
}

package com.squareup.sqldelight.drivers.native

import co.touchlab.sqliter.Cursor
import co.touchlab.sqliter.getBytesOrNull
import co.touchlab.sqliter.getDoubleOrNull
import co.touchlab.sqliter.getLongOrNull
import co.touchlab.sqliter.getStringOrNull
import com.squareup.sqldelight.db.SqlCursor

/**
 * Wrapper for cursor calls. Cursors point to real SQLite statements, so we need to be careful with
 * them. If dev closes the outer structure, this will get closed as well, which means it could start
 * throwing errors if you're trying to access it.
 */
internal class SqliterSqlCursor(
  private val cursor: Cursor,
  private val recycler: () -> Unit
) : SqlCursor {

  override fun close() {
    recycler()
  }

  override fun getBytes(index: Int): ByteArray? = cursor.getBytesOrNull(index)

  override fun getDouble(index: Int): Double? = cursor.getDoubleOrNull(index)

  override fun getLong(index: Int): Long? = cursor.getLongOrNull(index)

  override fun getString(index: Int): String? = cursor.getStringOrNull(index)

  override fun next(): Boolean = cursor.next()
}

package com.squareup.sqldelight.sqlite.driver

import com.squareup.sqldelight.sqlite.driver.ConnectionManager.Transaction
import com.squareup.sqldelight.sqlite.driver.JdbcSqliteDriver.Companion.IN_MEMORY
import java.sql.Connection
import java.sql.DriverManager
import java.util.Properties
import kotlin.concurrent.getOrSet

class JdbcSqliteDriver constructor(
  /**
   * Database connection URL in the form of `jdbc:sqlite:path` where `path` is either blank
   * (creating an in-memory database) or a path to a file.
   */
  url: String,
  properties: Properties = Properties()
) : JdbcDriver(), ConnectionManager by connectionManager(url, properties) {
  companion object {
    const val IN_MEMORY = "jdbc:sqlite:"
  }
}

private fun connectionManager(url: String, properties: Properties) = when (url) {
  IN_MEMORY -> InMemoryConnectionManager(properties)
  else -> ThreadedConnectionManager(url, properties)
}

private abstract class JdbcSqliteDriverConnectionManager : ConnectionManager {
  override fun Connection.beginTransaction() {
    prepareStatement("BEGIN TRANSACTION").execute()
  }

  override fun Connection.endTransaction() {
    prepareStatement("END TRANSACTION").execute()
  }

  override fun Connection.rollbackTransaction() {
    prepareStatement("ROLLBACK TRANSACTION").execute()
  }
}

private class InMemoryConnectionManager(
  properties: Properties
) : JdbcSqliteDriverConnectionManager() {
  override var transaction: Transaction? = null
  private val connection = DriverManager.getConnection(IN_MEMORY, properties)

  override fun getConnection() = connection
  override fun closeConnection(connection: Connection) = Unit
  override fun close() = connection.close()
}

private class ThreadedConnectionManager(
  private val url: String,
  private val properties: Properties,
) : JdbcSqliteDriverConnectionManager() {
  private val transactions = ThreadLocal<Transaction>()
  private val connections = ThreadLocal<Connection>()

  override var transaction: Transaction?
    get() = transactions.get()
    set(value) { transactions.set(value) }

  override fun getConnection() = connections.getOrSet {
    DriverManager.getConnection(url, properties)
  }

  override fun closeConnection(connection: Connection) {
    check(connections.get() == connection) { "Connections must be closed on the thread that opened them" }
    if (transaction == null) {
      connection.close()
      connections.remove()
    }
  }

  override fun close() = Unit
}

package com.squareup.sqldelight.drivers.sqljs

import com.squareup.sqldelight.Transacter
import com.squareup.sqldelight.TransacterImpl
import com.squareup.sqldelight.db.SqlCursor
import com.squareup.sqldelight.db.SqlDriver
import com.squareup.sqldelight.db.SqlPreparedStatement
import org.khronos.webgl.Int8Array
import org.khronos.webgl.Uint8Array
import kotlin.js.Promise

fun Promise<Database>.driver(): Promise<SqlDriver> = then { JsSqlDriver(it) }

fun Promise<SqlDriver>.withSchema(schema: SqlDriver.Schema? = null): Promise<SqlDriver> = then {
  schema?.create(it)
  it
}

fun Promise<SqlDriver>.transacter(): Promise<Transacter> = then { object : TransacterImpl(it) {} }

fun initSqlDriver(schema: SqlDriver.Schema? = null): Promise<SqlDriver> = initDb().driver().withSchema(schema)

class JsSqlDriver(private val db: Database) : SqlDriver {

  private val statements = mutableMapOf<Int, Statement>()
  private var transaction: Transacter.Transaction? = null

  override fun executeQuery(
    identifier: Int?,
    sql: String,
    parameters: Int,
    binders: (SqlPreparedStatement.() -> Unit)?
  ): SqlCursor = createOrGetStatement(identifier, sql).run {
    bind(binders)
    JsSqlCursor(this)
  }

  override fun execute(identifier: Int?, sql: String, parameters: Int, binders: (SqlPreparedStatement.() -> Unit)?) =
    createOrGetStatement(identifier, sql).run {
      bind(binders)
      step()
      freemem()
    }

  private fun Statement.bind(binders: (SqlPreparedStatement.() -> Unit)?) = binders?.let {
    val bound = JsSqlPreparedStatement()
    binders(bound)
    bind(bound.parameters.toTypedArray())
  }

  private fun createOrGetStatement(identifier: Int?, sql: String): Statement = if (identifier == null) {
    db.prepare(sql)
  } else {
    statements.getOrPut(identifier, { db.prepare(sql) }).apply { reset() }
  }

  override fun newTransaction(): Transacter.Transaction {
    val enclosing = transaction
    val transaction = Transaction(enclosing)
    this.transaction = transaction
    if (enclosing == null) {
      db.run("BEGIN TRANSACTION")
    }
    return transaction
  }

  override fun currentTransaction() = transaction

  override fun close() = db.close()

  private inner class Transaction(
    override val enclosingTransaction: Transacter.Transaction?
  ) : Transacter.Transaction() {
    override fun endTransaction(successful: Boolean) {
      if (enclosingTransaction == null) {
        if (successful) {
          db.run("END TRANSACTION")
        } else {
          db.run("ROLLBACK TRANSACTION")
        }
      }
      transaction = enclosingTransaction
    }
  }
}

private class JsSqlCursor(private val statement: Statement) : SqlCursor {
  override fun next(): Boolean = statement.step()
  override fun getString(index: Int): String? = statement.get()[index]
  override fun getLong(index: Int): Long? = (statement.get()[index] as? Double)?.toLong()
  override fun getBytes(index: Int): ByteArray? = (statement.get()[index] as? Uint8Array)?.let {
    Int8Array(it.buffer).unsafeCast<ByteArray>()
  }
  override fun getDouble(index: Int): Double? = statement.get()[index]
  override fun close() { statement.freemem() }
}

private class JsSqlPreparedStatement : SqlPreparedStatement {

  val parameters = mutableListOf<Any?>()

  override fun bindBytes(index: Int, bytes: ByteArray?) {
    parameters.add(bytes?.toTypedArray())
  }

  override fun bindLong(index: Int, long: Long?) {
    // We convert Long to Double because Kotlin's Double is mapped to JS number
    // whereas Kotlin's Long is implemented as a JS object
    parameters.add(long?.toDouble())
  }

  override fun bindDouble(index: Int, double: Double?) {
    parameters.add(double)
  }

  override fun bindString(index: Int, string: String?) {
    parameters.add(string)
  }
}

@file:JsModule("sql.js")
@file:JsNonModule
@file:Suppress("INTERFACE_WITH_SUPERCLASS", "OVERRIDING_FINAL_MEMBER", "RETURN_TYPE_MISMATCH_ON_OVERRIDE", "CONFLICTING_OVERLOADS", "EXTERNAL_DELEGATION")
package com.squareup.sqldelight.drivers.sqljs

import org.khronos.webgl.Uint8Array

external interface QueryResults {
  var columns: Array<String>
  var values: Array<Array<dynamic /* Number | String | Uint8Array | Nothing? */>>
}

external interface ParamsObject
external interface Buffer
external interface ParamsCallback
external interface Config

open external class Database() {
  constructor(data: Buffer?)
  constructor(data: Uint8Array?)
  constructor(data: Array<Number>?)
  open fun run(sql: String): Database
  open fun run(sql: String, params: ParamsObject): Database
  open fun run(sql: String, params: Array<dynamic /* Number | String | Uint8Array | Nothing? */>): Database
  open fun exec(sql: String): Array<QueryResults>
  open fun each(sql: String, callback: ParamsCallback, done: () -> Unit)
  open fun each(sql: String, params: ParamsObject, callback: ParamsCallback, done: () -> Unit)
  open fun each(sql: String, params: Array<dynamic /* Number | String | Uint8Array | Nothing? */>, callback: ParamsCallback, done: () -> Unit)
  open fun prepare(sql: String): Statement
  open fun prepare(sql: String, params: ParamsObject): Statement
  open fun prepare(sql: String, params: Array<dynamic /* Number | String | Uint8Array | Nothing? */>): Statement
  open fun export(): Uint8Array
  open fun close()
  open fun getRowsModified(): Number
  open fun create_function(name: String, func: Function<*>)
}

open external class Statement {
  open fun bind(): Boolean
  open fun bind(values: ParamsObject): Boolean
  open fun bind(values: Array<dynamic /* Number | String | Uint8Array | Nothing? */>): Boolean
  open fun step(): Boolean
  open fun get(): Array<dynamic /* Number | String | Uint8Array | Nothing? */>
  open fun get(params: ParamsObject): Array<dynamic /* Number | String | Uint8Array | Nothing? */>
  open fun get(params: Array<dynamic /* Number | String | Uint8Array | Nothing? */>): Array<dynamic /* Number | String | Uint8Array | Nothing? */>
  open fun getColumnNames(): Array<String>
  open fun getAsObject(): ParamsObject
  open fun getAsObject(params: ParamsObject): ParamsObject
  open fun getAsObject(params: Array<dynamic /* Number | String | Uint8Array | Nothing? */>): ParamsObject
  open fun run()
  open fun run(values: ParamsObject)
  open fun run(values: Array<dynamic /* Number | String | Uint8Array | Nothing? */>)
  open fun reset()
  open fun freemem()
  open fun free(): Boolean
}

external interface SqlJsStatic {
  var Database: InitDatabaseJsStatic
  var Statement: InitStatementJsStatic
}

external interface InitDatabaseJsStatic
external interface InitStatementJsStatic

external interface InitSqlJsStatic {
  var default: InitSqlJsStatic /* this */
}

package com.squareup.sqldelight.drivers.sqljs

import org.khronos.webgl.Uint8Array
import kotlin.js.Promise

operator fun InitStatementJsStatic.invoke(): Statement = createInstance(this)
operator fun InitDatabaseJsStatic.invoke(): Database = createInstance(this)
operator fun InitDatabaseJsStatic.invoke(data: Array<Number>): Database = createInstance(this, data)
operator fun InitDatabaseJsStatic.invoke(data: Uint8Array): Database = createInstance(this, data)
operator fun InitSqlJsStatic.invoke(): Promise<SqlJsStatic> = asDynamic()()
operator fun InitSqlJsStatic.invoke(config: Config?): Promise<SqlJsStatic> = asDynamic()(config)

@JsNonModule @JsModule("sql.js")
external val initSqlJs: InitSqlJsStatic

@Suppress("UNUSED_VARIABLE", "UNUSED_PARAMETER")
fun createInstance(type: dynamic, vararg args: dynamic): dynamic {
  val argsArray = (listOf(null) + args).toTypedArray()
  return js("new (Function.prototype.bind.apply(type, argsArray))")
}

fun initDb(config: Config? = js("{}")): Promise<Database> = initSqlJs(config).then { it.Database() }

package com.squareup.sqldelight.android.paging

import androidx.paging.DataSource
import androidx.paging.PositionalDataSource
import com.squareup.sqldelight.Query
import com.squareup.sqldelight.Transacter

class QueryDataSourceFactory<RowType : Any>(
  private val queryProvider: (limit: Long, offset: Long) -> Query<RowType>,
  private val countQuery: Query<Long>,
  private val transacter: Transacter
) : DataSource.Factory<Int, RowType>() {
  override fun create(): PositionalDataSource<RowType> =
    QueryDataSource(queryProvider, countQuery, transacter)
}

private class QueryDataSource<RowType : Any>(
  private val queryProvider: (limit: Long, offset: Long) -> Query<RowType>,
  private val countQuery: Query<Long>,
  private val transacter: Transacter
) : PositionalDataSource<RowType>(),
  Query.Listener {
  private var query: Query<RowType>? = null
  private val callbacks = linkedSetOf<InvalidatedCallback>()

  override fun queryResultsChanged() = invalidate()

  override fun invalidate() {
    query?.removeListener(this)
    query = null
    super.invalidate()
  }

  override fun addInvalidatedCallback(onInvalidatedCallback: InvalidatedCallback) {
    super.addInvalidatedCallback(onInvalidatedCallback)
    if (callbacks.isEmpty()) {
      query?.addListener(this)
    }
    callbacks.add(onInvalidatedCallback)
  }

  override fun removeInvalidatedCallback(onInvalidatedCallback: InvalidatedCallback) {
    super.removeInvalidatedCallback(onInvalidatedCallback)
    callbacks.remove(onInvalidatedCallback)
    if (callbacks.isEmpty()) {
      query?.removeListener(this)
    }
  }

  override fun loadRange(
    params: LoadRangeParams,
    callback: LoadRangeCallback<RowType>
  ) {
    query?.removeListener(this)
    queryProvider(params.loadSize.toLong(), params.startPosition.toLong()).let { query ->
      if (callbacks.isNotEmpty()) {
        query.addListener(this)
      }
      this.query = query
      if (!isInvalid) {
        callback.onResult(query.executeAsList())
      }
    }
  }

  override fun loadInitial(
    params: LoadInitialParams,
    callback: LoadInitialCallback<RowType>
  ) {
    query?.removeListener(this)
    queryProvider(params.requestedLoadSize.toLong(), params.requestedStartPosition.toLong()).let { query ->
      if (callbacks.isNotEmpty()) {
        query.addListener(this)
      }
      this.query = query
      if (!isInvalid) {
        transacter.transaction {
          callback.onResult(
            /* data = */ query.executeAsList(),
            /* position = */ params.requestedStartPosition,
            /* totalCount = */ countQuery.executeAsOne().toInt()
          )
        }
      }
    }
  }
}

package com.squareup.sqldelight.android.paging3

import androidx.paging.PagingState
import com.squareup.sqldelight.Query
import com.squareup.sqldelight.Transacter
import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.withContext

internal class KeyedQueryPagingSource<Key : Any, RowType : Any>(
  private val queryProvider: (beginInclusive: Key, endExclusive: Key?) -> Query<RowType>,
  private val pageBoundariesProvider: (anchor: Key?, limit: Long) -> Query<Key>,
  private val transacter: Transacter,
  private val dispatcher: CoroutineDispatcher,
) : QueryPagingSource<Key, RowType>() {

  private var pageBoundaries: List<Key>? = null
  override val jumpingSupported: Boolean get() = false

  override fun getRefreshKey(state: PagingState<Key, RowType>): Key? {
    val boundaries = pageBoundaries ?: return null
    val last = state.pages.lastOrNull() ?: return null
    val keyIndexFromNext = last.nextKey?.let { boundaries.indexOf(it) - 1 }
    val keyIndexFromPrev = last.prevKey?.let { boundaries.indexOf(it) + 1 }
    val keyIndex = keyIndexFromNext ?: keyIndexFromPrev ?: return null

    return boundaries.getOrNull(keyIndex)
  }

  override suspend fun load(params: LoadParams<Key>): LoadResult<Key, RowType> {
    return withContext(dispatcher) {
      try {
        transacter.transactionWithResult {
          val boundaries = pageBoundaries
            ?: pageBoundariesProvider(params.key, params.loadSize.toLong())
              .executeAsList()
              .also { pageBoundaries = it }

          val key = params.key ?: boundaries.first()

          require(key in boundaries)

          val keyIndex = boundaries.indexOf(key)
          val previousKey = boundaries.getOrNull(keyIndex - 1)
          val nextKey = boundaries.getOrNull(keyIndex + 1)
          val results = queryProvider(key, nextKey)
            .also { currentQuery = it }
            .executeAsList()

          LoadResult.Page(
            data = results,
            prevKey = previousKey,
            nextKey = nextKey,
          )
        }
      } catch (e: Exception) {
        if (e is IllegalArgumentException) throw e
        LoadResult.Error(e)
      }
    }
  }
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.android.paging3

import androidx.paging.PagingState
import com.squareup.sqldelight.Query
import com.squareup.sqldelight.Transacter
import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.withContext

internal class OffsetQueryPagingSource<RowType : Any>(
  private val queryProvider: (limit: Long, offset: Long) -> Query<RowType>,
  private val countQuery: Query<Long>,
  private val transacter: Transacter,
  private val dispatcher: CoroutineDispatcher,
) : QueryPagingSource<Long, RowType>() {

  override val jumpingSupported get() = true

  override suspend fun load(
    params: LoadParams<Long>
  ): LoadResult<Long, RowType> = withContext(dispatcher) {
    try {
      val key = params.key ?: 0L
      transacter.transactionWithResult {
        val count = countQuery.executeAsOne()
        if (count != 0L && key >= count) throw IndexOutOfBoundsException()

        val loadSize = if (key < 0) params.loadSize + key else params.loadSize

        val data = queryProvider(loadSize.toLong(), maxOf(0, key))
          .also { currentQuery = it }
          .executeAsList()

        LoadResult.Page(
          data = data,
          // allow one, and only one negative prevKey in a paging set. This is done for
          // misaligned prepend queries to avoid duplicates.
          prevKey = if (key <= 0L) null else key - params.loadSize,
          nextKey = if (key + params.loadSize >= count) null else key + params.loadSize,
          itemsBefore = maxOf(0L, key).toInt(),
          itemsAfter = maxOf(0, (count - (key + params.loadSize))).toInt()
        )
      }
    } catch (e: Exception) {
      if (e is IndexOutOfBoundsException) throw e
      LoadResult.Error(e)
    }
  }

  override fun getRefreshKey(state: PagingState<Long, RowType>) = state.anchorPosition?.toLong()
}

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * 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.squareup.sqldelight.android.paging3

import androidx.paging.PagingConfig
import androidx.paging.PagingSource
import com.squareup.sqldelight.Query
import com.squareup.sqldelight.Transacter
import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.Dispatchers
import kotlin.properties.Delegates

internal abstract class QueryPagingSource<Key : Any, RowType : Any> :
  PagingSource<Key, RowType>(),
  Query.Listener {

  protected var currentQuery: Query<RowType>? by Delegates.observable(null) { _, old, new ->
    old?.removeListener(this)
    new?.addListener(this)
  }

  init {
    registerInvalidatedCallback {
      currentQuery?.removeListener(this)
      currentQuery = null
    }
  }

  final override fun queryResultsChanged() = invalidate()
}

/**
 * Create a [PagingSource] that pages through results according to queries generated by
 * [queryProvider]. Queries returned by [queryProvider] should expect to do SQL offset/limit
 * based paging. For that reason, [countQuery] is required to calculate pages and page offsets.
 *
 * An example query returned by [queryProvider] could look like:
 *
 * ```sql
 * SELECT value FROM numbers
 * LIMIT 10
 * OFFSET 100;
 * ```
 *
 * Queries will be executed on [dispatcher].
 */
@Suppress("FunctionName")
fun <RowType : Any> QueryPagingSource(
  countQuery: Query<Long>,
  transacter: Transacter,
  dispatcher: CoroutineDispatcher = Dispatchers.IO,
  queryProvider: (limit: Long, offset: Long) -> Query<RowType>,
): PagingSource<Long, RowType> = OffsetQueryPagingSource(
  queryProvider,
  countQuery,
  transacter,
  dispatcher
)

/**
 * Create a [PagingSource] that pages through results according to queries generated by
 * [queryProvider]. Queries returned by [queryProvider] should expected to do keyset paging.
 * For that reason, queries should be arranged by an non-ambigious `ORDER BY` clause. [Key] must
 * be a unique clause that rows are ordered by. For performance reasons, an index should be present
 * on [Key].
 *
 * [pageBoundariesProvider] is a callback that produces a query containing [Key] items that specifies
 * where each page boundary exists within the full dataset. For example:
 *
 * The dataset `[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]` ordered ascending with a page size of 2 would produce
 * page boundaries `[0, 2, 4, 6, 8]`.
 *
 * The dataset `[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]` ordered ascending with a page size of 3 would produce
 * page boundaries `[0, 3, 6, 9]`.
 *
 * Callbacks made from [pageBoundariesProvider] may include an `anchor` key which, if present,
 * should appear in the resulting query.

 * Because page boundaries are computed ahead of time, [PagingConfig.initialLoadSize] should match
 * [PagingConfig.pageSize]. Failing to do so will result in unexpected page sizes, as
 * [pageBoundariesProvider] is called a single time during the first call to [PagingSource.load]
 * on this source.
 *
 * Generally, it's only feasible to produce page boundaries using SQLite window functions. An example
 * query to generate page boundaries like shown above would look like the following.
 *
 * ```sql
 * SELECT value
 * FROM (
 *   SELECT
 *     value,
 *     CASE
 *       WHEN ((row_number() OVER(ORDER BY value ASC) - 1) % :limit) = 0 THEN 1
 *       WHEN value = :anchor THEN 1
 *       ELSE 0
 *     END page_boundary
 *   FROM numbers
 *   ORDER BY value ASC
 * )
 * WHERE page_boundary = 1;
 * ```
 *
 * SQLite window queries became available as of version 3.25.0. For this reason, consuming
 * applications will likely need a minSdk of 30 set _or_ bundle a SQLite module separate from the OS
 * provided module.
 *
 * An example query returned by [queryProvider] could look like:
 *
 * ```sql
 * SELECT value FROM numbers
 * WHERE value >= :beginInclusive AND (value < :endExclusive OR :endExclusive IS NULL)
 * ORDER BY value ASC;
 * ```
 *
 * Queries will be executed on [dispatcher].
 *
 * This [PagingSource] _does not_ support jumping. If your use case requires jumping, use the
 * offset based [QueryPagingSource] function.
 */
@Suppress("FunctionName")
fun <Key : Any, RowType : Any> QueryPagingSource(
  transacter: Transacter,
  dispatcher: CoroutineDispatcher,
  pageBoundariesProvider: (anchor: Key?, limit: Long) -> Query<Key>,
  queryProvider: (beginInclusive: Key, endExclusive: Key?) -> Query<RowType>,
): PagingSource<Key, RowType> = KeyedQueryPagingSource(
  queryProvider,
  pageBoundariesProvider,
  transacter,
  dispatcher,
)

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.
 */

@file:JvmName("FlowQuery")

package com.squareup.sqldelight.runtime.coroutines

import com.squareup.sqldelight.Query
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.Channel.Factory.CONFLATED
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.map
import kotlinx.coroutines.flow.mapNotNull
import kotlinx.coroutines.withContext
import kotlin.coroutines.CoroutineContext
import kotlin.jvm.JvmName
import kotlin.jvm.JvmOverloads

/** Turns this [Query] into a [Flow] which emits whenever the underlying result set changes. */
@JvmName("toFlow")
fun <T : Any> Query<T>.asFlow(): Flow<Query<T>> = flow {
  val channel = Channel<Unit>(CONFLATED)
  channel.trySend(Unit)

  val listener = object : Query.Listener {
    override fun queryResultsChanged() {
      channel.trySend(Unit)
    }
  }

  addListener(listener)
  try {
    for (item in channel) {
      emit(this@asFlow)
    }
  } finally {
    removeListener(listener)
  }
}

@JvmOverloads
fun <T : Any> Flow<Query<T>>.mapToOne(
  context: CoroutineContext = Dispatchers.Default
): Flow<T> = map {
  withContext(context) {
    it.executeAsOne()
  }
}

@JvmOverloads
fun <T : Any> Flow<Query<T>>.mapToOneOrDefault(
  defaultValue: T,
  context: CoroutineContext = Dispatchers.Default
): Flow<T> = map {
  withContext(context) {
    it.executeAsOneOrNull() ?: defaultValue
  }
}

@JvmOverloads
fun <T : Any> Flow<Query<T>>.mapToOneOrNull(
  context: CoroutineContext = Dispatchers.Default
): Flow<T?> = map {
  withContext(context) {
    it.executeAsOneOrNull()
  }
}

@JvmOverloads
fun <T : Any> Flow<Query<T>>.mapToOneNotNull(
  context: CoroutineContext = Dispatchers.Default
): Flow<T> = mapNotNull {
  withContext(context) {
    it.executeAsOneOrNull()
  }
}

@JvmOverloads
fun <T : Any> Flow<Query<T>>.mapToList(
  context: CoroutineContext = Dispatchers.Default
): Flow<List<T>> = map {
  withContext(context) {
    it.executeAsList()
  }
}

@file:JvmName("RxQuery")

package com.squareup.sqldelight.runtime.rx

import com.squareup.sqldelight.Query
import io.reactivex.Observable
import io.reactivex.ObservableEmitter
import io.reactivex.ObservableOnSubscribe
import io.reactivex.Scheduler
import io.reactivex.annotations.CheckReturnValue
import io.reactivex.disposables.Disposable
import io.reactivex.schedulers.Schedulers
import java.util.Optional
import java.util.concurrent.atomic.AtomicBoolean

/**
 * Turns this [Query] into an [Observable] which emits whenever the underlying result set changes.
 *
 * @param scheduler By default, emissions occur on the [Schedulers.io] scheduler but can be
 * optionally overridden.
 */
@CheckReturnValue
@JvmOverloads
@JvmName("toObservable")
fun <T : Any> Query<T>.asObservable(scheduler: Scheduler = Schedulers.io()): Observable<Query<T>> {
  return Observable.create(QueryOnSubscribe(this)).observeOn(scheduler)
}

private class QueryOnSubscribe<T : Any>(
  private val query: Query<T>
) : ObservableOnSubscribe<Query<T>> {
  override fun subscribe(emitter: ObservableEmitter<Query<T>>) {
    val listenerAndDisposable = QueryListenerAndDisposable(emitter, query)
    query.addListener(listenerAndDisposable)
    emitter.setDisposable(listenerAndDisposable)
    emitter.onNext(query)
  }
}

private class QueryListenerAndDisposable<T : Any>(
  private val emitter: ObservableEmitter<Query<T>>,
  private val query: Query<T>
) : AtomicBoolean(), Query.Listener, Disposable {
  override fun queryResultsChanged() {
    emitter.onNext(query)
  }

  override fun isDisposed() = get()

  override fun dispose() {
    if (compareAndSet(false, true)) {
      query.removeListener(this)
    }
  }
}

@CheckReturnValue
fun <T : Any> Observable<Query<T>>.mapToOne(): Observable<T> {
  return map { it.executeAsOne() }
}

@CheckReturnValue
fun <T : Any> Observable<Query<T>>.mapToOneOrDefault(defaultValue: T): Observable<T> {
  return map { it.executeAsOneOrNull() ?: defaultValue }
}

@CheckReturnValue
fun <T : Any> Observable<Query<T>>.mapToOptional(): Observable<Optional<T>> {
  return map { Optional.ofNullable(it.executeAsOneOrNull()) }
}

@CheckReturnValue
fun <T : Any> Observable<Query<T>>.mapToList(): Observable<List<T>> {
  return map { it.executeAsList() }
}

@CheckReturnValue
fun <T : Any> Observable<Query<T>>.mapToOneNonNull(): Observable<T> {
  return flatMap {
    val result = it.executeAsOneOrNull()
    if (result == null) Observable.empty() else Observable.just(result)
  }
}

@file:JvmName("RxQuery")

package com.squareup.sqldelight.runtime.rx3

import com.squareup.sqldelight.Query
import io.reactivex.rxjava3.annotations.CheckReturnValue
import io.reactivex.rxjava3.core.Observable
import io.reactivex.rxjava3.core.ObservableEmitter
import io.reactivex.rxjava3.core.ObservableOnSubscribe
import io.reactivex.rxjava3.core.Scheduler
import io.reactivex.rxjava3.disposables.Disposable
import io.reactivex.rxjava3.schedulers.Schedulers
import java.util.Optional
import java.util.concurrent.atomic.AtomicBoolean

/**
 * Turns this [Query] into an [Observable] which emits whenever the underlying result set changes.
 *
 * @param scheduler By default, emissions occur on the [Schedulers.io] scheduler but can be
 * optionally overridden.
 */
@CheckReturnValue
@JvmOverloads
@JvmName("toObservable")
fun <T : Any> Query<T>.asObservable(scheduler: Scheduler = Schedulers.io()): Observable<Query<T>> {
  return Observable.create(QueryOnSubscribe(this)).observeOn(scheduler)
}

private class QueryOnSubscribe<T : Any>(
  private val query: Query<T>
) : ObservableOnSubscribe<Query<T>> {
  override fun subscribe(emitter: ObservableEmitter<Query<T>>) {
    val listenerAndDisposable = QueryListenerAndDisposable(emitter, query)
    query.addListener(listenerAndDisposable)
    emitter.setDisposable(listenerAndDisposable)
    emitter.onNext(query)
  }
}

private class QueryListenerAndDisposable<T : Any>(
  private val emitter: ObservableEmitter<Query<T>>,
  private val query: Query<T>
) : AtomicBoolean(), Query.Listener, Disposable {
  override fun queryResultsChanged() {
    emitter.onNext(query)
  }

  override fun isDisposed() = get()

  override fun dispose() {
    if (compareAndSet(false, true)) {
      query.removeListener(this)
    }
  }
}

@CheckReturnValue
fun <T : Any> Observable<Query<T>>.mapToOne(): Observable<T> {
  return map { it.executeAsOne() }
}

@CheckReturnValue
fun <T : Any> Observable<Query<T>>.mapToOneOrDefault(defaultValue: T): Observable<T> {
  return map { it.executeAsOneOrNull() ?: defaultValue }
}

@CheckReturnValue
fun <T : Any> Observable<Query<T>>.mapToOptional(): Observable<Optional<T>> {
  return map { Optional.ofNullable(it.executeAsOneOrNull()) }
}

@CheckReturnValue
fun <T : Any> Observable<Query<T>>.mapToList(): Observable<List<T>> {
  return map { it.executeAsList() }
}

@CheckReturnValue
fun <T : Any> Observable<Query<T>>.mapToOneNonNull(): Observable<T> {
  return flatMap {
    val result = it.executeAsOneOrNull()
    if (result == null) Observable.empty() else Observable.just(result)
  }
}

package com.example.sqldelight.hockey.data

import android.content.Context
import com.example.sqldelight.hockey.HockeyDb
import com.squareup.sqldelight.android.AndroidSqliteDriver

fun Db.getInstance(context: Context): HockeyDb {
  if (!Db.ready) {
    Db.dbSetup(AndroidSqliteDriver(Schema, context))
  }
  return Db.instance
}

package com.example.sqldelight.hockey.ui

import android.app.Activity
import android.content.Intent
import android.os.Bundle
import android.widget.Button
import com.example.sqldelight.hockey.R

class MainActivity : Activity() {
  override fun onCreate(savedInstanceState: Bundle?) {
    super.onCreate(savedInstanceState)
    setContentView(R.layout.main_activity)
    findViewById<Button>(R.id.list).setOnClickListener {
      startActivity(Intent(this, PlayersActivity::class.java))
    }
    findViewById<Button>(R.id.teams).setOnClickListener {
      startActivity(Intent(this, TeamsActivity::class.java))
    }
  }
}

package com.example.sqldelight.hockey.ui

import android.content.Context
import android.util.AttributeSet
import android.widget.LinearLayout
import android.widget.TextView
import com.example.sqldelight.hockey.R
import com.example.sqldelight.hockey.data.ForTeam

class PlayerRow(
  context: Context,
  attrs: AttributeSet
) : LinearLayout(context, attrs) {
  fun populate(row: ForTeam) {
    findViewById<TextView>(R.id.player_name).text = "${row.first_name} ${row.last_name}"
    findViewById<TextView>(R.id.player_number).text = row.number
    findViewById<TextView>(R.id.team_name).text = row.teamName
  }
}

package com.example.sqldelight.hockey.ui

import android.app.Activity
import android.os.Bundle
import android.view.ViewGroup
import androidx.recyclerview.widget.LinearLayoutManager
import androidx.recyclerview.widget.RecyclerView
import com.example.sqldelight.hockey.R
import com.example.sqldelight.hockey.data.Db
import com.example.sqldelight.hockey.data.ForTeam
import com.example.sqldelight.hockey.data.getInstance
import com.example.sqldelight.hockey.ui.PlayersActivity.PlayersAdapter.ViewHolder

class PlayersActivity : Activity() {
  override fun onCreate(savedInstanceState: Bundle?) {
    super.onCreate(savedInstanceState)
    setContentView(R.layout.list)

    val db = Db.getInstance(this).playerQueries
    val players = findViewById<RecyclerView>(R.id.list)
    players.layoutManager = LinearLayoutManager(this)
    players.adapter = PlayersAdapter(db.forTeam(intent.getLongExtra(TEAM_ID, -1)).executeAsList())
  }

  private inner class PlayersAdapter(
    private val data: List<ForTeam>
  ) : RecyclerView.Adapter<ViewHolder>() {
    override fun getItemCount() = data.size

    override fun onCreateViewHolder(parent: ViewGroup, viewType: Int): ViewHolder {
      return ViewHolder(layoutInflater.inflate(R.layout.player_row, parent, false) as PlayerRow)
    }

    override fun onBindViewHolder(holder: ViewHolder, position: Int) {
      holder.row.populate(data[position])
    }

    inner class ViewHolder(val row: PlayerRow) : RecyclerView.ViewHolder(row)
  }

  companion object {
    val TEAM_ID = "team_id"
  }
}

package com.example.sqldelight.hockey.ui

import android.content.Context
import android.util.AttributeSet
import android.widget.LinearLayout
import android.widget.TextView
import com.example.sqldelight.hockey.R
import com.example.sqldelight.hockey.data.Team
import com.example.sqldelight.hockey.platform.defaultFormatter

class TeamRow(
  context: Context,
  attrs: AttributeSet
) : LinearLayout(context, attrs) {
  fun populate(team: Team) {
    findViewById<TextView>(R.id.team_name).text = team.name
    findViewById<TextView>(R.id.coach_name).text = team.coach
    findViewById<TextView>(R.id.founded).text = context.getString(R.string.founded, defaultFormatter.format(team.founded))
  }
}

package com.example.sqldelight.hockey.ui

import android.app.Activity
import android.content.Intent
import android.os.Bundle
import android.view.ViewGroup
import androidx.recyclerview.widget.LinearLayoutManager
import androidx.recyclerview.widget.RecyclerView
import com.example.sqldelight.hockey.R
import com.example.sqldelight.hockey.data.Db
import com.example.sqldelight.hockey.data.Team
import com.example.sqldelight.hockey.data.getInstance
import com.example.sqldelight.hockey.ui.TeamsActivity.Adapter.ViewHolder

class TeamsActivity : Activity() {
  override fun onCreate(savedInstanceState: Bundle?) {
    super.onCreate(savedInstanceState)
    setContentView(R.layout.list)

    val db = Db.getInstance(this).teamQueries
    val adapter = Adapter(db.selectAll().executeAsList()) { teamClicked ->
      val intent = Intent(this@TeamsActivity, PlayersActivity::class.java)
      intent.putExtra(PlayersActivity.TEAM_ID, teamClicked.id)
      startActivity(intent)
    }
    val teams = findViewById<RecyclerView>(R.id.list)
    teams.layoutManager = LinearLayoutManager(this)
    teams.adapter = adapter
  }

  private inner class Adapter(
    private val data: List<Team>,
    private val clickListener: (Team) -> Unit
  ) : RecyclerView.Adapter<ViewHolder>() {
    override fun getItemCount() = data.size

    override fun onCreateViewHolder(parent: ViewGroup, viewType: Int): ViewHolder {
      return ViewHolder(layoutInflater.inflate(R.layout.team_row, parent, false) as TeamRow)
    }

    override fun onBindViewHolder(holder: ViewHolder, position: Int) {
      holder.setTeam(data[position])
    }

    inner class ViewHolder(val row: TeamRow) : RecyclerView.ViewHolder(row) {
      fun setTeam(team: Team) {
        row.populate(team)
        row.setOnClickListener { clickListener(team) }
      }
    }
  }
}

package com.example.sqldelight.hockey.data

import com.squareup.sqldelight.ColumnAdapter

expect class Date(year: Int, month: Int, day: Int)

expect class DateAdapter() : ColumnAdapter<Date, Long>

package com.example.sqldelight.hockey.data

interface PlayerVals {
  enum class Shoots {
    RIGHT,
    LEFT
  }

  enum class Position {
    LEFT_WING,
    RIGHT_WING,
    CENTER,
    DEFENSE,
    GOALIE
  }
}

package com.example.sqldelight.hockey.data

import com.example.sqldelight.hockey.HockeyDb
import com.example.sqldelight.hockey.data.PlayerVals.Position
import com.example.sqldelight.hockey.data.PlayerVals.Shoots
import com.squareup.sqldelight.EnumColumnAdapter
import com.squareup.sqldelight.db.SqlDriver

fun createQueryWrapper(driver: SqlDriver): HockeyDb {
  return HockeyDb(
    driver = driver,
    teamAdapter = Team.Adapter(
      foundedAdapter = DateAdapter()
    ),
    playerAdapter = Player.Adapter(
      shootsAdapter = EnumColumnAdapter(),
      positionAdapter = EnumColumnAdapter(),
      birth_dateAdapter = DateAdapter()
    )
  )
}

object Schema : SqlDriver.Schema by HockeyDb.Schema {
  override fun create(driver: SqlDriver) {
    HockeyDb.Schema.create(driver)

    // Seed data time!
    createQueryWrapper(driver).apply {

      val ducks = "Anaheim Ducks"
      val pens = "Pittsburgh Penguins"
      val sharks = "San Jose Sharks"

      // Populate teams.
      teamQueries.insertTeam(ducks, Date(1993, 3, 1), "Randy Carlyle", true)
      teamQueries.insertTeam(pens, Date(1966, 2, 8), "Mike Sullivan", true)
      teamQueries.insertTeam(sharks, Date(1990, 5, 5), "Peter DeBoer", false)

      playerQueries.insertPlayer(
        "Corey", "Perry", 10, ducks, 30, 210F, Date(1985, 5, 16),
        Shoots.RIGHT, Position.RIGHT_WING
      )
      playerQueries.insertPlayer(
        "Ryan", "Getzlaf", 15, ducks, 30, 221F, Date(1985, 5, 10),
        Shoots.RIGHT, Position.CENTER
      )
      teamQueries.setCaptain(15, ducks)

      playerQueries.insertPlayer(
        "Sidney", "Crosby", 87, pens, 28, 200F, Date(1987, 8, 7),
        Shoots.LEFT, Position.CENTER
      )
      teamQueries.setCaptain(87, pens)

      playerQueries.insertPlayer(
        "Erik", "Karlsson", 65, sharks, 28, 190F, Date(1990, 5, 31),
        Shoots.RIGHT, Position.DEFENSE
      )

      playerQueries.insertPlayer(
        "Joe", "Pavelski", 8, sharks, 31, 194F, Date(1984, 7, 18),
        Shoots.RIGHT, Position.CENTER
      )
      teamQueries.setCaptain(8, sharks)
    }
  }
}

package com.example.sqldelight.hockey.platform

import com.example.sqldelight.hockey.data.Date

expect class DateFormatHelper(format: String) {
  fun format(d: Date): String
}

val defaultFormatter = DateFormatHelper("dd/MM/yyyy")

package com.example.sqldelight.hockey.data

import com.squareup.sqldelight.ColumnAdapter
import platform.Foundation.NSCalendar
import platform.Foundation.NSDate
import platform.Foundation.NSDateComponents
import platform.Foundation.dateWithTimeIntervalSince1970
import platform.Foundation.timeIntervalSince1970
import kotlin.math.floor

actual class Date internal constructor(internal val nsDate: NSDate) {
  actual constructor(year: Int, month: Int, day: Int) : this(partsToDate(year, month + 1, day))
}

internal fun partsToDate(year: Int, month: Int, day: Int): NSDate {
  val cal = NSCalendar.currentCalendar
  val comps = NSDateComponents()
  comps.setDay(day.toLong())
  comps.setMonth(month.toLong())
  comps.setYear(year.toLong())
  return cal.dateFromComponents(comps)!!
}

actual class DateAdapter actual constructor() : ColumnAdapter<Date, Long> {
  override fun decode(databaseValue: Long): Date =
    Date(NSDate.dateWithTimeIntervalSince1970(databaseValue.toDouble() / 1000))

  override fun encode(value: Date): Long =
    floor(value.nsDate.timeIntervalSince1970 * 1000L).toLong()
}

package com.example.sqldelight.hockey.data

import com.example.sqldelight.hockey.HockeyDb
import com.squareup.sqldelight.db.SqlDriver
import com.squareup.sqldelight.drivers.native.NativeSqliteDriver
import kotlin.native.concurrent.AtomicReference
import kotlin.native.concurrent.freeze

object Db {
  private val driverRef = AtomicReference<SqlDriver?>(null)
  private val dbRef = AtomicReference<HockeyDb?>(null)

  internal fun dbSetup(driver: SqlDriver) {
    val db = createQueryWrapper(driver)
    driverRef.value = driver.freeze()
    dbRef.value = db.freeze()
  }

  internal fun dbClear() {
    driverRef.value!!.close()
    dbRef.value = null
    driverRef.value = null
  }

  // Called from Swift
  @Suppress("unused")
  fun defaultDriver() {
    Db.dbSetup(NativeSqliteDriver(Schema, "sampledb"))
  }

  val instance: HockeyDb
    get() = dbRef.value!!
}

package com.example.sqldelight.hockey.platform

import com.example.sqldelight.hockey.data.Date
import platform.Foundation.NSDateFormatter

actual class DateFormatHelper actual constructor(format: String) {
  private val formatter: NSDateFormatter

  init {
    formatter = NSDateFormatter()
    formatter.dateFormat = format
  }

  actual fun format(d: Date): String = formatter.stringFromDate(d.nsDate)
}

package com.example.sqldelight.hockey.ui

import com.example.sqldelight.hockey.data.Db

// For Swift
@Suppress("unused")
object PlayerData {
  fun players(teamId: Long) = Db.instance.playerQueries.forTeam(teamId).executeAsList()
}

package com.example.sqldelight.hockey.ui

import com.example.sqldelight.hockey.data.Db

// For Swift
object TeamData {
  fun teams() = Db.instance.teamQueries.selectAll().executeAsList()
}

package com.example.sqldelight.hockey.data

import com.squareup.sqldelight.ColumnAdapter

actual typealias Date = kotlin.js.Date

actual class DateAdapter actual constructor() : ColumnAdapter<Date, Long> {
  override fun encode(value: Date) = value.getTime().toLong()
  override fun decode(databaseValue: Long) = Date(databaseValue)
}

package com.example.sqldelight.hockey.data

import com.example.sqldelight.hockey.HockeyDb
import com.squareup.sqldelight.db.SqlDriver

object Db {

  private lateinit var driver: SqlDriver
  private lateinit var db: HockeyDb

  fun dbSetup(driver: SqlDriver) {
    this.driver = driver
    this.db = createQueryWrapper(driver)
  }

  internal fun dbClear() {
    driver.close()
  }

  val instance: HockeyDb
    get() = db
}

package com.example.sqldelight.hockey.platform

import com.example.sqldelight.hockey.data.Date

@JsModule("dateformat")
external fun dateFormat(date: Date, format: String): String

actual class DateFormatHelper actual constructor(private val format: String) {
  actual fun format(d: Date): String = dateFormat(d, format.replace('M', 'm'))
}

package com.example.sqldelight.hockey.data

import com.squareup.sqldelight.ColumnAdapter
import java.util.GregorianCalendar

actual typealias Date = GregorianCalendar

actual class DateAdapter actual constructor() : ColumnAdapter<Date, Long> {
  override fun encode(value: Date) = value.timeInMillis
  override fun decode(databaseValue: Long) = Date.getInstance().apply {
    timeInMillis = databaseValue
  } as Date
}

package com.example.sqldelight.hockey.data

import com.example.sqldelight.hockey.HockeyDb
import com.squareup.sqldelight.db.SqlDriver

object Db {
  private var driverRef: SqlDriver? = null
  private var dbRef: HockeyDb? = null

  val ready: Boolean
    get() = driverRef != null

  fun dbSetup(driver: SqlDriver) {
    val db = createQueryWrapper(driver)
    driverRef = driver
    dbRef = db
  }

  internal fun dbClear() {
    driverRef!!.close()
    dbRef = null
    driverRef = null
  }

  val instance: HockeyDb
    get() = dbRef!!
}

package com.example.sqldelight.hockey.platform

import com.example.sqldelight.hockey.data.Date
import java.text.DateFormat
import java.text.SimpleDateFormat

actual class DateFormatHelper actual constructor(format: String) {
  val dateFormatter = object : ThreadLocal<DateFormat>() {
    override fun initialValue(): DateFormat = SimpleDateFormat(format)
  }

  actual fun format(d: Date): String = dateFormatter.get()!!.format(d.time)
}

package com.example.sqldelight.hockey.data

import com.example.sqldelight.hockey.HockeyDb
import com.squareup.sqldelight.drivers.sqljs.initSqlDriver
import kotlin.js.Promise

fun Db.getInstance(): Promise<HockeyDb> = initSqlDriver(Schema).then {
  Db.dbSetup(it)
  Db.instance
}

package com.squareup.sqldelight.core.compiler.integration

import com.intellij.psi.PsiElement

/**
 * This pattern consists of 3 parts:
 *
 * - `/\\*\\*` - matches the first line of the Javadoc:
 *
 * ```
 * </**>
 *  * Javadoc
 *  */
 * ```
 *
 * - `\n \\*[ /]?` - matches every other line of Javadoc:
 *
 * ```
 * /**<
 *  * >Javadoc<
 *  */>
 * ```
 *
 * - ` \\**slash` - specifically matches the tail part of a single-line Javadoc:
 *
 * ```
 * /* Javadoc< */>
 * ```
 */
private val JAVADOC_TEXT_REGEX = Regex("/\\*\\*|\n \\*[ /]?| \\*/")

internal fun javadocText(javadoc: PsiElement): String? {
  return javadoc.text
    .split(JAVADOC_TEXT_REGEX)
    .dropWhile { it.isEmpty() }
    .joinToString(separator = "\n") { it.trim() }
}

package com.squareup.sqldelight.core.compiler.integration

import com.alecstrong.sql.psi.core.psi.LazyQuery
import com.alecstrong.sql.psi.core.psi.SqlColumnDef
import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.alecstrong.sql.psi.core.psi.SqlCreateViewStmt
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.KModifier
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.sqldelight.core.compiler.SqlDelightCompiler
import com.squareup.sqldelight.core.lang.ADAPTER_NAME
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.psi.ColumnTypeMixin

internal fun LazyQuery.needsAdapters() = when (tableName.parent) {
  is SqlCreateViewStmt -> false
  else -> columns().any { (it.columnType as ColumnTypeMixin).adapter() != null }
}

internal fun LazyQuery.adapterProperty(): PropertySpec {
  val adapterType = ClassName(
    (tableName.containingFile as SqlDelightFile).packageName!!,
    SqlDelightCompiler.allocateName(tableName).capitalize(),
    ADAPTER_NAME
  )
  return PropertySpec.builder(adapterName, adapterType, KModifier.INTERNAL)
    .initializer(adapterName)
    .build()
}

private fun LazyQuery.columns() = when (val parentRule = tableName.parent) {
  is SqlCreateTableStmt -> parentRule.columnDefList
  else -> query.columns.map { it.element.parent as SqlColumnDef }
}

internal val LazyQuery.adapterName
  get() = "${SqlDelightCompiler.allocateName(tableName)}$ADAPTER_NAME"

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.core.compiler.model

import com.alecstrong.sql.psi.core.psi.SqlBindExpr
import com.alecstrong.sql.psi.core.psi.SqlBindParameter
import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.alecstrong.sql.psi.core.psi.SqlIdentifier
import com.alecstrong.sql.psi.core.psi.SqlInsertStmt
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.intellij.psi.PsiElement
import com.squareup.sqldelight.core.compiler.SqlDelightCompiler.allocateName
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.ARGUMENT
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.NULL
import com.squareup.sqldelight.core.lang.IntermediateType
import com.squareup.sqldelight.core.lang.acceptsTableInterface
import com.squareup.sqldelight.core.lang.psi.ColumnTypeMixin
import com.squareup.sqldelight.core.lang.util.argumentType
import com.squareup.sqldelight.core.lang.util.childOfType
import com.squareup.sqldelight.core.lang.util.columns
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.lang.util.interfaceType
import com.squareup.sqldelight.core.lang.util.table
import java.util.concurrent.ConcurrentHashMap

abstract class BindableQuery(
  internal val identifier: PsiElement?,
  internal val statement: PsiElement
) {
  abstract val id: Int

  internal val javadoc: PsiElement? = identifier?.childOfType(SqlTypes.JAVADOC)

  /**
   * The collection of parameters exposed in the generated api for this query.
   */
  internal val parameters: List<IntermediateType> by lazy {
    if (statement is SqlInsertStmt && statement.acceptsTableInterface()) {
      val table = statement.table.tableName.parent as SqlCreateTableStmt
      return@lazy listOf(
        IntermediateType(
          ARGUMENT,
          table.interfaceType,
          name = allocateName(table.tableName)
        )
      )
    }
    return@lazy arguments.sortedBy { it.index }.map { it.type }
  }

  /**
   * The collection of all bind expressions in this query.
   */
  internal val arguments: List<Argument> by lazy {
    if (statement is SqlInsertStmt && statement.acceptsTableInterface()) {
      return@lazy statement.columns.mapIndexed { index, column ->
        Argument(
          index + 1,
          (column.columnType as ColumnTypeMixin).type().let {
            it.copy(
              name = "${allocateName(statement.tableName)}.${it.name}",
              extracted = true
            )
          }
        )
      }
    }

    val result = mutableListOf<Argument>()
    val indexesSeen = mutableSetOf<Int>()
    val manuallyNamedIndexes = mutableSetOf<Int>()
    val namesSeen = mutableSetOf<String>()
    var maxIndexSeen = 0
    statement.findChildrenOfType<SqlBindExpr>().forEach { bindArg ->
      bindArg.bindParameter.node.findChildByType(SqlTypes.DIGIT)?.text?.toInt()?.let { index ->
        if (!indexesSeen.add(index)) {
          result.findAndReplace(bindArg, index) { it.index == index }
          return@forEach
        }
        maxIndexSeen = maxOf(maxIndexSeen, index)
        result.add(Argument(index, bindArg.argumentType(), mutableListOf(bindArg)))
        return@forEach
      }
      bindArg.bindParameter.identifier?.let {
        if (!namesSeen.add(it.text)) {
          result.findAndReplace(bindArg) { (_, type, _) -> type.name == it.text }
          return@forEach
        }
        val index = ++maxIndexSeen
        indexesSeen.add(index)
        manuallyNamedIndexes.add(index)
        result.add(Argument(index, bindArg.argumentType().copy(name = it.text), mutableListOf(bindArg)))
        return@forEach
      }
      val index = ++maxIndexSeen
      indexesSeen.add(index)
      result.add(Argument(index, bindArg.argumentType(), mutableListOf(bindArg)))
    }

    // If there are still naming conflicts (edge case where the name we generate is the same as
    // the name a user specified for a different parameter), resolve those.
    result.replaceAll {
      var name = it.type.name
      while (it.index !in manuallyNamedIndexes && !namesSeen.add(name)) {
        name += "_"
      }
      it.copy(type = it.type.copy(name = name))
    }

    if (statement is SqlInsertStmt) {
      return@lazy result.map {
        val isPrimaryKey = it.type.column?.columnConstraintList
          ?.any { it.node?.findChildByType(SqlTypes.PRIMARY) != null } == true
        if (isPrimaryKey && it.type.column?.columnType?.typeName?.text == "INTEGER") {
          // INTEGER Primary keys can be inserted as null to be auto-assigned a primary key.
          return@map it.copy(type = it.type.asNullable())
        }
        return@map it
      }
    }

    return@lazy result
  }

  private fun MutableList<Argument>.findAndReplace(
    bindArg: SqlBindExpr,
    index: Int? = null,
    condition: (Argument) -> Boolean
  ) {
    val current = first(condition)
    current.bindArgs.add(bindArg)

    val newArgumentType = when {
      // If we currently have a NULL type for this argument but encounter a different type later,
      // then the new type must be nullable.
      // i.e. WHERE (:foo IS NULL OR data = :foo)
      current.type.dialectType == NULL -> bindArg.argumentType()
      // If we'd previously assigned a type to this argument other than NULL, and later encounter NULL,
      // we should update the existing type to be nullable.
      // i.e. WHERE (data = :foo OR :foo IS NULL)
      bindArg.argumentType().dialectType == NULL && current.type.dialectType != NULL -> current.type
      // Nothing to update
      else -> null
    }

    if (newArgumentType != null) {
      remove(current)
      add(
        current.copy(
          index = index ?: current.index,
          type = newArgumentType.run {
            copy(
              javaType = javaType.copy(nullable = true),
              name = bindArg.bindParameter.identifier?.text ?: name
            )
          }
        )
      )
    }
  }

  private val SqlBindParameter.identifier: SqlIdentifier?
    get() = childOfType(SqlTypes.IDENTIFIER) as? SqlIdentifier

  internal data class Argument(
    val index: Int,
    val type: IntermediateType,
    val bindArgs: MutableList<SqlBindExpr> = mutableListOf()
  )

  companion object {
    /**
     * The query id map use to avoid string hashcode collision. Ideally this map should be per module.
     */
    val queryIdMap = ConcurrentHashMap<String, Int>()

    /**
     * Use the hashcode of qualifiedQueryName to generate the unique identifier id for queries. Detect the
     * hashcode collision by caching the generated identifiers. Runtime exception will be thrown when collision happens.
     * Client would need to give a different query name to avoid the collision.
     */
    fun getUniqueQueryIdentifier(qualifiedQueryName: String): Int {
      return when (queryIdMap.containsKey(qualifiedQueryName)) {
        true -> queryIdMap[qualifiedQueryName]!!
        else -> {
          val queryId = qualifiedQueryName.hashCode()
          if (queryIdMap.values.contains(queryId)) {
            // throw an exception here to ask the client to give a different query name which will not cause hashcode collision.
            // this should not happen often, when it happens, should be an easy fix for the client
            // to give a different query than adding logic to generate deterministic identifier
            throw RuntimeException(
              "HashCode collision happened when generating unique identifier for $qualifiedQueryName." +
                "Please give a different name"
            )
          }
          queryIdMap[qualifiedQueryName] = queryId
          queryId
        }
      }
    }
  }
}

package com.squareup.sqldelight.core.compiler.model

import com.intellij.psi.PsiElement
import com.squareup.sqldelight.core.lang.psi.StmtIdentifierMixin
import com.squareup.sqldelight.core.lang.util.sqFile

open class NamedExecute(
  identifier: StmtIdentifierMixin,
  statement: PsiElement
) : BindableQuery(identifier, statement) {
  val name = identifier.name!!

  override val id: Int
    // the sqlFile package name -> com.example.
    // sqlFile.name -> test.sq
    // name -> query name
    get() = idForIndex(null)

  internal fun idForIndex(index: Int?): Int {
    val postFix = if (index == null) "" else "_$index"
    return getUniqueQueryIdentifier(
      statement.sqFile().let {
        "${it.packageName}:${it.name}:$name$postFix"
      }
    )
  }
}

/*
 * Copyright (C) 2018 Square, Inc.
 *
 * 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.squareup.sqldelight.core.compiler.model

import com.alecstrong.sql.psi.core.psi.SqlDeleteStmtLimited
import com.alecstrong.sql.psi.core.psi.SqlInsertStmt
import com.alecstrong.sql.psi.core.psi.SqlTableName
import com.alecstrong.sql.psi.core.psi.SqlUpdateStmtLimited
import com.alecstrong.sql.psi.core.sqlite_3_24.psi.SqliteUpsertClause
import com.intellij.psi.PsiElement
import com.squareup.sqldelight.core.lang.SqlDelightQueriesFile
import com.squareup.sqldelight.core.lang.psi.StmtIdentifierMixin
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.lang.util.referencedTables
import com.squareup.sqldelight.core.lang.util.sqFile

sealed class NamedMutator(
  statement: PsiElement,
  identifier: StmtIdentifierMixin,
  tableName: SqlTableName
) : NamedExecute(identifier, statement) {
  val containingFile = statement.sqFile() as SqlDelightQueriesFile

  internal val tableEffected: SqlTableName by lazy {
    tableName.referencedTables().single()
  }

  class Insert(
    insert: SqlInsertStmt,
    identifier: StmtIdentifierMixin
  ) : NamedMutator(insert, identifier, insert.tableName) {
    val hasUpsertClause get() = statement.findChildrenOfType<SqliteUpsertClause>().isNotEmpty()
  }

  class Delete(
    delete: SqlDeleteStmtLimited,
    identifier: StmtIdentifierMixin
  ) : NamedMutator(delete, identifier, delete.qualifiedTableName!!.tableName)

  class Update(
    internal val update: SqlUpdateStmtLimited,
    identifier: StmtIdentifierMixin
  ) : NamedMutator(update, identifier, update.qualifiedTableName.tableName)
}

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.core.compiler.model

import com.alecstrong.sql.psi.core.psi.LazyQuery
import com.alecstrong.sql.psi.core.psi.NamedElement
import com.alecstrong.sql.psi.core.psi.QueryElement
import com.alecstrong.sql.psi.core.psi.SqlCompoundSelectStmt
import com.alecstrong.sql.psi.core.psi.SqlExpr
import com.alecstrong.sql.psi.core.psi.SqlSelectStmt
import com.alecstrong.sql.psi.core.psi.SqlTableName
import com.alecstrong.sql.psi.core.psi.SqlValuesExpression
import com.intellij.psi.PsiElement
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.sqldelight.core.compiler.SqlDelightCompiler.allocateName
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.ARGUMENT
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.BLOB
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.INTEGER
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.NULL
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.REAL
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.TEXT
import com.squareup.sqldelight.core.lang.CUSTOM_DATABASE_NAME
import com.squareup.sqldelight.core.lang.IntermediateType
import com.squareup.sqldelight.core.lang.SqlDelightQueriesFile
import com.squareup.sqldelight.core.lang.queriesName
import com.squareup.sqldelight.core.lang.util.name
import com.squareup.sqldelight.core.lang.util.sqFile
import com.squareup.sqldelight.core.lang.util.tablesObserved
import com.squareup.sqldelight.core.lang.util.type
import java.util.Locale

data class NamedQuery(
  val name: String,
  val select: SqlCompoundSelectStmt,
  private val statementIdentifier: PsiElement? = null
) : BindableQuery(statementIdentifier, select) {
  /**
   * Explodes the sqlite query into an ordered list (same order as the query) of types to be exposed
   * by the generated api.
   */
  internal val resultColumns: List<IntermediateType> by lazy {
    val namesUsed = LinkedHashSet<String>()

    select.selectStmtList.fold(emptyList()) { results, select ->
      val compoundSelect = if (select.valuesExpressionList.isNotEmpty()) {
        resultColumns(select.valuesExpressionList)
      } else {
        select.typesExposed(namesUsed)
      }
      if (results.isEmpty()) {
        return@fold compoundSelect
      }
      return@fold results.zip(compoundSelect, this::superType)
    }
  }

  /**
   * If this query is a pure select from a table (virtual or otherwise), this returns the LazyQuery
   * which points to that table (Pure meaning it has exactly the same columns in the same order).
   */
  private val pureTable: LazyQuery? by lazy {
    fun List<QueryElement.QueryColumn>.flattenCompounded(): List<QueryElement.QueryColumn> {
      return map { column ->
        if (column.compounded.none { it.element != column.element || it.nullable != column.nullable }) {
          column.copy(compounded = emptyList())
        } else {
          column
        }
      }
    }

    val pureColumns = select.queryExposed().singleOrNull()?.columns?.flattenCompounded()
    return@lazy select.tablesAvailable(select).firstOrNull {
      it.query.columns.flattenCompounded() == pureColumns
    }
  }

  /**
   * The name of the generated interface that this query references. The linked interface will have
   * a default implementation subclass.
   */
  internal val interfaceType: ClassName by lazy {
    pureTable?.let {
      return@lazy ClassName(it.tableName.sqFile().packageName!!, allocateName(it.tableName).capitalize(Locale.ROOT))
    }
    var packageName = select.sqFile().packageName!!
    if (select.sqFile().parent?.files
      ?.filterIsInstance<SqlDelightQueriesFile>()?.flatMap { it.namedQueries }
      ?.filter { it.needsInterface() && it != this }
      ?.any { it.name == name } == true
    ) {
      packageName = "$packageName.${select.sqFile().virtualFile!!.nameWithoutExtension.decapitalize(
        Locale.ROOT
      )}"
    }
    return@lazy ClassName(packageName, name.capitalize())
  }

  /**
   * @return true if this query needs its own interface generated.
   */
  internal fun needsInterface() = needsWrapper() && pureTable == null

  internal fun needsWrapper() = (resultColumns.size > 1 || resultColumns[0].javaType.isNullable)

  internal val tablesObserved: List<SqlTableName> by lazy { select.tablesObserved() }

  internal val queryProperty =
    CodeBlock.of("$CUSTOM_DATABASE_NAME.${select.sqFile().queriesName}.$name")

  internal val customQuerySubtype = "${name.capitalize()}Query"

  private fun resultColumns(valuesList: List<SqlValuesExpression>): List<IntermediateType> {
    return valuesList.fold(
      emptyList(),
      { results, values ->
        val exposedTypes = values.exprList.map { it.type() }
        if (results.isEmpty()) return@fold exposedTypes
        return@fold results.zip(exposedTypes, this::superType)
      }
    )
  }

  private fun superType(typeOne: IntermediateType, typeTwo: IntermediateType): IntermediateType {
    // Arguments types always take the other type.
    if (typeOne.dialectType == ARGUMENT) {
      return typeTwo.copy(name = typeOne.name)
    } else if (typeTwo.dialectType == ARGUMENT) {
      return typeOne
    }

    // Nullable types take nullable version of the other type.
    if (typeOne.dialectType == NULL) {
      return typeTwo.asNullable().copy(name = typeOne.name)
    } else if (typeTwo.dialectType == NULL) {
      return typeOne.asNullable()
    }

    val nullable = typeOne.javaType.isNullable || typeTwo.javaType.isNullable

    if (typeOne.dialectType != typeTwo.dialectType) {
      // Incompatible sqlite types. Prefer the type which can contain the other.
      // NULL < INTEGER < REAL < TEXT < BLOB
      val type = listOf(NULL, INTEGER, REAL, TEXT, BLOB)
        .last { it == typeOne.dialectType || it == typeTwo.dialectType }
      return IntermediateType(dialectType = type, name = typeOne.name).nullableIf(nullable)
    }

    if (typeOne.column !== typeTwo.column &&
      typeOne.asNonNullable().cursorGetter(0) != typeTwo.asNonNullable().cursorGetter(0) &&
      typeOne.column != null && typeTwo.column != null
    ) {
      // Incompatible adapters. Revert to unadapted java type.
      return if (typeOne.javaType.copy(nullable = false) == typeTwo.javaType.copy(nullable = false)) {
        typeOne.copy(assumedCompatibleTypes = typeOne.assumedCompatibleTypes + typeTwo).nullableIf(nullable)
      } else {
        IntermediateType(dialectType = typeOne.dialectType, name = typeOne.name).nullableIf(nullable)
      }
    }

    return typeOne.nullableIf(nullable)
  }

  private fun PsiElement.functionName() = when (this) {
    is NamedElement -> allocateName(this)
    is SqlExpr -> name
    else -> throw IllegalStateException("Cannot get name for type ${this.javaClass}")
  }

  private fun SqlSelectStmt.typesExposed(
    namesUsed: LinkedHashSet<String>
  ): List<IntermediateType> {
    return queryExposed().flatMap {
      val table = it.table?.name
      return@flatMap it.columns.map { queryColumn ->
        var name = queryColumn.element.functionName()
        if (!namesUsed.add(name)) {
          if (table != null) name = "${table}_$name"
          while (!namesUsed.add(name)) name += "_"
        }

        return@map queryColumn.type().copy(name = name)
      }
    }
  }

  private fun QueryElement.QueryColumn.type(): IntermediateType {
    var rootType = element.type()
    nullable?.let { rootType = rootType.nullableIf(it) }
    return compounded.fold(rootType) { type, column -> superType(type, column.type()) }
  }

  override val id: Int
    // the sqlFile package name -> com.example.
    // sqlFile.name -> test.sq
    // name -> query name
    get() = getUniqueQueryIdentifier(statement.sqFile().let { "${it.packageName}:${it.name}:$name" })
}

package com.squareup.sqldelight.core.dialect.api

import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.TypeName

internal interface DialectType {

  val javaType: TypeName

  fun decode(value: CodeBlock): CodeBlock = value

  fun encode(value: CodeBlock): CodeBlock = value

  fun prepareStatementBinder(columnIndex: String, value: CodeBlock): CodeBlock

  fun cursorGetter(columnIndex: Int): CodeBlock
}

package com.squareup.sqldelight.core.dialect.hsql

import com.squareup.kotlinpoet.BOOLEAN
import com.squareup.kotlinpoet.BYTE
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.INT
import com.squareup.kotlinpoet.LONG
import com.squareup.kotlinpoet.SHORT
import com.squareup.kotlinpoet.TypeName
import com.squareup.sqldelight.core.dialect.api.DialectType
import com.squareup.sqldelight.core.lang.CURSOR_NAME

internal enum class HsqlType(override val javaType: TypeName) : DialectType {
  TINY_INT(BYTE) {
    override fun decode(value: CodeBlock) = CodeBlock.of("%L.toByte()", value)

    override fun encode(value: CodeBlock) = CodeBlock.of("%L.toLong()", value)
  },
  SMALL_INT(SHORT) {
    override fun decode(value: CodeBlock) = CodeBlock.of("%L.toShort()", value)

    override fun encode(value: CodeBlock) = CodeBlock.of("%L.toLong()", value)
  },
  INTEGER(INT) {
    override fun decode(value: CodeBlock) = CodeBlock.of("%L.toInt()", value)

    override fun encode(value: CodeBlock) = CodeBlock.of("%L.toLong()", value)
  },
  BIG_INT(LONG),
  BOOL(BOOLEAN) {
    override fun decode(value: CodeBlock) = CodeBlock.of("%L == 1L", value)

    override fun encode(value: CodeBlock) = CodeBlock.of("if (%L) 1L else 0L", value)
  };

  override fun prepareStatementBinder(columnIndex: String, value: CodeBlock): CodeBlock {
    return CodeBlock.builder()
      .add(
        when (this) {
          TINY_INT, SMALL_INT, INTEGER, BIG_INT, BOOL -> "bindLong"
        }
      )
      .add("($columnIndex, %L)\n", value)
      .build()
  }

  override fun cursorGetter(columnIndex: Int): CodeBlock {
    return CodeBlock.of(
      when (this) {
        TINY_INT, SMALL_INT, INTEGER, BIG_INT, BOOL -> "$CURSOR_NAME.getLong($columnIndex)"
      }
    )
  }
}

package com.squareup.sqldelight.core.dialect.mysql

import com.alecstrong.sql.psi.core.mysql.psi.MySqlExtensionExpr
import com.alecstrong.sql.psi.core.psi.SqlExpr
import com.squareup.kotlinpoet.BOOLEAN
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType
import com.squareup.sqldelight.core.lang.IntermediateType
import com.squareup.sqldelight.core.lang.util.encapsulatingType
import com.squareup.sqldelight.core.lang.util.findChildrenOfType

internal fun MySqlExtensionExpr.type(): IntermediateType {
  // if (expr, expr, expr)
  return encapsulatingType(
    ifExpr.findChildrenOfType<SqlExpr>().drop(1),
    SqliteType.INTEGER, SqliteType.REAL, SqliteType.TEXT, SqliteType.BLOB
  )
}

internal fun MySqlExtensionExpr.argumentType(expr: SqlExpr): IntermediateType {
  return if (expr == ifExpr.children[0]) IntermediateType(SqliteType.INTEGER, BOOLEAN)
  else IntermediateType(SqliteType.ARGUMENT)
}

package com.squareup.sqldelight.core.dialect.mysql

import com.squareup.kotlinpoet.BOOLEAN
import com.squareup.kotlinpoet.BYTE
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.INT
import com.squareup.kotlinpoet.LONG
import com.squareup.kotlinpoet.SHORT
import com.squareup.kotlinpoet.TypeName
import com.squareup.sqldelight.core.dialect.api.DialectType
import com.squareup.sqldelight.core.lang.CURSOR_NAME

internal enum class MySqlType(override val javaType: TypeName) : DialectType {
  TINY_INT(BYTE) {
    override fun decode(value: CodeBlock) = CodeBlock.of("%L.toByte()", value)

    override fun encode(value: CodeBlock) = CodeBlock.of("%L.toLong()", value)
  },
  TINY_INT_BOOL(BOOLEAN) {
    override fun decode(value: CodeBlock) = CodeBlock.of("%L == 1L", value)

    override fun encode(value: CodeBlock) = CodeBlock.of("if (%L) 1L else 0L", value)
  },
  SMALL_INT(SHORT) {
    override fun decode(value: CodeBlock) = CodeBlock.of("%L.toShort()", value)

    override fun encode(value: CodeBlock) = CodeBlock.of("%L.toLong()", value)
  },
  INTEGER(INT) {
    override fun decode(value: CodeBlock) = CodeBlock.of("%L.toInt()", value)

    override fun encode(value: CodeBlock) = CodeBlock.of("%L.toLong()", value)
  },
  BIG_INT(LONG),
  BIT(BOOLEAN) {
    override fun decode(value: CodeBlock) = CodeBlock.of("%L == 1L", value)

    override fun encode(value: CodeBlock) = CodeBlock.of("if (%L) 1L else 0L", value)
  };

  override fun prepareStatementBinder(columnIndex: String, value: CodeBlock): CodeBlock {
    return CodeBlock.builder()
      .add(
        when (this) {
          TINY_INT, TINY_INT_BOOL, SMALL_INT, INTEGER, BIG_INT, BIT -> "bindLong"
        }
      )
      .add("($columnIndex, %L)\n", value)
      .build()
  }

  override fun cursorGetter(columnIndex: Int): CodeBlock {
    return CodeBlock.of(
      when (this) {
        TINY_INT, TINY_INT_BOOL, SMALL_INT, INTEGER, BIG_INT, BIT -> "$CURSOR_NAME.getLong($columnIndex)"
      }
    )
  }
}

package com.squareup.sqldelight.core.dialect.postgresql

import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.INT
import com.squareup.kotlinpoet.LONG
import com.squareup.kotlinpoet.SHORT
import com.squareup.kotlinpoet.TypeName
import com.squareup.sqldelight.core.dialect.api.DialectType
import com.squareup.sqldelight.core.lang.CURSOR_NAME

internal enum class PostgreSqlType(override val javaType: TypeName) : DialectType {
  SMALL_INT(SHORT) {
    override fun decode(value: CodeBlock) = CodeBlock.of("%L.toShort()", value)

    override fun encode(value: CodeBlock) = CodeBlock.of("%L.toLong()", value)
  },
  INTEGER(INT) {
    override fun decode(value: CodeBlock) = CodeBlock.of("%L.toInt()", value)

    override fun encode(value: CodeBlock) = CodeBlock.of("%L.toLong()", value)
  },
  BIG_INT(LONG);

  override fun prepareStatementBinder(columnIndex: String, value: CodeBlock): CodeBlock {
    return CodeBlock.builder()
      .add(
        when (this) {
          SMALL_INT, INTEGER, BIG_INT -> "bindLong"
        }
      )
      .add("($columnIndex, %L)\n", value)
      .build()
  }

  override fun cursorGetter(columnIndex: Int): CodeBlock {
    return CodeBlock.of(
      when (this) {
        SMALL_INT, INTEGER, BIG_INT -> "$CURSOR_NAME.getLong($columnIndex)"
      }
    )
  }
}

package com.squareup.sqldelight.core.dialect.sqlite

import com.squareup.kotlinpoet.ANY
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.DOUBLE
import com.squareup.kotlinpoet.LONG
import com.squareup.kotlinpoet.TypeName
import com.squareup.kotlinpoet.asClassName
import com.squareup.kotlinpoet.asTypeName
import com.squareup.sqldelight.core.dialect.api.DialectType
import com.squareup.sqldelight.core.lang.CURSOR_NAME

internal enum class SqliteType(override val javaType: TypeName) : DialectType {
  ARGUMENT(ANY.copy(nullable = true)),
  NULL(Nothing::class.asClassName().copy(nullable = true)),
  INTEGER(LONG),
  REAL(DOUBLE),
  TEXT(String::class.asTypeName()),
  BLOB(ByteArray::class.asTypeName());

  override fun prepareStatementBinder(columnIndex: String, value: CodeBlock): CodeBlock {
    return CodeBlock.builder()
      .add(
        when (this) {
          INTEGER -> "bindLong"
          REAL -> "bindDouble"
          TEXT -> "bindString"
          BLOB -> "bindBytes"
          else -> throw IllegalArgumentException("Cannot bind unknown types or null")
        }
      )
      .add("($columnIndex, %L)\n", value)
      .build()
  }

  override fun cursorGetter(columnIndex: Int): CodeBlock {
    return CodeBlock.of(
      when (this) {
        NULL -> "null"
        INTEGER -> "$CURSOR_NAME.getLong($columnIndex)"
        REAL -> "$CURSOR_NAME.getDouble($columnIndex)"
        TEXT -> "$CURSOR_NAME.getString($columnIndex)"
        BLOB -> "$CURSOR_NAME.getBytes($columnIndex)"
        ARGUMENT -> throw IllegalArgumentException("Cannot retrieve argument from cursor")
      }
    )
  }
}

package com.squareup.sqldelight.core.lang.psi

import com.alecstrong.sql.psi.core.psi.LazyQuery
import com.alecstrong.sql.psi.core.psi.SqlCompositeElementImpl
import com.intellij.lang.ASTNode
import com.intellij.psi.PsiElement
import com.squareup.sqldelight.core.lang.util.sqFile
import com.squareup.sqldelight.core.psi.SqlDelightStmtClojureStmtList

abstract class ClojureStmtListMixin(
  node: ASTNode
) : SqlCompositeElementImpl(node),
  SqlDelightStmtClojureStmtList {
  override fun tablesAvailable(child: PsiElement): Collection<LazyQuery> {
    return sqFile().tablesAvailable(child)
  }
}

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.core.lang.psi

import com.alecstrong.sql.psi.core.SqlAnnotationHolder
import com.alecstrong.sql.psi.core.psi.SqlColumnDef
import com.alecstrong.sql.psi.core.psi.SqlIdentifier
import com.alecstrong.sql.psi.core.psi.SqlTypeName
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.alecstrong.sql.psi.core.psi.impl.SqlColumnTypeImpl
import com.intellij.lang.ASTNode
import com.intellij.psi.PsiWhiteSpace
import com.intellij.psi.tree.TokenSet
import com.squareup.kotlinpoet.ARRAY
import com.squareup.kotlinpoet.AnnotationSpec
import com.squareup.kotlinpoet.BOOLEAN
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.DOUBLE
import com.squareup.kotlinpoet.FLOAT
import com.squareup.kotlinpoet.INT
import com.squareup.kotlinpoet.LONG
import com.squareup.kotlinpoet.ParameterizedTypeName
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.kotlinpoet.SHORT
import com.squareup.kotlinpoet.TypeName
import com.squareup.kotlinpoet.asClassName
import com.squareup.kotlinpoet.joinToCode
import com.squareup.sqldelight.core.compiler.SqlDelightCompiler.allocateName
import com.squareup.sqldelight.core.lang.IntermediateType
import com.squareup.sqldelight.core.lang.util.parentOfType
import com.squareup.sqldelight.core.psi.SqlDelightAnnotation
import com.squareup.sqldelight.core.psi.SqlDelightAnnotationValue
import com.squareup.sqldelight.core.psi.SqlDelightColumnType
import com.squareup.sqldelight.core.psi.SqlDelightJavaType
import com.squareup.sqldelight.core.psi.SqlDelightJavaTypeName
import com.squareup.sqldelight.core.psi.SqlDelightParameterizedJavaType
import com.squareup.sqldelight.core.psi.SqlDelightStmtList

internal abstract class ColumnTypeMixin(
  node: ASTNode
) : SqlColumnTypeImpl(node),
  TypedColumn,
  SqlDelightColumnType {
  override fun type(): IntermediateType {
    val columnName = (parent as SqlColumnDef).columnName
    val columnConstraintList = (parent as SqlColumnDef).columnConstraintList

    var type = typeName.type().copy(column = (parent as SqlColumnDef), name = allocateName(columnName))
    javaTypeName?.type()?.let { type = type.copy(javaType = it) }
    if (columnConstraintList.none {
      (it.node.findChildByType(SqlTypes.NOT) != null && it.node.findChildByType(SqlTypes.NULL) != null) ||
        it.node.findChildByType(SqlTypes.PRIMARY) != null
    }
    ) {
      type = type.asNullable()
    }
    if (annotationList.isNotEmpty()) {
      type = type.copy(
        javaType = type.javaType
          .copy(annotations = type.javaType.annotations + annotationList.map { it.spec() })
      )
    }
    return type
  }

  override fun adapter(): PropertySpec? {
    val columnName = (parent as SqlColumnDef).columnName
    javaTypeName?.let {
      val customType = try {
        it.parameterizedJavaType?.type() ?: return null
      } catch (e: IllegalArgumentException) {
        // Found an invalid type.
        return null
      }
      return PropertySpec
        .builder(
          name = "${allocateName(columnName)}Adapter",
          type = columnAdapterType.parameterizedBy(customType, typeName.type().dialectType.javaType)
        )
        .build()
    }
    return null
  }

  private fun SqlDelightJavaTypeName.type(): TypeName? {
    try {
      return parameterizedJavaType?.type() ?: kotlinType(text)
    } catch (e: IllegalArgumentException) {
      // Found an invalid type.
      return null
    }
  }

  private fun SqlDelightJavaType.type(): ClassName {
    parentOfType<SqlDelightStmtList>().importStmtList.importStmtList.forEach { import ->
      val typePrefix = text.substringBefore('.')
      if (import.javaType.text.endsWith(".$typePrefix")) {
        return text.split(".").drop(1).fold(import.javaType.type()) { current, nested ->
          current.nestedClass(nested)
        }
      }
    }
    return ClassName.bestGuess(text)
  }

  private fun SqlDelightParameterizedJavaType.type(): TypeName {
    if (javaTypeNameList.isNotEmpty() || javaTypeName2List.isNotEmpty()) {
      var parameters = javaTypeNameList.mapNotNull { it.type() }.toTypedArray()
      if (javaTypeList.size == 2) {
        // Hack to get around '>>' character.
        parameters += javaTypeList[1].type()
          .parameterizedBy(*javaTypeName2List.map { it.javaType.type() }.toTypedArray())
      }
      return javaTypeList[0].type().parameterizedBy(*parameters)
    }
    return javaTypeList[0].type()
  }

  private fun SqlDelightAnnotation.spec(): AnnotationSpec {
    val annotation = AnnotationSpec.builder(javaType.type())
    val identifiers = children.filterIsInstance<SqlIdentifier>()
    if (identifiers.isEmpty() && annotationValueList.isNotEmpty()) {
      annotation.addMember(annotationValueList[0].value())
    }
    annotationValueList.zip(
      identifiers,
      { annotation_value, identifier ->
        annotation.addMember(
          CodeBlock.builder()
            .add("${identifier.text} = ")
            .add(annotation_value.value())
            .build()
        )
      }
    )
    return annotation.build()
  }

  private fun SqlDelightAnnotationValue.value(): CodeBlock {
    javaTypeName?.let {
      return CodeBlock.of("%T::class", it.type())
    }
    annotation?.let {
      return CodeBlock.of("%L", it.spec())
    }
    if (annotationValueList.isNotEmpty()) {
      return annotationValueList.map { it.value() }.joinToCode(",", "[", "]")
    }
    return CodeBlock.of(text)
  }

  override fun annotate(annotationHolder: SqlAnnotationHolder) {
    javaTypeName?.let { javaType ->
      if (javaType.type() == null) {
        annotationHolder.createErrorAnnotation(javaType, "Unknown type ${javaType.text}")
      }
    }
    val children = node.getChildren(TokenSet.create(SqlTypes.ID))
    children.filter { (it.text == "as" || it.text == "As") && it.prevVisibleSibling?.psi is SqlTypeName }
      .forEach {
        annotationHolder.createErrorAnnotation(it.psi, "Expected 'AS', got '${it.text}'")
      }
  }

  private val ASTNode.prevVisibleSibling: ASTNode?
    get() = generateSequence(treePrev) { it.treePrev }
      .firstOrNull { it.psi !is PsiWhiteSpace }

  companion object {
    private val columnAdapterType = ClassName("com.squareup.sqldelight", "ColumnAdapter")

    internal fun kotlinType(text: String) = when (text) {
      "Integer", "Int" -> INT
      "Boolean" -> BOOLEAN
      "Short" -> SHORT
      "Long" -> LONG
      "Float" -> FLOAT
      "Double" -> DOUBLE
      "String" -> String::class.asClassName()
      "ByteArray" -> ByteArray::class.asClassName()
      else -> null
    }

    internal val TypeName.isArrayType get() = when (this) {
      is ParameterizedTypeName -> rawType == ARRAY
      BooleanArray::class.asClassName(),
      ByteArray::class.asClassName(),
      CharArray::class.asClassName(),
      DoubleArray::class.asClassName(),
      FloatArray::class.asClassName(),
      IntArray::class.asClassName(),
      LongArray::class.asClassName(),
      ShortArray::class.asClassName() -> true
      else -> false
    }
  }
}

package com.squareup.sqldelight.core.lang.psi

import com.alecstrong.sql.psi.core.DialectPreset
import com.alecstrong.sql.psi.core.SqlAnnotationHolder
import com.alecstrong.sql.psi.core.psi.SqlExpr
import com.alecstrong.sql.psi.core.psi.SqlResultColumn
import com.alecstrong.sql.psi.core.psi.impl.SqlFunctionExprImpl
import com.intellij.lang.ASTNode
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType
import com.squareup.sqldelight.core.lang.IntermediateType
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.util.encapsulatingType
import com.squareup.sqldelight.core.lang.util.type

internal class FunctionExprMixin(node: ASTNode?) : SqlFunctionExprImpl(node) {
  fun argumentType(expr: SqlExpr) = when (functionName.text.toLowerCase()) {
    "instr" -> when (expr) {
      exprList.getOrNull(1) -> IntermediateType(SqliteType.TEXT)
      else -> functionType()
    }
    "ifnull", "coalesce" -> functionType()?.asNullable()
    else -> functionType()
  }

  fun functionType() = when (functionName.text.toLowerCase()) {
    "round" -> {
      // Single arg round function returns an int. Otherwise real.
      if (exprList.size == 1) {
        IntermediateType(SqliteType.INTEGER).nullableIf(exprList[0].type().javaType.isNullable)
      } else {
        IntermediateType(SqliteType.REAL).nullableIf(exprList.any { it.type().javaType.isNullable })
      }
    }

    /**
     * sum's output is always nullable because it returns NULL for an input that's empty or only contains NULLs.
     *
     * https://www.sqlite.org/lang_aggfunc.html#sumunc
     * >>> The result of sum() is an integer value if all non-NULL inputs are integers. If any input to sum() is neither
     * >>> an integer or a NULL then sum() returns a floating point value which might be an approximation to the true sum.
     *
     */
    "sum" -> {
      val type = exprList[0].type()
      if (type.dialectType == SqliteType.INTEGER && !type.javaType.isNullable) {
        type.asNullable()
      } else {
        IntermediateType(SqliteType.REAL).asNullable()
      }
    }

    "lower", "ltrim", "replace", "rtrim", "substr", "trim", "upper", "group_concat" -> {
      IntermediateType(SqliteType.TEXT).nullableIf(exprList[0].type().javaType.isNullable)
    }

    "date", "time", "char", "hex", "quote", "soundex", "typeof" -> {
      IntermediateType(SqliteType.TEXT)
    }

    "random", "count" -> {
      IntermediateType(SqliteType.INTEGER)
    }

    "instr", "length" -> {
      IntermediateType(SqliteType.INTEGER).nullableIf(exprList.any { it.type().javaType.isNullable })
    }

    "avg" -> IntermediateType(SqliteType.REAL).asNullable()
    "abs" -> exprList[0].type()
    "coalesce", "ifnull" -> encapsulatingType(exprList, SqliteType.INTEGER, SqliteType.REAL, SqliteType.TEXT, SqliteType.BLOB)
    "nullif" -> exprList[0].type().asNullable()
    "max" -> encapsulatingType(exprList, SqliteType.INTEGER, SqliteType.REAL, SqliteType.TEXT, SqliteType.BLOB).asNullable()
    "min" -> encapsulatingType(exprList, SqliteType.BLOB, SqliteType.TEXT, SqliteType.INTEGER, SqliteType.REAL).asNullable()
    else -> when ((containingFile as SqlDelightFile).dialect) {
      DialectPreset.SQLITE_3_18, DialectPreset.SQLITE_3_24, DialectPreset.SQLITE_3_25 -> sqliteFunctionType()
      DialectPreset.MYSQL -> mySqlFunctionType()
      DialectPreset.POSTGRESQL -> postgreSqlFunctionType()
      else -> null
    }
  }

  private fun sqliteFunctionType() = when (functionName.text.toLowerCase()) {
    "printf" -> IntermediateType(SqliteType.TEXT).nullableIf(exprList[0].type().javaType.isNullable)
    "datetime", "julianday", "strftime", "sqlite_compileoption_get", "sqlite_source_id", "sqlite_version" -> {
      IntermediateType(SqliteType.TEXT)
    }
    "changes", "last_insert_rowid", "sqlite_compileoption_used", "total_changes" -> {
      IntermediateType(SqliteType.INTEGER)
    }
    "unicode" -> {
      IntermediateType(SqliteType.INTEGER).nullableIf(exprList.any { it.type().javaType.isNullable })
    }
    "randomblob", "zeroblob" -> IntermediateType(SqliteType.BLOB)
    "total", "bm25" -> IntermediateType(SqliteType.REAL)
    "likelihood", "likely", "unlikely" -> exprList[0].type()
    "highlight", "snippet" -> IntermediateType(SqliteType.TEXT).asNullable()
    "offsets" -> IntermediateType(SqliteType.TEXT).asNullable()
    else -> null
  }

  private fun mySqlFunctionType() = when (functionName.text.toLowerCase()) {
    "greatest" -> encapsulatingType(exprList, SqliteType.INTEGER, SqliteType.REAL, SqliteType.TEXT, SqliteType.BLOB)
    "concat" -> encapsulatingType(exprList, SqliteType.TEXT)
    "last_insert_id" -> IntermediateType(SqliteType.INTEGER)
    "month", "year", "minute" -> IntermediateType(SqliteType.INTEGER)
    "sin", "cos", "tan" -> IntermediateType(SqliteType.REAL)
    else -> null
  }

  private fun postgreSqlFunctionType() = when (functionName.text.toLowerCase()) {
    "greatest" -> encapsulatingType(exprList, SqliteType.INTEGER, SqliteType.REAL, SqliteType.TEXT, SqliteType.BLOB)
    "concat" -> encapsulatingType(exprList, SqliteType.TEXT)
    "substring" -> IntermediateType(SqliteType.TEXT).nullableIf(exprList[0].type().javaType.isNullable)
    else -> null
  }

  override fun annotate(annotationHolder: SqlAnnotationHolder) {
    if (parent is SqlResultColumn && functionType() == null) {
      annotationHolder.createErrorAnnotation(this, "Unknown function ${functionName.text}")
    }
  }
}

package com.squareup.sqldelight.core.lang.psi

import com.alecstrong.sql.psi.core.SqlAnnotationHolder
import com.alecstrong.sql.psi.core.psi.SqlAnnotatedElement
import com.intellij.extapi.psi.ASTWrapperPsiElement
import com.intellij.lang.ASTNode
import com.squareup.sqldelight.core.psi.SqlDelightImportStmt
import com.squareup.sqldelight.core.psi.SqlDelightImportStmtList

abstract class ImportStmtMixin(
  node: ASTNode
) : ASTWrapperPsiElement(node),
  SqlDelightImportStmt,
  SqlAnnotatedElement {
  private fun type(): String {
    return javaType.text.substringAfterLast(".")
  }

  override fun annotate(annotationHolder: SqlAnnotationHolder) {
    if ((parent as SqlDelightImportStmtList).importStmtList
      .filterIsInstance<ImportStmtMixin>()
      .filter { it != this }
      .any { it.type() == type() }
    ) {
      annotationHolder.createErrorAnnotation(this, "Multiple imports for type ${type()}")
    }
  }
}

package com.squareup.sqldelight.core.lang.psi

import com.alecstrong.sql.psi.core.SqlAnnotationHolder
import com.alecstrong.sql.psi.core.psi.SqlColumnName
import com.alecstrong.sql.psi.core.psi.impl.SqlInsertStmtValuesImpl
import com.alecstrong.sql.psi.core.psi.mixins.ColumnDefMixin
import com.intellij.lang.ASTNode
import com.squareup.sqldelight.core.lang.acceptsTableInterface
import com.squareup.sqldelight.core.psi.SqlDelightInsertStmtValues

open class InsertStmtValuesMixin(
  node: ASTNode
) : SqlInsertStmtValuesImpl(node),
  SqlDelightInsertStmtValues {
  override fun annotate(annotationHolder: SqlAnnotationHolder) {
    val parent = parent ?: return
    if (parent.acceptsTableInterface()) {
      val table = tableAvailable(this, parent.tableName.name).firstOrNull() ?: return
      val columns = table.columns.map { (it.element as SqlColumnName).name }
      val setColumns =
        if (parent.columnNameList.isEmpty()) {
          columns
        } else {
          parent.columnNameList.mapNotNull { it.name }
        }

      val needsDefaultValue = table.columns
        .filter { (element, _) ->
          element is SqlColumnName &&
            element.name !in setColumns &&
            !(element.parent as ColumnDefMixin).hasDefaultValue()
        }
        .map { it.element as SqlColumnName }
      if (needsDefaultValue.size == 1) {
        annotationHolder.createErrorAnnotation(
          parent,
          "Cannot populate default value for column " +
            "${needsDefaultValue.first().name}, it must be specified in insert statement."
        )
      } else if (needsDefaultValue.size > 1) {
        annotationHolder.createErrorAnnotation(
          parent,
          "Cannot populate default values for columns " +
            "(${needsDefaultValue.joinToString { it.name }}), they must be specified in insert statement."
        )
      }

      // This is going to break error handling in sqlite-psi so just omit the superclass annotator.
      return
    }

    super.annotate(annotationHolder)
  }
}

package com.squareup.sqldelight.core.lang.psi

import com.intellij.extapi.psi.ASTWrapperPsiElement
import com.intellij.lang.ASTNode
import com.intellij.openapi.module.ModuleUtilCore.findModuleForPsiElement
import com.intellij.openapi.util.TextRange
import com.intellij.psi.JavaPsiFacade
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiReference
import com.intellij.psi.PsiReferenceBase
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.util.findChildrenOfType
import com.squareup.sqldelight.core.psi.SqlDelightImportStmt

abstract class JavaTypeMixin(
  node: ASTNode
) : ASTWrapperPsiElement(node) {
  override fun getReference(): PsiReference = JavaTypeReference()

  private inner class JavaTypeReference : PsiReferenceBase<JavaTypeMixin>(
    this, TextRange(0, textLength)
  ) {
    override fun getVariants(): Array<Any> = emptyArray()
    override fun resolve(): PsiElement? {
      val qualifiedType = if (parent is SqlDelightImportStmt) {
        text
      } else {
        val prefix = text.substringBefore('.')
        (containingFile as SqlDelightFile).sqlStmtList
          ?.findChildrenOfType<ImportStmtMixin>()
          ?.firstOrNull { it.javaType.text.endsWith(prefix) }
          ?.javaType?.text?.plus(text.removePrefix(prefix)) ?: typeForThisPackage(text)
      }

      return JavaPsiFacade.getInstance(project).findClass(
        qualifiedType,
        findModuleForPsiElement(element)!!.getModuleWithDependenciesAndLibrariesScope(false)
      )
    }

    private fun typeForThisPackage(text: String) = when (ColumnTypeMixin.kotlinType(text)) {
      null -> "${(containingFile as SqlDelightFile).packageName}.$text"
      else -> text
    }
  }
}

package com.squareup.sqldelight.core.lang.psi

import com.alecstrong.sql.psi.core.psi.SqlIdentifier
import com.intellij.psi.PsiNamedElement

interface StmtIdentifier : PsiNamedElement {
  fun identifier(): SqlIdentifier?
}

package com.squareup.sqldelight.core.lang.psi

import com.alecstrong.sql.psi.core.SqlAnnotationHolder
import com.alecstrong.sql.psi.core.SqlParser
import com.alecstrong.sql.psi.core.SqlParserDefinition
import com.alecstrong.sql.psi.core.psi.SqlAnnotatedElement
import com.alecstrong.sql.psi.core.psi.SqlIdentifier
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.intellij.extapi.psi.ASTWrapperPsiElement
import com.intellij.lang.ASTNode
import com.intellij.lang.LanguageParserDefinitions
import com.intellij.lang.PsiBuilderFactory
import com.intellij.lang.parser.GeneratedParserUtilBase
import com.intellij.psi.PsiElement
import com.intellij.psi.impl.GeneratedMarkerVisitor
import com.intellij.psi.impl.source.tree.TreeElement
import com.squareup.sqldelight.core.lang.SqlDelightQueriesFile
import com.squareup.sqldelight.core.psi.SqlDelightStmtIdentifier
import com.squareup.sqldelight.core.psi.SqlDelightStmtIdentifierClojure

abstract class StmtIdentifierMixin(
  node: ASTNode
) : ASTWrapperPsiElement(node),
  SqlDelightStmtIdentifier,
  SqlDelightStmtIdentifierClojure,
  SqlAnnotatedElement {
  override fun getName() = identifier()?.text

  override fun setName(name: String): PsiElement {
    val parserDefinition = LanguageParserDefinitions.INSTANCE.forLanguage(language) as SqlParserDefinition
    var builder = PsiBuilderFactory.getInstance().createBuilder(
      project, node, parserDefinition.createLexer(project), language, name
    )
    builder = GeneratedParserUtilBase.adapt_builder_(
      SqlTypes.IDENTIFIER, builder, SqlParser(), SqlParser.EXTENDS_SETS_
    )
    GeneratedParserUtilBase.ErrorState.get(builder).currentFrame = GeneratedParserUtilBase.Frame()

    SqlParser.identifier_real(builder, 0)
    val element = builder.treeBuilt
    (element as TreeElement).acceptTree(GeneratedMarkerVisitor())
    node.replaceChild(identifier()!!.node, element)
    return this
  }

  override fun annotate(annotationHolder: SqlAnnotationHolder) {
    if (name != null && (containingFile as SqlDelightQueriesFile).sqliteStatements()
      .filterNot { it.identifier == this }
      .any { it.identifier.name == name }
    ) {
      annotationHolder.createErrorAnnotation(this, "Duplicate SQL identifier")
    }
  }

  override fun identifier() = children.filterIsInstance<SqlIdentifier>().singleOrNull()
}

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.core.lang.psi

import com.alecstrong.sql.psi.core.psi.SqlColumnName
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.sqldelight.core.lang.IntermediateType
import com.squareup.sqldelight.core.lang.util.type

internal interface TypedColumn {
  /**
   * @return the java type for this column.
   */
  fun type(): IntermediateType

  /**
   * @return the adapter property which will include the type and name of the adapter. If this
   * column does not require an adapter, returns null.
   */
  fun adapter(): PropertySpec?
}

fun SqlColumnName.isColumnSameAs(other: SqlColumnName) = type().let { it.column != null && it.column == other.type().column }
fun SqlColumnName.isTypeSameAs(other: SqlColumnName) = type().javaType == other.type().javaType

package com.squareup.sqldelight.core.lang.psi

import com.alecstrong.sql.psi.core.hsql.psi.HsqlTypeName
import com.alecstrong.sql.psi.core.mysql.psi.MySqlTypeName
import com.alecstrong.sql.psi.core.postgresql.psi.PostgreSqlTypeName
import com.alecstrong.sql.psi.core.psi.SqlTypeName
import com.alecstrong.sql.psi.core.sqlite_3_18.psi.SqliteTypeName
import com.squareup.sqldelight.core.dialect.hsql.HsqlType
import com.squareup.sqldelight.core.dialect.mysql.MySqlType
import com.squareup.sqldelight.core.dialect.postgresql.PostgreSqlType
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType
import com.squareup.sqldelight.core.lang.IntermediateType

internal fun SqlTypeName.type(): IntermediateType {
  return when (this) {
    is SqliteTypeName -> type()
    is MySqlTypeName -> type()
    is PostgreSqlTypeName -> type()
    is HsqlTypeName -> type()
    else -> throw IllegalArgumentException("Unknown sql type $this")
  }
}

private fun SqliteTypeName.type(): IntermediateType {
  return when {
    textDataType != null -> IntermediateType(SqliteType.TEXT)
    blobDataType != null -> IntermediateType(SqliteType.BLOB)
    intDataType != null -> IntermediateType(SqliteType.INTEGER)
    realDataType != null -> IntermediateType(SqliteType.REAL)
    else -> throw IllegalArgumentException("Unknown sql type $text")
  }
}

private fun MySqlTypeName.type(): IntermediateType {
  return when {
    approximateNumericDataType != null -> IntermediateType(SqliteType.REAL)
    binaryDataType != null -> IntermediateType(SqliteType.BLOB)
    dateDataType != null -> IntermediateType(SqliteType.TEXT)
    tinyIntDataType != null -> if (tinyIntDataType!!.text == "BOOLEAN") {
      IntermediateType(MySqlType.TINY_INT_BOOL)
    } else {
      IntermediateType(MySqlType.TINY_INT)
    }
    smallIntDataType != null -> IntermediateType(MySqlType.SMALL_INT)
    mediumIntDataType != null -> IntermediateType(MySqlType.INTEGER)
    intDataType != null -> IntermediateType(MySqlType.INTEGER)
    bigIntDataType != null -> IntermediateType(MySqlType.BIG_INT)
    fixedPointDataType != null -> IntermediateType(SqliteType.INTEGER)
    jsonDataType != null -> IntermediateType(SqliteType.TEXT)
    enumSetType != null -> IntermediateType(SqliteType.TEXT)
    characterType != null -> IntermediateType(SqliteType.TEXT)
    bitDataType != null -> IntermediateType(MySqlType.BIT)
    else -> throw IllegalArgumentException("Unknown kotlin type for sql type ${this.text}")
  }
}

private fun PostgreSqlTypeName.type(): IntermediateType {
  return when {
    smallIntDataType != null -> IntermediateType(PostgreSqlType.SMALL_INT)
    intDataType != null -> IntermediateType(PostgreSqlType.INTEGER)
    bigIntDataType != null -> IntermediateType(PostgreSqlType.BIG_INT)
    numericDataType != null -> IntermediateType(SqliteType.REAL)
    approximateNumericDataType != null -> IntermediateType(SqliteType.REAL)
    stringDataType != null -> IntermediateType(SqliteType.TEXT)
    smallSerialDataType != null -> IntermediateType(PostgreSqlType.SMALL_INT)
    serialDataType != null -> IntermediateType(PostgreSqlType.INTEGER)
    bigSerialDataType != null -> IntermediateType(PostgreSqlType.BIG_INT)
    dateDataType != null -> IntermediateType(SqliteType.TEXT)
    jsonDataType != null -> IntermediateType(SqliteType.TEXT)
    else -> throw IllegalArgumentException("Unknown kotlin type for sql type ${this.text}")
  }
}

private fun HsqlTypeName.type(): IntermediateType {
  return when {
    approximateNumericDataType != null -> IntermediateType(SqliteType.REAL)
    binaryStringDataType != null -> IntermediateType(SqliteType.BLOB)
    dateDataType != null -> IntermediateType(SqliteType.TEXT)
    tinyIntDataType != null -> IntermediateType(HsqlType.TINY_INT)
    smallIntDataType != null -> IntermediateType(HsqlType.SMALL_INT)
    intDataType != null -> IntermediateType(HsqlType.INTEGER)
    bigIntDataType != null -> IntermediateType(HsqlType.BIG_INT)
    fixedPointDataType != null -> IntermediateType(SqliteType.INTEGER)
    characterStringDataType != null -> IntermediateType(SqliteType.TEXT)
    booleanDataType != null -> IntermediateType(HsqlType.BOOL)
    bitStringDataType != null -> IntermediateType(SqliteType.BLOB)
    intervalDataType != null -> IntermediateType(SqliteType.BLOB)
    else -> throw IllegalArgumentException("Unknown kotlin type for sql type ${this.text}")
  }
}

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.core.lang.util

import com.alecstrong.sql.psi.core.mysql.psi.MySqlExtensionExpr
import com.alecstrong.sql.psi.core.psi.SqlBetweenExpr
import com.alecstrong.sql.psi.core.psi.SqlBinaryExpr
import com.alecstrong.sql.psi.core.psi.SqlBinaryLikeExpr
import com.alecstrong.sql.psi.core.psi.SqlBindExpr
import com.alecstrong.sql.psi.core.psi.SqlCaseExpr
import com.alecstrong.sql.psi.core.psi.SqlCastExpr
import com.alecstrong.sql.psi.core.psi.SqlCollateExpr
import com.alecstrong.sql.psi.core.psi.SqlCompoundSelectStmt
import com.alecstrong.sql.psi.core.psi.SqlExpr
import com.alecstrong.sql.psi.core.psi.SqlInExpr
import com.alecstrong.sql.psi.core.psi.SqlInsertStmt
import com.alecstrong.sql.psi.core.psi.SqlInsertStmtValues
import com.alecstrong.sql.psi.core.psi.SqlIsExpr
import com.alecstrong.sql.psi.core.psi.SqlLimitingTerm
import com.alecstrong.sql.psi.core.psi.SqlNullExpr
import com.alecstrong.sql.psi.core.psi.SqlParenExpr
import com.alecstrong.sql.psi.core.psi.SqlResultColumn
import com.alecstrong.sql.psi.core.psi.SqlSelectStmt
import com.alecstrong.sql.psi.core.psi.SqlSetterExpression
import com.alecstrong.sql.psi.core.psi.SqlUnaryExpr
import com.alecstrong.sql.psi.core.psi.SqlUpdateStmt
import com.alecstrong.sql.psi.core.psi.SqlUpdateStmtLimited
import com.alecstrong.sql.psi.core.psi.SqlUpdateStmtSubsequentSetter
import com.alecstrong.sql.psi.core.psi.SqlValuesExpression
import com.alecstrong.sql.psi.core.sqlite_3_24.psi.SqliteUpsertDoUpdate
import com.intellij.psi.PsiElement
import com.intellij.psi.util.PsiTreeUtil
import com.squareup.kotlinpoet.BOOLEAN
import com.squareup.kotlinpoet.asClassName
import com.squareup.sqldelight.core.compiler.model.NamedQuery
import com.squareup.sqldelight.core.dialect.mysql.argumentType
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.ARGUMENT
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.INTEGER
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.NULL
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.TEXT
import com.squareup.sqldelight.core.lang.IntermediateType
import com.squareup.sqldelight.core.lang.psi.FunctionExprMixin

/**
 * Return the expected type for this expression, which is the argument type exposed in the generated
 * api.
 */
internal fun SqlBindExpr.argumentType(): IntermediateType {
  return inferredType().copy(bindArg = this)
}

internal fun SqlBindExpr.isArrayParameter(): Boolean {
  return (parent is SqlInExpr && this == parent.lastChild)
}

private fun SqlExpr.inferredType(): IntermediateType {
  return when (val parentRule = parent!!) {
    is SqlExpr -> {
      val result = parentRule.argumentType(this)
      if (result.dialectType == ARGUMENT) {
        parentRule.inferredType()
      } else {
        result
      }
    }

    is SqlValuesExpression -> parentRule.argumentType(this)
    is SqlSetterExpression -> parentRule.argumentType()
    is SqlLimitingTerm -> IntermediateType(INTEGER)
    is SqlResultColumn -> {
      (parentRule.parent as SqlSelectStmt).argumentType(parentRule)
        ?: IntermediateType(NULL, Any::class.asClassName())
    }
    else -> IntermediateType(NULL, Any::class.asClassName())
  }
}

/**
 * Return the expected type for [argument], which is the argument type exposed in the generated api.
 */
private fun SqlExpr.argumentType(argument: SqlExpr): IntermediateType {
  return when (this) {
    is SqlInExpr -> {
      if (argument === firstChild) return IntermediateType(ARGUMENT)

      return exprList.first().type()
    }

    is SqlCaseExpr -> {
      fun PsiElement.isCaseResult() = PsiTreeUtil.skipWhitespacesBackward(this)?.text in listOf("THEN", "ELSE")
      fun PsiElement.isCondition() = PsiTreeUtil.skipWhitespacesBackward(this)?.text in listOf("CASE", "WHEN")

      return if (argument.isCaseResult()) {
        val validOtherArg = children.lastOrNull { it is SqlExpr && it !== argument && it !is SqlBindExpr && it.isCaseResult() }
        return validOtherArg?.type() ?: inferredType()
      } else if (argument.isCondition()) {
        val validOtherCondition = children.lastOrNull { it is SqlExpr && it !== argument && it !is SqlBindExpr && it.isCondition() }
        return validOtherCondition?.type() ?: IntermediateType(INTEGER, BOOLEAN)
      } else IntermediateType(INTEGER, BOOLEAN)
    }
    is SqlBetweenExpr, is SqlIsExpr, is SqlBinaryExpr -> {
      val validArg = children.lastOrNull { it is SqlExpr && it !== argument && it !is SqlBindExpr }
      validArg?.type() ?: children.last { it is SqlExpr && it !== argument }.type()
    }

    is SqlNullExpr -> IntermediateType(NULL).asNullable()
    is SqlBinaryLikeExpr -> {
      val other = children.last { it is SqlExpr && it !== argument }.type()
      IntermediateType(TEXT).copy(name = other.name)
    }

    is SqlCollateExpr, is SqlCastExpr, is SqlParenExpr, is SqlUnaryExpr -> {
      return IntermediateType(ARGUMENT)
    }

    is FunctionExprMixin -> {
      return argumentType(argument) ?: IntermediateType(NULL)
    }

    is MySqlExtensionExpr -> return argumentType(argument)
    else -> throw AssertionError()
  }
}

private fun SqlValuesExpression.argumentType(expression: SqlExpr): IntermediateType {
  val argumentIndex = children.indexOf(expression)
  if (argumentIndex == -1) throw AssertionError()

  return when (val parentRule = parent!!) {
    is SqlInsertStmtValues -> {
      val insertStmt = parentRule.parent as SqlInsertStmt
      if (insertStmt.columnNameList.isNotEmpty()) {
        insertStmt.columnNameList[argumentIndex].type()
      } else {
        insertStmt.columns[argumentIndex].type()
      }
    }
    is SqlSelectStmt -> {
      val compoundSelect = parentRule.parent as SqlCompoundSelectStmt
      NamedQuery("temp", compoundSelect).resultColumns[argumentIndex]
    }

    else -> throw AssertionError()
  }
}

private fun SqlSelectStmt.argumentType(result: SqlResultColumn): IntermediateType? {
  val index = resultColumnList.indexOf(result)
  val compoundSelect = parent!! as SqlCompoundSelectStmt

  val parentRule = compoundSelect.parent ?: return null
  return when (parentRule) {
    is SqlInsertStmtValues -> {
      val insertStmt = parentRule.parent as SqlInsertStmt
      insertStmt.columns[index].type()
    }

    else -> {
      // Check if this is part of an inner expression of a resulit column.
      val parentResult = PsiTreeUtil.getParentOfType(parentRule, SqlResultColumn::class.java)
        ?: return null
      (parentResult.parent as SqlSelectStmt).argumentType(parentResult)
    }
  }
}

private fun SqlSetterExpression.argumentType(): IntermediateType {
  return when (val parentRule = parent!!) {
    is SqlUpdateStmt -> parentRule.columnName!!.type()
    is SqlUpdateStmtLimited -> parentRule.columnName!!.type()
    is SqlUpdateStmtSubsequentSetter -> parentRule.columnName!!.type()
    is SqliteUpsertDoUpdate -> expr.type()
    else -> throw AssertionError()
  }
}

package com.squareup.sqldelight.core.lang.util

import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.squareup.kotlinpoet.ClassName
import com.squareup.sqldelight.core.compiler.SqlDelightCompiler.allocateName

internal val SqlCreateTableStmt.interfaceType: ClassName
  get() = ClassName(sqFile().packageName!!, allocateName(tableName).capitalize())

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.core.lang.util

import com.alecstrong.sql.psi.core.mysql.psi.MySqlExtensionExpr
import com.alecstrong.sql.psi.core.psi.SqlBetweenExpr
import com.alecstrong.sql.psi.core.psi.SqlBinaryAddExpr
import com.alecstrong.sql.psi.core.psi.SqlBinaryExpr
import com.alecstrong.sql.psi.core.psi.SqlBinaryLikeExpr
import com.alecstrong.sql.psi.core.psi.SqlBinaryMultExpr
import com.alecstrong.sql.psi.core.psi.SqlBindExpr
import com.alecstrong.sql.psi.core.psi.SqlCaseExpr
import com.alecstrong.sql.psi.core.psi.SqlCastExpr
import com.alecstrong.sql.psi.core.psi.SqlCollateExpr
import com.alecstrong.sql.psi.core.psi.SqlColumnExpr
import com.alecstrong.sql.psi.core.psi.SqlExistsExpr
import com.alecstrong.sql.psi.core.psi.SqlExpr
import com.alecstrong.sql.psi.core.psi.SqlFunctionExpr
import com.alecstrong.sql.psi.core.psi.SqlInExpr
import com.alecstrong.sql.psi.core.psi.SqlIsExpr
import com.alecstrong.sql.psi.core.psi.SqlLiteralExpr
import com.alecstrong.sql.psi.core.psi.SqlNullExpr
import com.alecstrong.sql.psi.core.psi.SqlOtherExpr
import com.alecstrong.sql.psi.core.psi.SqlParenExpr
import com.alecstrong.sql.psi.core.psi.SqlRaiseExpr
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.alecstrong.sql.psi.core.psi.SqlUnaryExpr
import com.intellij.psi.tree.TokenSet
import com.squareup.kotlinpoet.BOOLEAN
import com.squareup.sqldelight.core.compiler.SqlDelightCompiler.allocateName
import com.squareup.sqldelight.core.dialect.mysql.type
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.ARGUMENT
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.BLOB
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.INTEGER
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.NULL
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.REAL
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.TEXT
import com.squareup.sqldelight.core.lang.IntermediateType
import com.squareup.sqldelight.core.lang.psi.FunctionExprMixin
import com.squareup.sqldelight.core.lang.psi.type

internal val SqlExpr.name: String get() = when (this) {
  is SqlCastExpr -> expr.name
  is SqlParenExpr -> expr?.name ?: "value"
  is SqlFunctionExpr -> functionName.text
  is SqlColumnExpr -> allocateName(columnName)
  else -> "expr"
}

/**
 * If javaType is true, this will return a possible more descriptive type for column expressions.
 *
 * Order of operations:
 * expr ::= ( raise_expr
 *          | case_expr
 *          | exists_expr
 *          | in_expr
 *          | between_expr
 *          | is_expr
 *          | null_expr
 *          | like_expr
 *          | collate_expr
 *          | cast_expr
 *          | paren_expr
 *          | function_expr
 *          | binary_expr
 *          | unary_expr
 *          | bind_expr
 *          | literal_expr
 *          | column_expr )
 */
internal fun SqlExpr.type(): IntermediateType = when (this) {
  is SqlRaiseExpr -> IntermediateType(NULL)
  is SqlCaseExpr -> childOfType(SqlTypes.THEN)!!.nextSiblingOfType<SqlExpr>().type()

  is SqlExistsExpr -> {
    val isExists = childOfType(SqlTypes.EXISTS) != null
    if (isExists) {
      IntermediateType(INTEGER, BOOLEAN)
    } else {
      compoundSelectStmt.queryExposed().single().columns.single().element.type()
    }
  }

  is SqlInExpr -> IntermediateType(INTEGER, BOOLEAN)
  is SqlBetweenExpr -> IntermediateType(INTEGER, BOOLEAN)
  is SqlIsExpr -> IntermediateType(INTEGER, BOOLEAN)
  is SqlNullExpr -> IntermediateType(INTEGER, BOOLEAN)
  is SqlBinaryLikeExpr -> IntermediateType(INTEGER, BOOLEAN)
  is SqlCollateExpr -> expr.type()
  is SqlCastExpr -> typeName.type().nullableIf(expr.type().javaType.isNullable)
  is SqlParenExpr -> expr?.type() ?: IntermediateType(NULL)
  is FunctionExprMixin -> functionType() ?: IntermediateType(NULL)

  is SqlBinaryExpr -> {
    if (childOfType(
        TokenSet.create(
            SqlTypes.EQ, SqlTypes.EQ2, SqlTypes.NEQ,
            SqlTypes.NEQ2, SqlTypes.AND, SqlTypes.OR, SqlTypes.GT, SqlTypes.GTE,
            SqlTypes.LT, SqlTypes.LTE
          )
      ) != null
    ) {
      IntermediateType(INTEGER, BOOLEAN)
    } else {
      encapsulatingType(
        exprList = getExprList(),
        nullableIfAny = (this is SqlBinaryAddExpr || this is SqlBinaryMultExpr),
        INTEGER, REAL, TEXT, BLOB
      )
    }
  }

  is SqlUnaryExpr -> expr.type()

  is SqlBindExpr -> IntermediateType(ARGUMENT)

  is SqlLiteralExpr -> when {
    (literalValue.stringLiteral != null) -> IntermediateType(TEXT)
    (literalValue.blobLiteral != null) -> IntermediateType(BLOB)
    (literalValue.numericLiteral != null) -> {
      if (literalValue.text.contains('.')) {
        IntermediateType(REAL)
      } else {
        IntermediateType(INTEGER)
      }
    }
    (
      literalValue.childOfType(
        TokenSet.create(
          SqlTypes.CURRENT_TIMESTAMP,
          SqlTypes.CURRENT_TIME, SqlTypes.CURRENT_DATE
        )
      ) != null
      ) -> IntermediateType(TEXT)
    (literalValue.childOfType(SqlTypes.NULL) != null) -> IntermediateType(NULL)
    else -> IntermediateType(BLOB).asNullable()
  }

  is SqlColumnExpr -> columnName.type()

  is SqlOtherExpr -> {
    extensionExpr.type()
  }

  is MySqlExtensionExpr -> type()
  else -> throw IllegalStateException("Unknown expression type $this")
}

/**
 * @return the type from the expr list which is the highest order in the typeOrder list
 */
internal fun encapsulatingType(
  exprList: List<SqlExpr>,
  vararg typeOrder: SqliteType
) = encapsulatingType(exprList = exprList, nullableIfAny = false, typeOrder = typeOrder)

/**
 * @return the type from the expr list which is the highest order in the typeOrder list
 */
internal fun encapsulatingType(
  exprList: List<SqlExpr>,
  nullableIfAny: Boolean,
  vararg typeOrder: SqliteType
): IntermediateType {
  val types = exprList.map { it.type() }
  val sqlTypes = types.map { it.dialectType }

  val type = typeOrder.last { it in sqlTypes }
  if (!nullableIfAny && types.all { it.javaType.isNullable } ||
    nullableIfAny && types.any { it.javaType.isNullable }
  ) {
    return IntermediateType(type).asNullable()
  }
  return IntermediateType(type)
}

package com.squareup.sqldelight.core.lang.util

import com.alecstrong.sql.psi.core.psi.LazyQuery
import com.alecstrong.sql.psi.core.psi.SqlColumnDef
import com.alecstrong.sql.psi.core.psi.SqlColumnName
import com.alecstrong.sql.psi.core.psi.SqlInsertStmt

/**
 * The list of columns that values are being provided for.
 */
internal val SqlInsertStmt.columns: List<SqlColumnDef>
  get() {
    val columns = table.query.columns
      .mapNotNull { (it.element as? SqlColumnName)?.parent as? SqlColumnDef }
    if (columnNameList.isEmpty()) return columns

    val columnMap = linkedMapOf(*columns.map { it.columnName.name to it }.toTypedArray())
    return columnNameList.mapNotNull { columnMap[it.name] }
  }

internal val SqlInsertStmt.table: LazyQuery
  get() = tablesAvailable(this).first { it.tableName.name == tableName.name }

package com.squareup.sqldelight.core.lang.util

import com.alecstrong.sql.psi.core.psi.SqlCompoundSelectStmt
import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.alecstrong.sql.psi.core.psi.SqlCreateViewStmt
import com.alecstrong.sql.psi.core.psi.SqlCreateVirtualTableStmt
import com.alecstrong.sql.psi.core.psi.SqlCteTableName
import com.alecstrong.sql.psi.core.psi.SqlTableAlias
import com.alecstrong.sql.psi.core.psi.SqlTableName
import com.alecstrong.sql.psi.core.psi.SqlViewName
import com.alecstrong.sql.psi.core.psi.SqlWithClause
import com.intellij.psi.PsiElement

internal fun SqlCompoundSelectStmt.tablesObserved() = findChildrenOfType<SqlTableName>()
  .mapNotNull { it.reference?.resolve() }
  .distinct()
  .flatMap { it.referencedTables(this) }
  .distinct()

internal fun PsiElement.referencedTables(
  compoundSelectStmt: SqlCompoundSelectStmt? = null
): List<SqlTableName> = when (this) {
  is SqlCompoundSelectStmt -> tablesObserved()
  is SqlTableAlias -> source().referencedTables()
  is SqlTableName, is SqlViewName -> {
    when (val parentRule = parent!!) {
      is SqlCreateTableStmt -> listOf(parentRule.tableName)
      is SqlCreateVirtualTableStmt -> listOf(parentRule.tableName)
      is SqlCreateViewStmt -> parentRule.compoundSelectStmt?.tablesObserved().orEmpty()
      is SqlCteTableName -> {
        val withClause = parentRule.parent as SqlWithClause
        val index = withClause.cteTableNameList.indexOf(parentRule)
        val withSelect = withClause.withClauseAuxiliaryStmtList[index]
        if (withSelect.compoundSelectStmt == compoundSelectStmt) {
          // Recursive subquery. We've already resolved the other tables in this recursive query
          // so quit out.
          emptyList()
        } else {
          withClause.withClauseAuxiliaryStmtList[index].compoundSelectStmt.tablesObserved()
        }
      }
      else -> reference!!.resolve()!!.referencedTables()
    }
  }
  else -> throw IllegalStateException("Cannot get reference table for psi type ${this.javaClass}")
}

/*
 * Copyright (C) 2017 Square, Inc.
 *
 * 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.squareup.sqldelight.core.lang.util

import com.alecstrong.sql.psi.core.psi.AliasElement
import com.alecstrong.sql.psi.core.psi.SqlColumnName
import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
import com.alecstrong.sql.psi.core.psi.SqlCreateViewStmt
import com.alecstrong.sql.psi.core.psi.SqlCreateVirtualTableStmt
import com.alecstrong.sql.psi.core.psi.SqlExpr
import com.alecstrong.sql.psi.core.psi.SqlModuleArgument
import com.alecstrong.sql.psi.core.psi.SqlTableName
import com.alecstrong.sql.psi.core.psi.SqlTypes
import com.alecstrong.sql.psi.core.psi.mixins.ColumnDefMixin
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiWhiteSpace
import com.intellij.psi.tree.IElementType
import com.intellij.psi.tree.TokenSet
import com.intellij.psi.util.PsiTreeUtil
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.INTEGER
import com.squareup.sqldelight.core.dialect.sqlite.SqliteType.TEXT
import com.squareup.sqldelight.core.lang.IntermediateType
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.SqlDelightQueriesFile
import com.squareup.sqldelight.core.lang.acceptsTableInterface
import com.squareup.sqldelight.core.lang.psi.ColumnTypeMixin
import com.squareup.sqldelight.core.lang.psi.InsertStmtValuesMixin

internal inline fun <reified R : PsiElement> PsiElement.parentOfType(): R {
  return PsiTreeUtil.getParentOfType(this, R::class.java)!!
}

internal fun PsiElement.type(): IntermediateType = when (this) {
  is AliasElement -> source().type().copy(name = name)
  is ColumnDefMixin -> (columnType as ColumnTypeMixin).type()
  is SqlColumnName -> {
    when (val parentRule = parent!!) {
      is ColumnDefMixin -> parentRule.type()
      is SqlCreateVirtualTableStmt -> IntermediateType(TEXT, name = this.name)
      else -> {
        when (val resolvedReference = reference!!.resolve()!!) {
          // Synthesized columns refer directly to the table
          is SqlCreateTableStmt,
          is SqlCreateVirtualTableStmt -> synthesizedColumnType(this.name)
          else -> {
            val columnSelected = queryAvailable(this).flatMap { it.columns }
              .firstOrNull { it.element == resolvedReference }
            columnSelected?.nullable?.let {
              resolvedReference.type().nullableIf(it)
            } ?: resolvedReference.type()
          }
        }
      }
    }
  }
  is SqlExpr -> type()
  else -> throw IllegalStateException("Cannot get function type for psi type ${this.javaClass}")
}

private fun synthesizedColumnType(columnName: String): IntermediateType {
  val sqliteType = when (columnName) {
    "docid", "rowid", "oid", "_rowid_" -> INTEGER
    else -> TEXT
  }

  return IntermediateType(sqliteType, name = columnName)
}

internal fun PsiElement.sqFile(): SqlDelightFile = containingFile as SqlDelightFile

inline fun <reified T : PsiElement> PsiElement.findChildrenOfType(): Collection<T> {
  return PsiTreeUtil.findChildrenOfType(this, T::class.java)
}

inline fun <reified T : PsiElement> PsiElement.findChildOfType(): T? {
  return PsiTreeUtil.findChildOfType(this, T::class.java)
}

fun PsiElement.childOfType(type: IElementType): PsiElement? {
  return node.findChildByType(type)?.psi
}

fun PsiElement.childOfType(types: TokenSet): PsiElement? {
  return node.findChildByType(types)?.psi
}

inline fun <reified T : PsiElement> PsiElement.nextSiblingOfType(): T {
  return PsiTreeUtil.getNextSiblingOfType(this, T::class.java)!!
}

private fun PsiElement.rangesToReplace(): List<Pair<IntRange, String>> {
  return if (this is ColumnTypeMixin && javaTypeName != null) {
    listOf(
      Pair(
        first = (typeName.node.startOffset + typeName.node.textLength) until
          (javaTypeName!!.node.startOffset + javaTypeName!!.node.textLength),
        second = ""
      )
    )
  } else if (this is SqlModuleArgument && moduleArgumentDef?.columnDef != null && (parent as SqlCreateVirtualTableStmt).moduleName?.text?.toLowerCase() == "fts5") {
    val columnDef = moduleArgumentDef!!.columnDef!!
    // If there is a space at the end of the constraints, preserve it.
    val lengthModifier = if (columnDef.columnConstraintList.isNotEmpty() && columnDef.columnConstraintList.last()?.lastChild?.prevSibling is PsiWhiteSpace) 1 else 0
    listOf(
      Pair(
        first = (columnDef.columnName.node.startOffset + columnDef.columnName.node.textLength) until
          (columnDef.columnName.node.startOffset + columnDef.node.textLength - lengthModifier),
        second = ""
      )
    )
  } else if (this is InsertStmtValuesMixin && parent?.acceptsTableInterface() == true) {
    listOf(
      Pair(
        first = childOfType(SqlTypes.BIND_EXPR)!!.range,
        second = parent!!.columns.joinToString(separator = ", ", prefix = "(", postfix = ")") { "?" }
      )
    )
  } else {
    children.flatMap { it.rangesToReplace() }
  }
}

private operator fun IntRange.minus(amount: Int): IntRange {
  return IntRange(start - amount, endInclusive - amount)
}

private val IntRange.length: Int
  get() = endInclusive - start + 1

fun PsiElement.rawSqlText(
  replacements: List<Pair<IntRange, String>> = emptyList()
): String {
  return (replacements + rangesToReplace())
    .sortedBy { it.first.first }
    .map { (range, replacement) -> (range - node.startOffset) to replacement }
    .fold(
      0 to text,
      { (totalRemoved, sqlText), (range, replacement) ->
        (totalRemoved + (range.length - replacement.length)) to sqlText.replaceRange(range - totalRemoved, replacement)
      }
    ).second
}

internal val PsiElement.range: IntRange
  get() = node.startOffset until (node.startOffset + node.textLength)

fun Collection<SqlDelightQueriesFile>.forInitializationStatements(
  body: (sqlText: String) -> Unit
) {
  val views = ArrayList<SqlCreateViewStmt>()
  val creators = ArrayList<PsiElement>()

  forEach { file ->
    file.sqliteStatements()
      .filter { (label, _) -> label.name == null }
      .forEach { (_, sqliteStatement) ->
        when {
          sqliteStatement.createViewStmt != null -> views.add(sqliteStatement.createViewStmt!!)
          sqliteStatement.createTriggerStmt != null -> creators.add(sqliteStatement.createTriggerStmt!!)
          sqliteStatement.createIndexStmt != null -> creators.add(sqliteStatement.createIndexStmt!!)
          else -> body(sqliteStatement.rawSqlText())
        }
      }
  }

  val viewsLeft = views.map { it.viewName.name }.toMutableSet()
  while (views.isNotEmpty()) {
    views.removeAll { view ->
      if (view.compoundSelectStmt!!.findChildrenOfType<SqlTableName>()
        .any { it.name in viewsLeft }
      ) {
        return@removeAll false
      }
      body(view.rawSqlText())
      viewsLeft.remove(view.viewName.name)
      return@removeAll true
    }
  }

  creators.forEach { body(it.rawSqlText()) }
}

package com.squareup.sqldelight.intellij.actions.file

import com.alecstrong.sql.psi.core.psi.SqlTableName
import com.alecstrong.sql.psi.core.psi.SqlTypes.TABLE
import com.intellij.ide.actions.CreateFileFromTemplateAction
import com.intellij.ide.actions.CreateFileFromTemplateDialog
import com.intellij.ide.fileTemplates.FileTemplate
import com.intellij.ide.fileTemplates.FileTemplateManager
import com.intellij.ide.fileTemplates.actions.AttributesDefaults
import com.intellij.ide.fileTemplates.ui.CreateFromTemplateDialog
import com.intellij.openapi.actionSystem.CommonDataKeys
import com.intellij.openapi.actionSystem.DataContext
import com.intellij.openapi.application.runWriteAction
import com.intellij.openapi.editor.LogicalPosition
import com.intellij.openapi.fileEditor.FileEditorManager
import com.intellij.openapi.project.DumbAware
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.InputValidatorEx
import com.intellij.psi.PsiDirectory
import com.intellij.psi.PsiFile
import com.intellij.psi.util.PsiTreeUtil
import com.intellij.util.IncorrectOperationException
import com.squareup.sqldelight.core.SqlDelightFileIndex
import com.squareup.sqldelight.core.lang.SqlDelightFile
import com.squareup.sqldelight.core.lang.SqlDelightFileType
import com.squareup.sqldelight.core.lang.SqlDelightLanguage
import com.squareup.sqldelight.intellij.SqlDelightFileIconProvider
import org.jetbrains.kotlin.idea.util.findModule
import java.util.Properties

/**
 * Creates a new SqlDelight file/table/migration from a template (see [fileTemplates.internal])
 *
 * Adopted from [NewKotlinFileAction](https://github.com/JetBrains/kotlin/blob/f6a739bbc5b66a8f2cb5cb61acb1e87c899fa11e/idea/src/org/jetbrains/kotlin/idea/actions/NewKotlinFileAction.kt#L56)
 */
class SqlDelightCreateFileAction :
  CreateFileFromTemplateAction(
    CAPTION,
    "Creates new SqlDelight file or table",
    SqlDelightFileType.icon
  ),
  DumbAware {

  override fun postProcess(
    createdElement: PsiFile,
    templateName: String,
    customProperties: MutableMap<String, String>?
  ) {
    super.postProcess(createdElement, templateName, customProperties)

    if (createdElement is SqlDelightFile) {
      val isTableTemplate = PsiTreeUtil.collectElements(createdElement) { it is SqlTableName }.isNotEmpty()
      val editor = FileEditorManager.getInstance(createdElement.project).selectedTextEditor ?: return
      if (editor.document == createdElement.viewProvider.document) {
        if (isTableTemplate) {
          // move caret inside "create table" sql statement
          val lineCount = editor.document.lineCount
          if (lineCount > 0) {
            editor.caretModel.moveToLogicalPosition(LogicalPosition(lineCount - 2, 0))
          }
        } else {
          // move caret after "create table", so the user can enter the table name conveniently
          val tableKeyword = PsiTreeUtil.collectElements(createdElement) { it.node.elementType == TABLE }.firstOrNull() ?: return
          editor.caretModel.moveToOffset(tableKeyword.textRange.endOffset + 1)
        }
      }
    }
  }

  override fun getActionName(directory: PsiDirectory?, newName: String, templateName: String?): String = CAPTION

  override fun buildDialog(project: Project, directory: PsiDirectory, builder: CreateFileFromTemplateDialog.Builder) {
    val icon = SqlDelightFileIconProvider.getIcon(SqlDelightLanguage, project)
    builder.setTitle("New $CAPTION")
      .addKind("File", icon, "SqlDelight File")
      .addKind("Table", icon, "SqlDelight Table")
      .addKind("Migration", icon, "SqlDelight Migration")
      .setValidator(NameValidator)
  }

  override fun isAvailable(dataContext: DataContext): Boolean {
    if (!super.isAvailable(dataContext)) return false

    val project = CommonDataKeys.PROJECT.getData(dataContext) ?: return false
    val file = CommonDataKeys.VIRTUAL_FILE.getData(dataContext) ?: return false

    val module = file.findModule(project) ?: return false

    return SqlDelightFileIndex.getInstance(module).isConfigured
  }

  override fun startInWriteAction() = false

  override fun createFileFromTemplate(name: String, template: FileTemplate, dir: PsiDirectory): PsiFile? {
    val (className, targetDir) = findOrCreateTarget(dir, name, template.extension)
    return createFromTemplate(targetDir, className, template)
  }

  companion object {
    private const val CAPTION = "SqlDelight File/Table"

    private val FQNAME_SEPARATORS = charArrayOf('/', '\\', '.')

    private object NameValidator : InputValidatorEx {
      private val numberRegex = Regex(".*\\d.*")

      override fun getErrorText(inputString: String): String? {
        if (inputString.trim().isEmpty()) {
          return "Name can't be empty"
        }

        val parts: List<String> = inputString.split(*FQNAME_SEPARATORS)
        if (parts.any { it.trim().isEmpty() }) {
          return "Name can't have empty parts"
        }

        if (parts.any { it.contains("migration") || it.contains("sqm") } &&
          parts.none { it.matches(numberRegex) }
        ) {
          return "Migration filenames must contain a number"
        }

        return null
      }

      override fun checkInput(inputString: String): Boolean = true

      override fun canClose(inputString: String): Boolean = getErrorText(inputString) == null
    }

    private fun findOrCreateTarget(dir: PsiDirectory, name: String, extension: String): Pair<String, PsiDirectory> {
      var className = name.removeSuffix(".$extension")
      var targetDir = dir

      for (splitChar in FQNAME_SEPARATORS) {
        if (splitChar in className) {
          val names = className.trim().split(splitChar)

          for (dirName in names.dropLast(1)) {
            targetDir = targetDir.findSubdirectory(dirName) ?: runWriteAction {
              targetDir.createSubdirectory(dirName)
            }
          }

          className = names.last()
          break
        }
      }
      return className to targetDir
    }

    private fun createFromTemplate(dir: PsiDirectory, className: String, template: FileTemplate): PsiFile? {
      val project = dir.project
      val defaultProperties = FileTemplateManager.getInstance(project).defaultProperties

      val properties = Properties(defaultProperties)

      val element = try {
        CreateFromTemplateDialog(
          project, dir, template,
          AttributesDefaults(className).withFixedName(true),
          properties
        ).create()
      } catch (e: IncorrectOperationException) {
        throw e
      } catch (e: Exception) {
        LOG.error(e)
        return null
      }

      return element?.containingFile
    }
  }
}

package com.squareup.sqldelight.drivers.native.util

import co.touchlab.stately.collections.SharedHashMap
import co.touchlab.stately.collections.frozenHashMap

internal actual fun <T : Any> nativeCache(): NativeCache<T> =
  NativeCacheImpl()

private class NativeCacheImpl<T : Any> : NativeCache<T> {
  private val dictionary = frozenHashMap<String, T?>() as SharedHashMap<String, T?>

  override fun put(key: String, value: T?): T? = dictionary.put(key, value)
  override fun getOrCreate(key: String, block: () -> T): T = dictionary.getOrPut(key, block)!!
  override fun remove(key: String): T? = dictionary.remove(key)
  override fun cleanUp(block: (T) -> Unit) {
    dictionary.values.filterNotNull().forEach(block)
  }
}

package com.squareup.sqldelight.drivers.native.util

import co.touchlab.stately.concurrency.AtomicBoolean
import kotlinx.cinterop.alloc
import kotlinx.cinterop.free
import kotlinx.cinterop.nativeHeap
import kotlinx.cinterop.ptr
import platform.posix.pthread_cond_destroy
import platform.posix.pthread_cond_init
import platform.posix.pthread_cond_signal
import platform.posix.pthread_cond_tVar
import platform.posix.pthread_cond_wait
import platform.posix.pthread_mutex_destroy
import platform.posix.pthread_mutex_init
import platform.posix.pthread_mutex_lock
import platform.posix.pthread_mutex_tVar
import platform.posix.pthread_mutex_unlock

internal actual class PoolLock actual constructor() {
  private val isActive = AtomicBoolean(true)
  private val mutex = nativeHeap.alloc<pthread_mutex_tVar>()
    .apply { pthread_mutex_init(ptr, null) }
  private val cond = nativeHeap.alloc<pthread_cond_tVar>()
    .apply { pthread_cond_init(ptr, null) }

  actual fun <R> withLock(
    action: CriticalSection.() -> R
  ): R {
    check(isActive.value)
    pthread_mutex_lock(mutex.ptr)

    val result: R

    try {
      result = action(CriticalSection())
    } finally {
      pthread_mutex_unlock(mutex.ptr)
    }

    return result
  }

  actual fun notifyConditionChanged() {
    pthread_cond_signal(cond.ptr)
  }

  actual fun close(): Boolean {
    if (isActive.compareAndSet(expected = true, new = false)) {
      pthread_cond_destroy(cond.ptr)
      pthread_mutex_destroy(mutex.ptr)
      nativeHeap.free(cond)
      nativeHeap.free(mutex)
      return true
    }

    return false
  }

  actual inner class CriticalSection {
    actual fun <R> loopForConditionalResult(block: () -> R?): R {
      check(isActive.value)

      var result = block()

      while (result == null) {
        pthread_cond_wait(cond.ptr, mutex.ptr)
        result = block()
      }

      return result
    }

    actual fun loopUntilConditionalResult(block: () -> Boolean) {
      check(isActive.value)

      while (!block()) {
        pthread_cond_wait(cond.ptr, mutex.ptr)
      }
    }
  }
}

package com.squareup.sqldelight.drivers.native.util

import co.touchlab.stately.collections.SharedHashMap
import co.touchlab.stately.collections.SharedLinkedList

/**
 * This should probably be in Stately
 */
internal fun <T> SharedLinkedList<T>.cleanUp(block: (T) -> Unit) {
  val extractList = ArrayList<T>(size)
  extractList.addAll(this)
  this.clear()
  extractList.forEach(block)
}

internal fun <K, V> SharedHashMap<K, V>.cleanUp(block: (Map.Entry<K, V>) -> Unit) {
  val extractMap = HashMap<K, V>(this.size)
  extractMap.putAll(this)
  this.clear()
  extractMap.forEach(block)
}

package com.squareup.sqldelight.drivers.native.util

import platform.Foundation.NSMutableDictionary
import platform.Foundation.allValues
import platform.Foundation.setValue
import platform.Foundation.valueForKey
import kotlin.native.concurrent.freeze

internal actual fun <T : Any> nativeCache(): NativeCache<T> =
  NativeCacheImpl()

private class NativeCacheImpl<T : Any> : NativeCache<T> {
  private val dictionary = NSMutableDictionary()
  private val lock = PoolLock()

  @Suppress("UNCHECKED_CAST")
  override fun put(key: String, value: T?): T? {
    value.freeze()
    return lock.withLock {
      val r = dictionary.valueForKey(key)
      dictionary.setValue(value, key)
      r as? T
    }
  }

  @Suppress("UNCHECKED_CAST")
  override fun getOrCreate(key: String, block: () -> T): T = lock.withLock {
    val r = dictionary.valueForKey(key) as? T
    r ?: block().let { newValue ->
      newValue.freeze()
      dictionary.setValue(newValue, key)
      newValue
    }
  }

  @Suppress("UNCHECKED_CAST")
  override fun remove(key: String): T? = lock.withLock {
    val r = dictionary.valueForKey(key)
    dictionary.removeObjectForKey(key)
    r as? T
  }

  @Suppress("UNCHECKED_CAST")
  override fun cleanUp(block: (T) -> Unit) = lock.withLock {
    dictionary.allValues.forEach { entry ->
      entry?.let { block(entry as T) }
    }
  }
}

package com.squareup.sqldelight.drivers.native.util

import co.touchlab.stately.concurrency.AtomicBoolean
import kotlinx.cinterop.alloc
import kotlinx.cinterop.free
import kotlinx.cinterop.nativeHeap
import kotlinx.cinterop.ptr
import platform.posix.pthread_cond_destroy
import platform.posix.pthread_cond_init
import platform.posix.pthread_cond_signal
import platform.posix.pthread_cond_t
import platform.posix.pthread_cond_wait
import platform.posix.pthread_mutex_destroy
import platform.posix.pthread_mutex_init
import platform.posix.pthread_mutex_lock
import platform.posix.pthread_mutex_t
import platform.posix.pthread_mutex_unlock

internal actual class PoolLock actual constructor() {
  private val isActive = AtomicBoolean(true)
  private val mutex = nativeHeap.alloc<pthread_mutex_t>()
    .apply { pthread_mutex_init(ptr, null) }
  private val cond = nativeHeap.alloc<pthread_cond_t>()
    .apply { pthread_cond_init(ptr, null) }

  actual fun <R> withLock(
    action: CriticalSection.() -> R
  ): R {
    check(isActive.value)
    pthread_mutex_lock(mutex.ptr)

    val result: R

    try {
      result = action(CriticalSection())
    } finally {
      pthread_mutex_unlock(mutex.ptr)
    }

    return result
  }

  actual fun notifyConditionChanged() {
    pthread_cond_signal(cond.ptr)
  }

  actual fun close(): Boolean {
    if (isActive.compareAndSet(expected = true, new = false)) {
      pthread_cond_destroy(cond.ptr)
      pthread_mutex_destroy(mutex.ptr)
      nativeHeap.free(cond)
      nativeHeap.free(mutex)
      return true
    }

    return false
  }

  actual inner class CriticalSection {
    actual fun <R> loopForConditionalResult(block: () -> R?): R {
      check(isActive.value)

      var result = block()

      while (result == null) {
        pthread_cond_wait(cond.ptr, mutex.ptr)
        result = block()
      }

      return result
    }

    actual fun loopUntilConditionalResult(block: () -> Boolean) {
      check(isActive.value)

      while (!block()) {
        pthread_cond_wait(cond.ptr, mutex.ptr)
      }
    }
  }
}

package com.squareup.sqldelight.drivers.native.util

internal interface NativeCache<T : Any> {
  fun put(key: String, value: T?): T?
  fun getOrCreate(key: String, block: () -> T): T
  fun remove(key: String): T?
  fun cleanUp(block: (T) -> Unit)
}

internal expect fun <T : Any> nativeCache(): NativeCache<T>

package com.squareup.sqldelight.drivers.native.util

@Suppress("NO_ACTUAL_FOR_EXPECT")
internal expect class PoolLock() {
  fun <R> withLock(
    action: CriticalSection.() -> R
  ): R

  /**
   * Select one blocked thread in [CriticalSection.loopForConditionalResult] to be woken up for
   * re-evaluation, if any.
   */
  fun notifyConditionChanged()

  fun close(): Boolean

  inner class CriticalSection {
    /**
     * Evaluate the given lambda of a conditional result in an infinite loop, until the result is
     * available.
     *
     * If null is produced, the current thread enters suspension, and are only woken up for
     * re-evaluation by a subsequent [PoolLock.notifyConditionChanged] call. Note that the lock
     * would not be held by the current thread during its suspension. This allows resources
     * protected by the same lock to remain accessible by other threads, provided that they do not
     * depend on the same conditional result.
     */
    fun <R> loopForConditionalResult(block: () -> R?): R

    fun loopUntilConditionalResult(block: () -> Boolean)
  }
}