1 files changed, 48 insertions, 53 deletions

diff --git a/QuickLook/PreviewViewController.swift b/QuickLook/PreviewViewController.swift
index c93b2cb..ecacee0 100644
--- a/QuickLook/PreviewViewController.swift
+++ b/QuickLook/PreviewViewController.swift
@@ -1,13 +1,10 @@
-//
-//  PreviewViewController.swift
-//  QuickLook
-//
-//  Created by Ruben Beltran del Rio on 2025-07-11.
-//
-
 import Cocoa
 import Quartz
 import SwiftUI
+import WmapParser
+import SwiftTreeSitter
+import SwiftTreeSitterLayer
+import TreeSitterWmap
 
 class PreviewViewController: NSViewController, QLPreviewingController {
 
@@ -48,43 +45,12 @@ class PreviewViewController: NSViewController, QLPreviewingController {
           previewContent = AnyView(createPlainTextView(content: content))
 
         case .highlightedText:
-          // Parse for syntax highlighting
-          let lexer = Wmap.Lexer(content)
-          let parser = Wmap.Parser(lexer: lexer)
-          let parsedMap = parser.parse()
           previewContent = AnyView(
-            createHighlightedTextView(content: content, parsedMap: parsedMap))
+            createHighlightedTextView(content: content))
 
         case .map:
-          // Parse the content to create a ParsedMap
-          let lexer = Wmap.Lexer(content)
-          let parser = Wmap.Parser(lexer: lexer)
-          let parsedMap = parser.parse()
-
-          if parsedMap.entities.isEmpty {
-            // Fallback view for empty or invalid maps
-            previewContent = AnyView(
-              VStack {
-                Image(systemName: "map")
-                  .font(.system(size: 48))
-                  .foregroundColor(.gray)
-                Text("quick_look.empty_map.title")
-                  .font(.Theme.Title.emphasized)
-                  .foregroundColor(.Theme.UI.foreground)
-                Text("quick_look.empty_map.description")
-                  .font(.Theme.Body.regular)
-                  .foregroundColor(.Theme.UI.foreground)
-              }
-              .frame(maxWidth: .infinity, maxHeight: .infinity)
-              .foregroundColor(.Theme.UI.background)
-            )
-          } else {
-            // Limit entities for memory efficiency in QuickLook
-            let limitedEntities = Array(parsedMap.entities.prefix(100))  // Limit to first 100 entities
-            let limitedMap = Wmap.ParsedMap(
-              entities: limitedEntities, vertexLabels: parsedMap.vertexLabels)
-            previewContent = AnyView(MapPreviewView(parsedMap: limitedMap))
-          }
+          let parsedMap = WmapParser.parse(content)
+          previewContent = AnyView(MapPreviewView(parsedMap: parsedMap))
         }
 
         let hostingView = NSHostingView(rootView: previewContent)
@@ -117,8 +83,8 @@ class PreviewViewController: NSViewController, QLPreviewingController {
     .background(Color.white)
   }
 
-  private func createHighlightedTextView(content: String, parsedMap: Wmap.ParsedMap) -> some View {
-    QuickLookTextEditor(content: content, parsedMap: parsedMap)
+  private func createHighlightedTextView(content: String) -> some View {
+    QuickLookTextEditor(content: content)
       .frame(maxWidth: .infinity, maxHeight: .infinity)
       .background(Color.white)
   }
@@ -136,7 +102,7 @@ class PreviewViewController: NSViewController, QLPreviewingController {
 
 // A SwiftUI view specifically for QuickLook preview
 struct MapPreviewView: View {
-  let parsedMap: Wmap.ParsedMap
+  let parsedMap: WmapParser.Map
   @State private var selectedEvolution: StageType = .general
 
   var body: some View {
@@ -153,7 +119,7 @@ struct MapPreviewView: View {
 
       ScrollView([.horizontal, .vertical]) {
         MapRenderView(
-          entities: parsedMap.entities,
+          parsedMap: parsedMap,
           evolution: $selectedEvolution
         )
         .drawingGroup(opaque: false)
@@ -172,7 +138,6 @@ struct MapPreviewView: View {
 // A read-only text editor for QuickLook that shows syntax highlighting
 struct QuickLookTextEditor: NSViewRepresentable {
   let content: String
-  let parsedMap: Wmap.ParsedMap
 
   func makeNSView(context: Context) -> NSScrollView {
     let scrollView = NSTextView.scrollableTextView()
@@ -195,14 +160,9 @@ struct QuickLookTextEditor: NSViewRepresentable {
     textView.textContainer?.containerSize = NSSize(
       width: CGFloat.greatestFiniteMagnitude, height: CGFloat.greatestFiniteMagnitude)
 
-    // Apply syntax highlighting
+    // Apply syntax highlighting using tree-sitter
     if let textStorage = textView.textStorage {
-      let syntaxHighlighter = Wmap.SyntaxHighlighter()
-      syntaxHighlighter.updateSyntaxElements(from: parsedMap)
-
-      // Apply highlighting to the entire text
-      let fullRange = NSRange(location: 0, length: textStorage.length)
-      syntaxHighlighter.applySyntaxHighlighting(textStorage: textStorage, range: fullRange)
+      applyTreeSitterHighlighting(textStorage: textStorage, content: content)
     }
 
     return scrollView
@@ -211,4 +171,39 @@ struct QuickLookTextEditor: NSViewRepresentable {
   func updateNSView(_ nsView: NSScrollView, context: Context) {
     // No updates needed for read-only view
   }
+
+  private func applyTreeSitterHighlighting(textStorage: NSTextStorage, content: String) {
+    // Set up tree-sitter language layer
+    guard let wmapConfiguration = try? LanguageConfiguration(tree_sitter_wmap(), name: "wmap")
+    else { return }
+
+    let languageConfiguration = LanguageLayer.Configuration(languageProvider: { name in
+      switch name {
+      case "wmap": return wmapConfiguration
+      default: return nil
+      }
+    })
+
+    guard
+      let rootLayer = try? LanguageLayer(
+        languageConfig: wmapConfiguration, configuration: languageConfiguration)
+    else { return }
+
+    // Update layer with content
+    rootLayer.replaceContent(with: content)
+
+    // Set default text color
+    let fullRange = NSRange(location: 0, length: textStorage.length)
+    textStorage.addAttribute(.foregroundColor, value: NSColor.textColor, range: fullRange)
+
+    // Apply tree-sitter highlighting
+    guard let highlights = try? rootLayer.highlights(in: fullRange, provider: content.predicateTextProvider)
+    else { return }
+
+    for highlight in highlights {
+      let color = NSColor.Theme.Syntax.colorForSyntax(highlight.name)
+      textStorage.addAttribute(.foregroundColor, value: color, range: highlight.range)
+    }
+  }
+
 }