//===--- DiagnosticsBridge.swift ------------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2022-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 ASTBridging
import BasicBridging
import SwiftDiagnostics
import SwiftSyntax

fileprivate func emitDiagnosticParts(
  diagnosticEngine: BridgedDiagnosticEngine,
  sourceFileBuffer: UnsafeBufferPointer<UInt8>,
  message: String,
  severity: DiagnosticSeverity,
  position: AbsolutePosition,
  offset: Int,
  highlights: [Syntax] = [],
  edits: [SourceEdit] = []
) {
  // Map severity
  let bridgedSeverity = severity.bridged

  func bridgedSourceLoc(at position: AbsolutePosition) -> BridgedSourceLoc {
    return BridgedSourceLoc(at: position.advanced(by: offset), in: sourceFileBuffer)
  }

  // Emit the diagnostic
  var mutableMessage = message
  let diag = mutableMessage.withBridgedString { bridgedMessage in
    BridgedDiagnostic(
      at: bridgedSourceLoc(at: position),
      message: bridgedMessage,
      severity: bridgedSeverity,
      engine: diagnosticEngine
    )
  }

  // Emit highlights
  for highlight in highlights {
    diag.highlight(
      start: bridgedSourceLoc(at: highlight.positionAfterSkippingLeadingTrivia),
      end: bridgedSourceLoc(at: highlight.endPositionBeforeTrailingTrivia)
    )
  }

  // Emit changes for a Fix-It.
  for edit in edits {
    var newText: String = edit.replacement
    newText.withBridgedString { bridgedMessage in
      diag.fixItReplace(
        start: bridgedSourceLoc(at: edit.range.lowerBound),
        end: bridgedSourceLoc(at: edit.range.upperBound),
        replacement: bridgedMessage
      )
    }
  }

  diag.finish();
}

/// Emit the given diagnostic via the diagnostic engine.
public func emitDiagnostic(
  diagnosticEngine: BridgedDiagnosticEngine,
  sourceFileBuffer: UnsafeBufferPointer<UInt8>,
  sourceFileBufferOffset: Int = 0,
  diagnostic: Diagnostic,
  diagnosticSeverity: DiagnosticSeverity,
  messageSuffix: String? = nil
) {
  // Emit the main diagnostic
  emitDiagnosticParts(
    diagnosticEngine: diagnosticEngine,
    sourceFileBuffer: sourceFileBuffer,
    message: diagnostic.diagMessage.message + (messageSuffix ?? ""),
    severity: diagnosticSeverity,
    position: diagnostic.position,
    offset: sourceFileBufferOffset,
    highlights: diagnostic.highlights
  )

  // Emit Fix-Its.
  // FIXME: Ths assumes the fixIt is on the same tree/buffer, which is not guaranteed.
  for fixIt in diagnostic.fixIts {
    emitDiagnosticParts(
      diagnosticEngine: diagnosticEngine,
      sourceFileBuffer: sourceFileBuffer,
      message: fixIt.message.message,
      severity: .note,
      position: diagnostic.position,
      offset: sourceFileBufferOffset,
      edits: fixIt.edits
    )
  }

  // Emit any notes as follow-ons.
  // FIXME: Ths assumes the node is on the same tree/buffer, which is not guaranteed.
  for note in diagnostic.notes {
    emitDiagnosticParts(
      diagnosticEngine: diagnosticEngine,
      sourceFileBuffer: sourceFileBuffer,
      message: note.message,
      severity: .note,
      position: note.position,
      offset: sourceFileBufferOffset
    )
  }
}

extension DiagnosticSeverity {
  public var bridged: BridgedDiagnosticSeverity {
    switch self {
    case .error: return .error
    case .note: return .note
    case .warning: return .warning
    case .remark: return .remark
    }
  }
}

struct QueuedDiagnostics {
  var grouped: GroupedDiagnostics = GroupedDiagnostics()

  /// The source file IDs we allocated, mapped from the buffer IDs used
  /// by the C++ source manager.
  var sourceFileIDs: [Int: UnsafeMutablePointer<GroupedDiagnostics.SourceFileID>] = [:]

  /// The known source files
  var sourceFiles: [ExportedSourceFile] = []
}

/// Create a grouped diagnostics structure in which we can add osou
@_cdecl("swift_ASTGen_createQueuedDiagnostics")
public func createQueuedDiagnostics() -> UnsafeRawPointer {
  let ptr = UnsafeMutablePointer<QueuedDiagnostics>.allocate(capacity: 1)
  ptr.initialize(to: .init())
  return UnsafeRawPointer(ptr)
}

/// Destroy the queued diagnostics.
@_cdecl("swift_ASTGen_destroyQueuedDiagnostics")
public func destroyQueuedDiagnostics(
  queuedDiagnosticsPtr: UnsafeMutableRawPointer
) {
  let queuedDiagnostics = queuedDiagnosticsPtr.assumingMemoryBound(to: QueuedDiagnostics.self)
  for (_, sourceFileID) in queuedDiagnostics.pointee.sourceFileIDs {
    sourceFileID.deinitialize(count: 1)
    sourceFileID.deallocate()
  }

  queuedDiagnostics.deinitialize(count: 1)
  queuedDiagnostics.deallocate()
}

/// Diagnostic message used for thrown errors.
fileprivate struct SimpleDiagnostic: DiagnosticMessage {
  let message: String

  let severity: DiagnosticSeverity

  var diagnosticID: MessageID {
    .init(domain: "SwiftCompiler", id: "SimpleDiagnostic")
  }
}

extension BridgedDiagnosticSeverity {
  var asSeverity: DiagnosticSeverity {
    switch self {
    case .fatalError: return .error
    case .error: return .error
    case .warning: return .warning
    case .remark: return .remark
    case .note: return .note
    @unknown default: return .error
    }
  }
}

/// Register a source file wih the queued diagnostics.
@_cdecl("swift_ASTGen_addQueuedSourceFile")
public func addQueuedSourceFile(
  queuedDiagnosticsPtr: UnsafeMutableRawPointer,
  bufferID: Int,
  sourceFilePtr: UnsafeRawPointer,
  displayNamePtr: UnsafePointer<UInt8>,
  displayNameLength: Int,
  parentID: Int,
  positionInParent: Int
) {
  let queuedDiagnostics = queuedDiagnosticsPtr.assumingMemoryBound(to: QueuedDiagnostics.self)
  // Determine the parent link, for a child buffer.
  let parent: (GroupedDiagnostics.SourceFileID, AbsolutePosition)?
  if parentID >= 0,
    let parentSourceFileID = queuedDiagnostics.pointee.sourceFileIDs[parentID]
  {
    parent = (parentSourceFileID.pointee, AbsolutePosition(utf8Offset: positionInParent))
  } else {
    parent = nil
  }

  let displayName = String(
    decoding: UnsafeBufferPointer(
      start: displayNamePtr,
      count: displayNameLength
    ),
    as: UTF8.self
  )

  // Add the source file.
  let sourceFile = sourceFilePtr.assumingMemoryBound(to: ExportedSourceFile.self)
  let sourceFileID = queuedDiagnostics.pointee.grouped.addSourceFile(
    tree: sourceFile.pointee.syntax,
    sourceLocationConverter: sourceFile.pointee.sourceLocationConverter,
    displayName: displayName,
    parent: parent
  )
  queuedDiagnostics.pointee.sourceFiles.append(sourceFile.pointee)

  // Record the buffer ID.
  let allocatedSourceFileID = UnsafeMutablePointer<GroupedDiagnostics.SourceFileID>.allocate(capacity: 1)
  allocatedSourceFileID.initialize(to: sourceFileID)
  queuedDiagnostics.pointee.sourceFileIDs[bufferID] = allocatedSourceFileID
}

/// Add a new diagnostic to the queue.
@_cdecl("swift_ASTGen_addQueuedDiagnostic")
public func addQueuedDiagnostic(
  queuedDiagnosticsPtr: UnsafeMutableRawPointer,
  text: UnsafePointer<UInt8>,
  textLength: Int,
  severity: BridgedDiagnosticSeverity,
  cLoc: BridgedSourceLoc,
  highlightRangesPtr: UnsafePointer<BridgedSourceLoc>?,
  numHighlightRanges: Int
) {
  let queuedDiagnostics = queuedDiagnosticsPtr.assumingMemoryBound(
    to: QueuedDiagnostics.self
  )

  guard let rawPosition = cLoc.getOpaquePointerValue() else {
    return
  }

  // Find the source file that contains this location.
  let sourceFile = queuedDiagnostics.pointee.sourceFiles.first { sf in
    guard let baseAddress = sf.buffer.baseAddress else {
      return false
    }

    return rawPosition >= baseAddress && rawPosition <= baseAddress + sf.buffer.count
  }
  guard let sourceFile = sourceFile else {
    // FIXME: Hard to report an error here...
    return
  }

  let sourceFileBaseAddress = UnsafeRawPointer(sourceFile.buffer.baseAddress!)
  let sourceFileEndAddress = sourceFileBaseAddress + sourceFile.buffer.count
  let offset = rawPosition - sourceFileBaseAddress
  let position = AbsolutePosition(utf8Offset: offset)

  // Find the token at that offset.
  let node: Syntax
  if let token = sourceFile.syntax.token(at: position) {
    node = Syntax(token)
  } else if position == sourceFile.syntax.endPosition {
    // FIXME: EOF token is not included in '.token(at: position)'
    // We might want to include it, but want to avoid special handling.
    // Also 'sourceFile.syntax' is not guaranteed to be 'SourceFileSyntax'.
    if let token = sourceFile.syntax.lastToken(viewMode: .all) {
      node = Syntax(token)
    } else {
      node = sourceFile.syntax
    }
  } else {
    // position out of range.
    return
  }

  // Map the highlights.
  var highlights: [Syntax] = []
  let highlightRanges = UnsafeBufferPointer<BridgedSourceLoc>(
    start: highlightRangesPtr,
    count: numHighlightRanges * 2
  )
  for index in 0..<numHighlightRanges {
    // Make sure both the start and the end land within this source file.
    guard let start = highlightRanges[index * 2].getOpaquePointerValue(),
      let end = highlightRanges[index * 2 + 1].getOpaquePointerValue()
    else {
      continue
    }

    guard start >= sourceFileBaseAddress && start < sourceFileEndAddress,
      end >= sourceFileBaseAddress && end <= sourceFileEndAddress
    else {
      continue
    }

    // Find start tokens in the source file.
    let startPos = AbsolutePosition(utf8Offset: start - sourceFileBaseAddress)
    guard let startToken = sourceFile.syntax.token(at: startPos) else {
      continue
    }

    // Walk up from the start token until we find a syntax node that matches
    // the highlight range.
    let endPos = AbsolutePosition(utf8Offset: end - sourceFileBaseAddress)
    var highlightSyntax = Syntax(startToken)
    while true {
      // If this syntax matches our starting/ending positions, add the
      // highlight and we're done.
      if highlightSyntax.positionAfterSkippingLeadingTrivia == startPos
        && highlightSyntax.endPositionBeforeTrailingTrivia == endPos
      {
        highlights.append(highlightSyntax)
        break
      }

      // Go up to the parent.
      guard let parent = highlightSyntax.parent else {
        break
      }

      highlightSyntax = parent
    }
  }

  let textBuffer = UnsafeBufferPointer(start: text, count: textLength)
  let diagnostic = Diagnostic(
    node: node,
    position: position,
    message: SimpleDiagnostic(
      message: String(decoding: textBuffer, as: UTF8.self),
      severity: severity.asSeverity
    ),
    highlights: highlights
  )

  queuedDiagnostics.pointee.grouped.addDiagnostic(diagnostic)
}

/// Render the queued diagnostics into a UTF-8 string.
@_cdecl("swift_ASTGen_renderQueuedDiagnostics")
public func renderQueuedDiagnostics(
  queuedDiagnosticsPtr: UnsafeMutableRawPointer,
  contextSize: Int,
  colorize: Int,
  renderedStringOutPtr: UnsafeMutablePointer<BridgedStringRef>
) {
  let queuedDiagnostics = queuedDiagnosticsPtr.assumingMemoryBound(to: QueuedDiagnostics.self)
  let formatter = DiagnosticsFormatter(contextSize: contextSize, colorize: colorize != 0)
  let renderedStr = formatter.annotateSources(in: queuedDiagnostics.pointee.grouped)

  renderedStringOutPtr.pointee = allocateBridgedString(renderedStr)
}
