//===----------------------------------------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 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
//
//===----------------------------------------------------------------------===//

#if compiler(>=6)
internal import SwiftBasicFormat
internal import SwiftDiagnostics
@_spi(ExperimentalLanguageFeatures) internal import SwiftIfConfig
internal import SwiftOperators
@_spi(ExperimentalLanguageFeatures) internal import SwiftParser
internal import SwiftSyntax
@_spi(MacroExpansion) @_spi(ExperimentalLanguageFeatures) internal import SwiftSyntaxMacroExpansion
@_spi(ExperimentalLanguageFeatures) internal import SwiftSyntaxMacros
#else
import SwiftBasicFormat
import SwiftDiagnostics
@_spi(ExperimentalLanguageFeatures) import SwiftIfConfig
import SwiftOperators
@_spi(ExperimentalLanguageFeatures) import SwiftParser
import SwiftSyntax
@_spi(MacroExpansion) @_spi(ExperimentalLanguageFeatures) import SwiftSyntaxMacroExpansion
@_spi(ExperimentalLanguageFeatures) import SwiftSyntaxMacros
#endif

extension PluginProviderMessageHandler {
  /// Get concrete macro type from a pair of module name and type name.
  private func resolveMacro(_ ref: PluginMessage.MacroReference) throws -> Macro.Type {
    try provider.resolveMacro(moduleName: ref.moduleName, typeName: ref.typeName)
  }

  /// Resolve the lexical context
  private static func resolveLexicalContext(
    _ lexicalContext: [PluginMessage.Syntax]?,
    sourceManager: SourceManager,
    swiftVersion: Parser.SwiftVersion?,
    languageFeatures: Parser.LanguageFeatures?,
    operatorTable: OperatorTable,
    fallbackSyntax: some SyntaxProtocol
  ) -> [Syntax] {
    // If we weren't provided with a lexical context, retrieve it from the
    // syntax node we were given. This is for dealing with older compilers.
    guard let lexicalContext else {
      return fallbackSyntax.allMacroLexicalContexts()
    }

    return lexicalContext.map {
      sourceManager.add(
        $0,
        swiftVersion: swiftVersion,
        languageFeatures: languageFeatures,
        foldingWith: operatorTable
      )
    }
  }

  /// Expand `@freestainding(XXX)` macros.
  func expandFreestandingMacro(
    macro: PluginMessage.MacroReference,
    macroRole pluginMacroRole: PluginMessage.MacroRole?,
    discriminator: String,
    staticBuildConfiguration: StaticBuildConfiguration?,
    expandingSyntax: PluginMessage.Syntax,
    lexicalContext: [PluginMessage.Syntax]?
  ) -> PluginToHostMessage {
    let sourceManager = SourceManager(syntaxRegistry: syntaxRegistry)
    let swiftVersion = staticBuildConfiguration?.parserSwiftVersion
    let languageFeatures = staticBuildConfiguration?.parserLanguageFeatures
    let syntax = sourceManager.add(
      expandingSyntax,
      swiftVersion: swiftVersion,
      languageFeatures: languageFeatures,
      foldingWith: .standardOperators
    )

    let context = PluginMacroExpansionContext(
      sourceManager: sourceManager,
      lexicalContext: Self.resolveLexicalContext(
        lexicalContext,
        sourceManager: sourceManager,
        swiftVersion: swiftVersion,
        languageFeatures: languageFeatures,
        operatorTable: .standardOperators,
        fallbackSyntax: syntax
      ),
      expansionDiscriminator: discriminator,
      staticBuildConfiguration: staticBuildConfiguration
    )

    let expandedSource: String?
    do {
      guard let macroSyntax = syntax.asProtocol(FreestandingMacroExpansionSyntax.self) else {
        throw MacroExpansionError.freestandingMacroSyntaxIsNotMacro
      }
      let macroDefinition = try resolveMacro(macro)
      let macroRole: MacroRole
      if let pluginMacroRole {
        macroRole = MacroRole(messageMacroRole: pluginMacroRole)
      } else {
        macroRole = try inferFreestandingMacroRole(definition: macroDefinition)
      }

      expandedSource = SwiftSyntaxMacroExpansion.expandFreestandingMacro(
        definition: macroDefinition,
        macroRole: macroRole,
        node: macroSyntax,
        in: context
      )
    } catch {
      context.addDiagnostics(from: error, node: syntax)
      expandedSource = nil
    }

    let diagnostics = context.diagnostics.map {
      PluginMessage.Diagnostic(from: $0, in: sourceManager)
    }

    let response: PluginToHostMessage
    if hostCapability.hasExpandMacroResult {
      response = .expandMacroResult(expandedSource: expandedSource, diagnostics: diagnostics)
    } else {
      // TODO: Remove this  when all compilers have 'hasExpandMacroResult'.
      response = .expandFreestandingMacroResult(expandedSource: expandedSource, diagnostics: diagnostics)
    }
    return response
  }

  /// Expand `@attached(XXX)` macros.
  func expandAttachedMacro(
    macro: PluginMessage.MacroReference,
    macroRole: PluginMessage.MacroRole,
    discriminator: String,
    staticBuildConfiguration: StaticBuildConfiguration?,
    attributeSyntax: PluginMessage.Syntax,
    declSyntax: PluginMessage.Syntax,
    parentDeclSyntax: PluginMessage.Syntax?,
    extendedTypeSyntax: PluginMessage.Syntax?,
    conformanceListSyntax: PluginMessage.Syntax?,
    lexicalContext: [PluginMessage.Syntax]?
  ) -> PluginToHostMessage {
    let sourceManager = SourceManager(syntaxRegistry: syntaxRegistry)
    let swiftVersion = staticBuildConfiguration?.parserSwiftVersion
    let languageFeatures = staticBuildConfiguration?.parserLanguageFeatures

    func addToSourceManager(_ syntax: PluginMessage.Syntax, foldOperators: Bool) -> Syntax {
      sourceManager.add(
        syntax,
        swiftVersion: swiftVersion,
        languageFeatures: languageFeatures,
        foldingWith: foldOperators ? .standardOperators : nil
      )
    }

    let attributeNode = addToSourceManager(attributeSyntax, foldOperators: true)
      .cast(AttributeSyntax.self)
    let declarationNode = addToSourceManager(declSyntax, foldOperators: false)
    let parentDeclNode = parentDeclSyntax.map {
      addToSourceManager($0, foldOperators: false).cast(DeclSyntax.self)
    }
    let extendedType = extendedTypeSyntax.map {
      addToSourceManager($0, foldOperators: false).cast(TypeSyntax.self)
    }
    let conformanceList = conformanceListSyntax.map {
      let placeholderStruct = addToSourceManager($0, foldOperators: false).cast(StructDeclSyntax.self)
      return placeholderStruct.inheritanceClause!.inheritedTypes
    }

    let context = PluginMacroExpansionContext(
      sourceManager: sourceManager,
      lexicalContext: Self.resolveLexicalContext(
        lexicalContext,
        sourceManager: sourceManager,
        swiftVersion: swiftVersion,
        languageFeatures: languageFeatures,
        operatorTable: .standardOperators,
        fallbackSyntax: declarationNode
      ),
      expansionDiscriminator: discriminator,
      staticBuildConfiguration: staticBuildConfiguration
    )

    // TODO: Make this a 'String?' and remove non-'hasExpandMacroResult' branches.
    let expandedSources: [String]?
    do {
      let macroDefinition = try resolveMacro(macro)
      let role = MacroRole(messageMacroRole: macroRole)

      let expansions = SwiftSyntaxMacroExpansion.expandAttachedMacroWithoutCollapsing(
        definition: macroDefinition,
        macroRole: role,
        attributeNode: attributeNode,
        node: declarationNode,
        parentDeclNode: parentDeclNode,
        extendedType: extendedType,
        conformanceList: conformanceList,
        in: context
      )
      if let expansions, hostCapability.hasExpandMacroResult {
        let collapseIndentationWidth: Trivia?
        switch macroDefinition.formatMode {
        case .auto: collapseIndentationWidth = nil
        case .disabled: collapseIndentationWidth = []
        #if RESILIENT_LIBRARIES
        @unknown default: fatalError()
        #endif
        }
        // Make a single element array by collapsing the results into a string.
        expandedSources = [
          SwiftSyntaxMacroExpansion.collapse(
            expansions: expansions,
            for: role,
            attachedTo: declarationNode,
            indentationWidth: collapseIndentationWidth
          )
        ]
      } else {
        expandedSources = expansions
      }
    } catch {
      context.addDiagnostics(from: error, node: attributeNode)
      expandedSources = nil
    }

    let diagnostics = context.diagnostics.map {
      PluginMessage.Diagnostic(from: $0, in: sourceManager)
    }

    let response: PluginToHostMessage
    if hostCapability.hasExpandMacroResult {
      response = .expandMacroResult(expandedSource: expandedSources?.first, diagnostics: diagnostics)
    } else {
      response = .expandAttachedMacroResult(expandedSources: expandedSources, diagnostics: diagnostics)
    }
    return response
  }
}

private extension MacroRole {
  init(messageMacroRole: PluginMessage.MacroRole) {
    switch messageMacroRole {
    case .expression: self = .expression
    case .declaration: self = .declaration
    case .accessor: self = .accessor
    case .memberAttribute: self = .memberAttribute
    case .member: self = .member
    case .peer: self = .peer
    case .conformance: self = .extension
    case .codeItem: self = .codeItem
    case .extension: self = .extension
    case .preamble: self = .preamble
    case .body: self = .body
    }
  }
}
