// Copyright (C) 2024 Rubén Beltrán del Río

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program. If not, see https://map.tranquil.systems.
import AppKit

extension Wmap {
  class SyntaxHighlighter {
    private var cachedSyntaxElements: [SyntaxElement] = []
    private var lineToElementsMap: [Int: [SyntaxElement]] = [:]

    func updateSyntaxElements(from parsedMap: [Entity]) {
      cachedSyntaxElements = parsedMap.flatMap { $0.syntaxElements }

      // Build line-to-elements mapping for fast lookup
      lineToElementsMap.removeAll()
      for element in cachedSyntaxElements {
        let line = element.range.start.line
        lineToElementsMap[line, default: []].append(element)
      }
    }

    func applySyntaxHighlighting(textStorage: NSTextStorage, range: NSRange) {
      // Get the line range that encompasses the edited range
      let string = textStorage.string as NSString
      let lineRange = string.lineRange(for: range)

      // Apply highlighting to elements that intersect with the line range
      for element in cachedSyntaxElements {
        // Safely convert string range to NSRange with bounds checking
        guard element.range.stringRange.lowerBound <= element.range.stringRange.upperBound,
          element.range.stringRange.upperBound <= textStorage.string.endIndex
        else {
          continue
        }

        let nsRange = NSRange(element.range.stringRange, in: textStorage.string)

        // Only process elements that intersect with our update area
        guard nsRange.location != NSNotFound,
          nsRange.location + nsRange.length <= textStorage.length,
          NSIntersectionRange(nsRange, lineRange).length > 0
        else {
          continue
        }

        let color = colorForSyntaxType(element.type)
        textStorage.addAttribute(.foregroundColor, value: color, range: nsRange)
      }
    }

    private func colorForSyntaxType(_ type: SyntaxType) -> NSColor {
      switch type {
      case .keyword:
        return NSColor.Theme.Syntax.option
      case .vertexLabel:
        return NSColor.Theme.Syntax.vertex
      case .number:
        return NSColor.Theme.Syntax.number
      case .symbol:
        return NSColor.Theme.Syntax.symbol
      case .shape:
        return NSColor.Theme.Syntax.option
      case .text:
        return NSColor.Theme.Syntax.text
      case .punctuation:
        return NSColor.Theme.Syntax.punctuation
      case .stageNumber:
        return NSColor.Theme.Syntax.option
      }
    }
  }
}