src/main/java/com/squareup/kotlinpoet/TypeSpec.kt

/*
 * Copyright (C) 2015 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.kotlinpoet

import com.squareup.kotlinpoet.Util.requireExactlyOneOf
import java.io.IOException
import java.io.StringWriter
import java.lang.reflect.Type
import java.util.Locale
import javax.lang.model.SourceVersion
import javax.lang.model.element.Element
import kotlin.reflect.KClass

/** A generated class, interface, or enum declaration.  */
class TypeSpec private constructor(builder: TypeSpec.Builder) {
  val kind = builder.kind
  val name = builder.name
  val anonymousTypeArguments = builder.anonymousTypeArguments
  val kdoc = builder.kdoc.build()
  val annotations: List<AnnotationSpec> = Util.immutableList(builder.annotations)
  val modifiers: Set<KModifier> = Util.immutableSet(builder.modifiers)
  val typeVariables: List<TypeVariableName> = Util.immutableList(builder.typeVariables)
  val primaryConstructor = builder.primaryConstructor
  val superclass = builder.superclass
  val superinterfaces: List<TypeName> = Util.immutableList(builder.superinterfaces)
  val enumConstants: Map<String, TypeSpec> = Util.immutableMap(builder.enumConstants)
  val propertySpecs: List<PropertySpec> = Util.immutableList(builder.propertySpecs)
  val initializerBlock = builder.initializerBlock.build()
  val funSpecs: List<FunSpec> = Util.immutableList(builder.funSpecs)
  val typeSpecs: List<TypeSpec> = Util.immutableList(builder.typeSpecs)
  val originatingElements: List<Element>

  init {
    val originatingElementsMutable = mutableListOf<Element>()
    originatingElementsMutable += builder.originatingElements
    for (typeSpec in builder.typeSpecs) {
      originatingElementsMutable += typeSpec.originatingElements
    }
    this.originatingElements = Util.immutableList(originatingElementsMutable)
  }

  fun toBuilder(): Builder {
    val builder = Builder(kind, name, anonymousTypeArguments)
    builder.kdoc.add(kdoc)
    builder.annotations += annotations
    builder.modifiers += modifiers
    builder.typeVariables += typeVariables
    builder.superclass = superclass
    builder.superinterfaces += superinterfaces
    builder.enumConstants += enumConstants
    builder.propertySpecs += propertySpecs
    builder.funSpecs += funSpecs
    builder.typeSpecs += typeSpecs
    builder.initializerBlock.add(initializerBlock)
    return builder
  }

  @Throws(IOException::class)
  internal fun emit(codeWriter: CodeWriter, enumName: String?, implicitModifiers: Set<KModifier>) {
    // Nested classes interrupt wrapped line indentation. Stash the current wrapping state and put
    // it back afterwards when this type is complete.
    val previousStatementLine = codeWriter.statementLine
    codeWriter.statementLine = -1

    try {
      if (enumName != null) {
        codeWriter.emitKdoc(kdoc)
        codeWriter.emitAnnotations(annotations, false)
        codeWriter.emit("%L", enumName)
        if (!anonymousTypeArguments!!.formatParts.isEmpty()) {
          codeWriter.emit("(")
          codeWriter.emit(anonymousTypeArguments)
          codeWriter.emit(")")
        }
        if (propertySpecs.isEmpty() && funSpecs.isEmpty() && typeSpecs.isEmpty()) {
          return  // Avoid unnecessary braces "{}".
        }
        codeWriter.emit(" {\n")
      } else if (anonymousTypeArguments != null) {
        val supertype = if (!superinterfaces.isEmpty()) superinterfaces[0] else superclass
        codeWriter.emit("new %T(", supertype)
        codeWriter.emit(anonymousTypeArguments)
        codeWriter.emit(") {\n")
      } else {
        codeWriter.emitKdoc(kdoc)
        codeWriter.emitAnnotations(annotations, false)
        codeWriter.emitJavaModifiers(modifiers, implicitModifiers)
        if (kind == Kind.ANNOTATION) {
          codeWriter.emit("%L %L", "@interface", name)
        } else {
          codeWriter.emit("%L %L", kind.name.toLowerCase(Locale.US), name)
        }
        codeWriter.emitTypeVariables(typeVariables)

        primaryConstructor?.emitParameterList(codeWriter)

        val extendsTypes: List<TypeName>
        val implementsTypes: List<TypeName>
        if (kind == Kind.INTERFACE) {
          extendsTypes = superinterfaces
          implementsTypes = emptyList()
        } else {
          extendsTypes = if (superclass == ANY) emptyList() else listOf(superclass)
          implementsTypes = superinterfaces
        }

        if (!extendsTypes.isEmpty()) {
          codeWriter.emit(" extends")
          var firstType = true
          for (type in extendsTypes) {
            if (!firstType) codeWriter.emit(",")
            codeWriter.emit(" %T", type)
            firstType = false
          }
        }

        if (!implementsTypes.isEmpty()) {
          codeWriter.emit(" implements")
          var firstType = true
          for (type in implementsTypes) {
            if (!firstType) codeWriter.emit(",")
            codeWriter.emit(" %T", type)
            firstType = false
          }
        }

        codeWriter.emit(" {\n")
      }

      codeWriter.pushType(this)
      codeWriter.indent()
      var firstMember = true
      val i = enumConstants.entries.iterator()
      while (i.hasNext()) {
        val enumConstant = i.next()
        if (!firstMember) codeWriter.emit("\n")
        enumConstant.value
            .emit(codeWriter, enumConstant.key, emptySet<KModifier>())
        firstMember = false
        if (i.hasNext()) {
          codeWriter.emit(",\n")
        } else if (!propertySpecs.isEmpty() || !funSpecs.isEmpty() || !typeSpecs.isEmpty()) {
          codeWriter.emit(";\n")
        } else {
          codeWriter.emit("\n")
        }
      }

      // Non-static properties.
      for (propertySpec in propertySpecs) {
        if (!firstMember) codeWriter.emit("\n")
        propertySpec.emit(codeWriter, kind.implicitPropertyModifiers)
        firstMember = false
      }

      if (primaryConstructor != null && !primaryConstructor.code.isEmpty()) {
        codeWriter.emit("init {\n")
        codeWriter.indent()
        codeWriter.emit(primaryConstructor.code)
        codeWriter.unindent()
        codeWriter.emit("}\n")
      }

      // Initializer block.
      if (!initializerBlock.isEmpty()) {
        if (!firstMember) codeWriter.emit("\n")
        codeWriter.emit(initializerBlock)
        firstMember = false
      }

      // Constructors.
      for (funSpec in funSpecs) {
        if (!funSpec.isConstructor) continue
        if (!firstMember) codeWriter.emit("\n")
        funSpec.emit(codeWriter, name!!, kind.implicitFunctionModifiers)
        firstMember = false
      }

      // Functions (static and non-static).
      for (funSpec in funSpecs) {
        if (funSpec.isConstructor) continue
        if (!firstMember) codeWriter.emit("\n")
        funSpec.emit(codeWriter, name, kind.implicitFunctionModifiers)
        firstMember = false
      }

      // Types.
      for (typeSpec in typeSpecs) {
        if (!firstMember) codeWriter.emit("\n")
        typeSpec.emit(codeWriter, null, setOf())
        firstMember = false
      }

      codeWriter.unindent()
      codeWriter.popType()

      codeWriter.emit("}")
      if (enumName == null && anonymousTypeArguments == null) {
        codeWriter.emit("\n") // If this type isn't also a value, include a trailing newline.
      }
    } finally {
      codeWriter.statementLine = previousStatementLine
    }
  }

  override fun equals(other: Any?): Boolean {
    if (this === other) return true
    if (other == null) return false
    if (javaClass != other.javaClass) return false
    return toString() == other.toString()
  }

  override fun hashCode(): Int {
    return toString().hashCode()
  }

  override fun toString(): String {
    val out = StringWriter()
    try {
      val codeWriter = CodeWriter(out)
      emit(codeWriter, null, emptySet<KModifier>())
      return out.toString()
    } catch (e: IOException) {
      throw AssertionError()
    }
  }

  enum class Kind(
      internal val implicitPropertyModifiers: Set<KModifier>,
      internal val implicitFunctionModifiers: Set<KModifier>) {
    CLASS(
        setOf(KModifier.PUBLIC),
        setOf(KModifier.PUBLIC)),

    INTERFACE(
        setOf(KModifier.PUBLIC),
        setOf(KModifier.PUBLIC, KModifier.ABSTRACT)),

    ENUM(
        setOf(KModifier.PUBLIC),
        setOf(KModifier.PUBLIC)),

    ANNOTATION(
        emptySet(),
        setOf(KModifier.PUBLIC, KModifier.ABSTRACT))
  }

  class Builder internal constructor(
      internal val kind: Kind,
      internal val name: String?,
      internal val anonymousTypeArguments: CodeBlock?) {
    internal val kdoc = CodeBlock.builder()
    internal val annotations = mutableListOf<AnnotationSpec>()
    internal val modifiers = mutableListOf<KModifier>()
    internal val typeVariables = mutableListOf<TypeVariableName>()
    internal var primaryConstructor : FunSpec? = null
    internal var superclass: TypeName = ANY
    internal val superinterfaces = mutableListOf<TypeName>()
    internal val enumConstants = mutableMapOf<String, TypeSpec>()
    internal val propertySpecs = mutableListOf<PropertySpec>()
    internal val initializerBlock = CodeBlock.builder()
    internal val funSpecs = mutableListOf<FunSpec>()
    internal val typeSpecs = mutableListOf<TypeSpec>()
    internal val originatingElements = mutableListOf<Element>()

    init {
      require(name == null || SourceVersion.isName(name)) { "not a valid name: $name" }
    }

    fun addKdoc(format: String, vararg args: Any): Builder {
      kdoc.add(format, *args)
      return this
    }

    fun addKdoc(block: CodeBlock): Builder {
      kdoc.add(block)
      return this
    }

    fun addAnnotations(annotationSpecs: Iterable<AnnotationSpec>): Builder {
      annotations += annotationSpecs
      return this
    }

    fun addAnnotation(annotationSpec: AnnotationSpec): Builder {
      annotations += annotationSpec
      return this
    }

    fun addAnnotation(annotation: ClassName)
        = addAnnotation(AnnotationSpec.builder(annotation).build())

    fun addAnnotation(annotation: Class<*>) = addAnnotation(ClassName.get(annotation))

    fun addAnnotation(annotation: KClass<*>) = addAnnotation(ClassName.get(annotation))

    fun addModifiers(vararg modifiers: KModifier): Builder {
      check(anonymousTypeArguments == null) { "forbidden on anonymous types." }
      this.modifiers += modifiers
      return this
    }

    fun addTypeVariables(typeVariables: Iterable<TypeVariableName>): Builder {
      check(anonymousTypeArguments == null) { "forbidden on anonymous types." }
      this.typeVariables += typeVariables
      return this
    }

    fun addTypeVariable(typeVariable: TypeVariableName): Builder {
      check(anonymousTypeArguments == null) { "forbidden on anonymous types." }
      typeVariables += typeVariable
      return this
    }

    fun primaryConstructor(primaryConstructor: FunSpec?): Builder {
      check(kind == Kind.CLASS || kind == Kind.ENUM) { "$kind can't have initializer blocks" }
      if (primaryConstructor != null) {
        require (primaryConstructor.isConstructor) {
          "expected a constructor but was ${primaryConstructor.name}"
        }
      }
      this.primaryConstructor = primaryConstructor
      return this
    }

    fun superclass(superclass: TypeName): Builder {
      check(kind == Kind.CLASS) { "only classes have super classes, not $kind" }
      check(this.superclass === ANY) { "superclass already set to ${this.superclass}" }
      this.superclass = superclass
      return this
    }

    fun superclass(superclass: Type): Builder {
      return superclass(TypeName.get(superclass))
    }

    fun addSuperinterfaces(superinterfaces: Iterable<TypeName>): Builder {
      this.superinterfaces += superinterfaces
      return this
    }

    fun addSuperinterface(superinterface: TypeName): Builder {
      superinterfaces += superinterface
      return this
    }

    fun addSuperinterface(superinterface: Type)
        = addSuperinterface(TypeName.get(superinterface))

    fun addSuperinterface(superinterface: KClass<*>)
        = addSuperinterface(TypeName.get(superinterface))

    @JvmOverloads fun addEnumConstant(
        name: String,
        typeSpec: TypeSpec = anonymousClassBuilder("").build()): Builder {
      check(kind == Kind.ENUM) { "${this.name} is not enum" }
      require(typeSpec.anonymousTypeArguments != null) {
          "enum constants must have anonymous type arguments" }
      require(SourceVersion.isName(name)) { "not a valid enum constant: $name" }
      enumConstants.put(name, typeSpec)
      return this
    }

    fun addProperties(propertySpecs: Iterable<PropertySpec>): Builder {
      this.propertySpecs += propertySpecs
      return this
    }

    fun addProperty(propertySpec: PropertySpec): Builder {
      check(kind == Kind.CLASS || kind == Kind.ENUM) {
        "cannot add property $propertySpec to $kind"
      }
      propertySpecs += propertySpec
      return this
    }

    fun addProperty(type: TypeName, name: String, vararg modifiers: KModifier)
        = addProperty(PropertySpec.builder(type, name, *modifiers).build())

    fun addProperty(type: Type, name: String, vararg modifiers: KModifier)
        = addProperty(TypeName.get(type), name, *modifiers)

    fun addProperty(type: KClass<*>, name: String, vararg modifiers: KModifier)
        = addProperty(TypeName.get(type), name, *modifiers)

    fun addInitializerBlock(block: CodeBlock): Builder {
      check(kind == Kind.CLASS || kind == Kind.ENUM) { "$kind can't have initializer blocks" }
      initializerBlock.add("{\n")
          .indent()
          .add(block)
          .unindent()
          .add("}\n")
      return this
    }

    fun addFunctions(funSpecs: Iterable<FunSpec>): Builder {
      this.funSpecs += funSpecs
      return this
    }

    fun addFun(funSpec: FunSpec): Builder {
      if (kind == Kind.INTERFACE) {
        requireExactlyOneOf(funSpec.modifiers, KModifier.ABSTRACT)
        requireExactlyOneOf(funSpec.modifiers, KModifier.PUBLIC, KModifier.PRIVATE)
      } else if (kind == Kind.ANNOTATION) {
        check(funSpec.modifiers == kind.implicitFunctionModifiers) {
            "$kind $name.${funSpec.name} requires modifiers ${kind.implicitFunctionModifiers}" }
      }
      if (kind != Kind.ANNOTATION) {
        check(funSpec.defaultValue == null) {
          "$kind $name.${funSpec.name} cannot have a default value" }
      }
      funSpecs += funSpec
      return this
    }

    fun addTypes(typeSpecs: Iterable<TypeSpec>): Builder {
      this.typeSpecs += typeSpecs
      return this
    }

    fun addType(typeSpec: TypeSpec): Builder {
      typeSpecs += typeSpec
      return this
    }

    fun addOriginatingElement(originatingElement: Element): Builder {
      originatingElements += originatingElement
      return this
    }

    fun build(): TypeSpec {
      require(kind != Kind.ENUM || !enumConstants.isEmpty()) {
          "at least one enum constant is required for $name" }

      val isAbstract = modifiers.contains(KModifier.ABSTRACT) || kind != Kind.CLASS
      for (funSpec in funSpecs) {
        require(isAbstract || !funSpec.modifiers.contains(KModifier.ABSTRACT)) {
            "non-abstract type $name cannot declare abstract function ${funSpec.name}" }
      }

      val superclassIsAny = superclass == ANY
      val interestingSupertypeCount = (if (superclassIsAny) 0 else 1) + superinterfaces.size
      require(anonymousTypeArguments == null || interestingSupertypeCount <= 1) {
          "anonymous type has too many supertypes" }

      return TypeSpec(this)
    }
  }

  companion object {
    @JvmStatic fun classBuilder(name: String) = Builder(Kind.CLASS, name, null)

    @JvmStatic fun classBuilder(className: ClassName) = classBuilder(className.simpleName())

    @JvmStatic fun interfaceBuilder(name: String) = Builder(Kind.INTERFACE, name, null)

    @JvmStatic fun interfaceBuilder(className: ClassName) = interfaceBuilder(className.simpleName())

    @JvmStatic fun enumBuilder(name: String) = Builder(Kind.ENUM, name, null)

    @JvmStatic fun enumBuilder(className: ClassName) = enumBuilder(className.simpleName())

    @JvmStatic fun anonymousClassBuilder(typeArgumentsFormat: String, vararg args: Any): Builder {
      return Builder(Kind.CLASS, null, CodeBlock.builder()
          .add(typeArgumentsFormat, *args)
          .build())
    }

    @JvmStatic fun annotationBuilder(name: String) = Builder(Kind.ANNOTATION, name, null)

    @JvmStatic fun annotationBuilder(className: ClassName)
        = annotationBuilder(className.simpleName())
  }
}