//===----------------------------------------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

import SwiftBasicFormat
import SwiftSyntax

/// A format style for files generated by CodeGeneration.
public class CodeGenerationFormat: BasicFormat {
  /// The maximum number of array/dictionary/function parameters/tuple elements
  /// that should be put on the same line.
  private let maxElementsOnSameLine: Int

  public init(maxElementsOnSameLine: Int = 3) {
    self.maxElementsOnSameLine = maxElementsOnSameLine
    super.init(indentationWidth: .spaces(2))
  }

  var indentedNewline: Trivia {
    .newline + currentIndentationLevel
  }

  public override func visit(_ node: ArrayElementListSyntax) -> ArrayElementListSyntax {
    let children = node.children(viewMode: .all)
    // Short array literals are presented on one line, list each element on a different line.
    if children.count > maxElementsOnSameLine {
      return ArrayElementListSyntax(
        formatChildrenSeparatedByNewline(children: children, elementType: ArrayElementSyntax.self)
      )
    } else {
      return super.visit(node)
    }
  }

  public override func visit(_ node: CodeBlockItemSyntax) -> CodeBlockItemSyntax {
    if node.parent?.parent?.is(SourceFileSyntax.self) == true, !shouldBeSeparatedByTwoNewlines(node: node) {
      let formatted = super.visit(node)
      return ensuringTwoLeadingNewlines(node: formatted)
    } else {
      return super.visit(node)
    }
  }

  public override func visit(_ node: DictionaryElementListSyntax) -> DictionaryElementListSyntax {
    let children = node.children(viewMode: .all)
    // Short dictionary literals are presented on one line, list each element on a different line.
    if children.count > maxElementsOnSameLine {
      return DictionaryElementListSyntax(
        formatChildrenSeparatedByNewline(children: children, elementType: DictionaryElementSyntax.self)
      )
    } else {
      return super.visit(node)
    }
  }

  public override func visit(_ node: FunctionParameterListSyntax) -> FunctionParameterListSyntax {
    let children = node.children(viewMode: .all)
    // Short function parameter literals are presented on one line, list each element on a different line.
    if children.count > maxElementsOnSameLine {
      return FunctionParameterListSyntax(
        formatChildrenSeparatedByNewline(children: children, elementType: FunctionParameterSyntax.self)
      )
    } else {
      return super.visit(node)
    }
  }

  public override func visit(_ node: MemberBlockSyntax) -> MemberBlockSyntax {
    if node.members.count == 0 {
      return node.with(\.leftBrace, .leftBraceToken())
    } else {
      return super.visit(node)
    }
  }

  public override func visit(_ node: MemberBlockItemSyntax) -> MemberBlockItemSyntax {
    let formatted = super.visit(node)
    if node != node.parent?.children(viewMode: .sourceAccurate).first?.as(MemberBlockItemSyntax.self)
      && !node.decl.is(EnumCaseDeclSyntax.self)
    {
      return ensuringTwoLeadingNewlines(node: formatted)
    } else {
      return formatted
    }
  }

  public override func visit(_ node: LabeledExprListSyntax) -> LabeledExprListSyntax {
    let children = node.children(viewMode: .all)
    // Short tuple element list literals are presented on one line, list each element on a different line.
    if children.count > maxElementsOnSameLine {
      let inMethodCallThatStartsOnNewline =
        node.parent?.as(FunctionCallExprSyntax.self)?.calledExpression.as(MemberAccessExprSyntax.self)?.period
        .startsOnNewline ?? false
      if inMethodCallThatStartsOnNewline {
        increaseIndentationLevel()
      }
      defer {
        if inMethodCallThatStartsOnNewline {
          decreaseIndentationLevel()
        }
      }

      return LabeledExprListSyntax(
        formatChildrenSeparatedByNewline(children: children, elementType: LabeledExprSyntax.self)
      )
    } else {
      return super.visit(node)
    }
  }

  public override func requiresIndent(_ node: some SyntaxProtocol) -> Bool {
    switch node.kind {
    case .arrayElementList, .dictionaryElementList, .functionParameterList, .labeledExprList:
      let indentManually = node.children(viewMode: .sourceAccurate).count > maxElementsOnSameLine
      if indentManually {
        return false
      }
      if !node.startsOnNewline {
        return false
      }
    default:
      break
    }
    return super.requiresIndent(node)
  }

  // MARK: - Private

  private func shouldBeSeparatedByTwoNewlines(node: CodeBlockItemSyntax) -> Bool {
    // First item in the ``CodeBlockItemListSyntax`` don't need a newline or indentation if the parent is a ``SourceFileSyntax``.
    // We want to group imports so newline between them should be omitted
    return node.parent?.as(CodeBlockItemListSyntax.self)?.first == node || node.item.is(ImportDeclSyntax.self)
  }

  private func ensuringTwoLeadingNewlines<NodeType: SyntaxProtocol>(node: NodeType) -> NodeType {
    if node.leadingTrivia.first?.isNewline ?? false {
      return node.with(\.leadingTrivia, .newline + node.leadingTrivia)
    } else {
      return node.with(\.leadingTrivia, .newlines(2) + node.leadingTrivia)
    }
  }

  private func formatChildrenSeparatedByNewline<SyntaxType: SyntaxProtocol>(
    children: SyntaxChildren,
    elementType: SyntaxType.Type
  ) -> [SyntaxType] {
    increaseIndentationLevel()
    var formattedChildren = children.map {
      return self.rewrite($0.cast(SyntaxType.self)).cast(SyntaxType.self)
    }
    formattedChildren = formattedChildren.map { child in
      var child = child

      if let firstNonSpaceOrTabIndex = child.trailingTrivia.firstIndex(where: { !$0.isSpaceOrTab }) {
        if child.trailingTrivia[firstNonSpaceOrTabIndex].isNewline {
          child.trailingTrivia = Trivia(pieces: child.trailingTrivia.suffix(from: firstNonSpaceOrTabIndex))
        }
      } else {
        child.trailingTrivia = Trivia()
      }

      if !child.startsOnNewline {
        child.leadingTrivia = indentedNewline + child.leadingTrivia
      }
      return child
    }
    decreaseIndentationLevel()
    if let lastChild = formattedChildren.last,
      !lastChild.trailingTrivia.contains(where: \.isNewline)
    {
      let nextTokenStartsWithNewline =
        lastChild.nextToken(viewMode: .sourceAccurate)?.leadingTrivia.first?.isNewline ?? false
      if !nextTokenStartsWithNewline {
        formattedChildren[formattedChildren.count - 1] = lastChild.with(
          \.trailingTrivia,
          lastChild.trailingTrivia + indentedNewline
        )
      }
    }
    return formattedChildren
  }
}

private extension SyntaxProtocol {
  var startsOnNewline: Bool {
    return self.leadingTrivia.contains(where: \.isNewline)
      || self.previousToken(viewMode: .sourceAccurate)?.trailingTrivia.contains(where: \.isNewline) ?? false
  }
}
