;; ==========================================================================
;; tree-tools.sx — Structural comprehension tools for s-expression files
;; ==========================================================================
;;
;; Pure functions for navigating, annotating, and understanding parsed
;; s-expression trees. Operates on the output of sx-parse (native SX
;; values: lists, symbols, strings, numbers, keywords, dicts).
;;
;; Phase 1: Read-only comprehension tools.
;; Phase 2 will add structural edit operations.

;; --------------------------------------------------------------------------
;; Path utilities
;; --------------------------------------------------------------------------

;; Format a path list as a display string: (0 2 1) → "[0,2,1]"
(define path-str :effects []
  (fn ((path :as list))
    (str "[" (join "," (map str path)) "]")))

;; --------------------------------------------------------------------------
;; annotate-tree — the primary comprehension tool
;; --------------------------------------------------------------------------
;;
;; Takes a parsed s-expression (or list of top-level expressions) and
;; returns a string with every node labeled by its path. This makes tree
;; structure explicit — no bracket counting needed.
;;
;; Example output:
;;   [0] (defcomp ~card
;;     [0,0] defcomp
;;     [0,1] ~card
;;     [0,2] (&key title)
;;     [0,3] (div :class "card"
;;       [0,3,0] div
;;       [0,3,1] :class
;;       [0,3,2] "card"
;;       [0,3,3] (h2 title)))

(define annotate-tree :effects []
  (fn (exprs)
    ;; Accept a single expression or a list of top-level expressions
    (let ((nodes (if (list? exprs) exprs (list exprs))))
      (let ((result (list)))
        (for-each (fn (i)
          (annotate-node (nth nodes i) (list i) 0 result))
          (range 0 (len nodes)))
        (join "\n" result)))))

(define annotate-node :effects []
  (fn (node path depth result)
    (let ((indent (join "" (map (fn (_) "  ") (range 0 depth))))
          (label (path-str path)))
      (if (list? node)
        ;; List node — show opening with path, then recurse into children
        (if (empty? node)
          (append! result (str indent label " ()"))
          (let ((head (first node))
                (head-str (node-display head)))
            ;; Compact form for short, simple lists (no nested lists)
            (if (and (<= (len node) 4)
                     (not (some (fn (c) (list? c)) (rest node))))
              ;; Single-line: [path] (head args...)
              (append! result
                (str indent label " ("
                  (join " " (map node-display node)) ")"))
              ;; Multi-line: show head on first line, children indented
              (do
                (append! result (str indent label " (" head-str))
                (for-each (fn (i)
                  (annotate-node (nth node i) (concat path (list i))
                                 (+ depth 1) result))
                  (range 1 (len node)))
                (append! result (str indent "  )"))))))
        ;; Atom node — display inline
        (append! result (str indent label " " (node-display node)))))))

;; Display a single node value as a short string
(define node-display :effects []
  (fn (node)
    (cond
      (nil? node) "nil"
      (= (type-of node) "symbol") (symbol-name node)
      (= (type-of node) "keyword") (str ":" (keyword-name node))
      (= (type-of node) "string")
        (let ((s (if (> (len node) 40) (str (slice node 0 37) "...") node)))
          (str "\"" s "\""))
      (= (type-of node) "number") (str node)
      (= (type-of node) "boolean") (if node "true" "false")
      (list? node)
        (if (empty? node) "()"
          (str "(" (node-display (first node))
            (if (> (len node) 1) " ..." "")
            ")"))
      (= (type-of node) "dict") "{...}"
      :else (str node))))


;; --------------------------------------------------------------------------
;; summarise — folded structural overview
;; --------------------------------------------------------------------------
;;
;; Shows the shape of a file without full detail. Each top-level form
;; is shown with its head and path. Nested forms are folded to a
;; configurable depth.

(define summarise :effects []
  (fn (exprs max-depth)
    (let ((nodes (if (list? exprs) exprs (list exprs)))
          (result (list)))
      (for-each (fn (i)
        (summarise-node (nth nodes i) (list i) 0 max-depth result))
        (range 0 (len nodes)))
      (join "\n" result))))

(define summarise-node :effects []
  (fn (node path depth max-depth result)
    (let ((indent (join "" (map (fn (_) "  ") (range 0 depth))))
          (label (path-str path)))
      (if (list? node)
        (if (empty? node)
          (append! result (str indent label " ()"))
          (let ((head (first node))
                (head-str (node-display head)))
            (if (>= depth max-depth)
              ;; At max depth — show head + child count
              (append! result
                (str indent label " (" head-str
                  (if (> (len node) 1)
                    (str " ... " (- (len node) 1) " children")
                    "")
                  ")"))
              ;; Below max depth — recurse
              (do
                (append! result (str indent label " (" head-str))
                (for-each (fn (i)
                  (summarise-node (nth node i) (concat path (list i))
                                  (+ depth 1) max-depth result))
                  (range 1 (len node)))
                (append! result (str indent "  )"))))))
        ;; Atom — show inline
        (append! result (str indent label " " (node-display node)))))))


;; --------------------------------------------------------------------------
;; read-subtree — expand a specific subtree
;; --------------------------------------------------------------------------

(define read-subtree :effects []
  (fn (exprs path)
    (let ((node (navigate exprs path)))
      (if (nil? node)
        (str "Error: path " (path-str path) " not found")
        (annotate-tree (list node))))))


;; --------------------------------------------------------------------------
;; get-context — show enclosing chain from root to target
;; --------------------------------------------------------------------------

(define get-context :effects []
  (fn (exprs path)
    (let ((result (list))
          (nodes (if (list? exprs) exprs (list exprs))))
      ;; Walk each prefix of the path
      (for-each (fn (depth)
        (let ((prefix (slice path 0 (+ depth 1)))
              (node (navigate nodes prefix)))
          (when (not (nil? node))
            (let ((label (path-str prefix))
                  (indent (join "" (map (fn (_) "  ") (range 0 depth))))
                  (marker (if (= (+ depth 1) (len path)) "→ " "  ")))
              (if (list? node)
                (append! result
                  (str indent marker label " " (node-summary node)))
                (append! result
                  (str indent marker label " " (node-display node))))))))
        (range 0 (len path)))
      (join "\n" result))))

;; One-line summary of a list node (head + key info)
(define node-summary :effects []
  (fn (node)
    (if (or (not (list? node)) (empty? node))
      (node-display node)
      (let ((head (node-display (first node)))
            (child-count (- (len node) 1)))
        (if (<= child-count 3)
          (str "(" (join " " (map node-display node)) ")")
          (str "(" head " " (node-display (nth node 1))
            (when (> child-count 1)
              (str " ... +" (- child-count 1)))
            ")"))))))


;; --------------------------------------------------------------------------
;; find-all — search for nodes matching a pattern
;; --------------------------------------------------------------------------

(define find-all :effects []
  (fn (exprs pattern)
    (let ((results (list))
          (nodes (if (list? exprs) exprs (list exprs))))
      (for-each (fn (i)
        (find-in-node (nth nodes i) (list i) pattern results))
        (range 0 (len nodes)))
      results)))

(define find-in-node :effects []
  (fn (node path pattern results)
    ;; Check if this node matches
    (when (node-matches? node pattern)
      (append! results (list path (node-summary-short node))))
    ;; Recurse into list children
    (when (list? node)
      (for-each (fn (i)
        (find-in-node (nth node i) (concat path (list i)) pattern results))
        (range 0 (len node))))))

;; Match: string pattern matches symbol names, string content
(define node-matches? :effects []
  (fn (node pattern)
    (cond
      ;; Symbol matches if name contains pattern
      (= (type-of node) "symbol")
        (contains? (symbol-name node) pattern)
      ;; String matches if value contains pattern
      (string? node)
        (contains? node pattern)
      ;; List matches if head is a symbol matching pattern
      (and (list? node) (not (empty? node))
           (= (type-of (first node)) "symbol"))
        (contains? (symbol-name (first node)) pattern)
      :else false)))

;; Short summary for search results
(define node-summary-short :effects []
  (fn (node)
    (if (list? node)
      (if (empty? node) "()"
        (let ((head (node-display (first node))))
          (if (> (len node) 3)
            (str "(" head " " (node-display (nth node 1)) " ...)")
            (str "(" (join " " (map node-display node)) ")"))))
      (node-display node))))


;; --------------------------------------------------------------------------
;; get-siblings — show siblings of a node
;; --------------------------------------------------------------------------

(define get-siblings :effects []
  (fn (exprs path)
    (if (empty? path)
      "Error: root has no siblings"
      (let ((parent-path (slice path 0 (- (len path) 1)))
            (target-idx (last path))
            (parent (navigate exprs parent-path)))
        (if (or (nil? parent) (not (list? parent)))
          (str "Error: parent at " (path-str parent-path) " not found or not a list")
          (let ((result (list)))
            (for-each (fn (i)
              (let ((child (nth parent i))
                    (child-path (concat parent-path (list i)))
                    (marker (if (= i target-idx) "→ " "  ")))
                (append! result
                  (str marker (path-str child-path) " "
                    (if (list? child)
                      (node-summary child)
                      (node-display child))))))
              (range 0 (len parent)))
            (join "\n" result)))))))


;; --------------------------------------------------------------------------
;; validate — structural integrity checks
;; --------------------------------------------------------------------------

(define validate :effects []
  (fn (exprs)
    (let ((errors (list))
          (nodes (if (list? exprs) exprs (list exprs))))
      (for-each (fn (i)
        (validate-node (nth nodes i) (list i) errors))
        (range 0 (len nodes)))
      (if (empty? errors)
        "OK"
        (join "\n" errors)))))

(define validate-node :effects []
  (fn (node path errors)
    (when (list? node)
      (when (not (empty? node))
        (let ((head (first node)))
          ;; Check: letrec bindings should all be pairs
          (when (and (= (type-of head) "symbol")
                     (= (symbol-name head) "letrec")
                     (>= (len node) 2))
            (let ((bindings (nth node 1)))
              (when (list? bindings)
                (for-each (fn (i)
                  (let ((pair (nth bindings i)))
                    (when (not (and (list? pair) (>= (len pair) 2)
                                   (= (type-of (first pair)) "symbol")))
                      (append! errors
                        (str "WARNING " (path-str (concat path (list 1 i)))
                          ": letrec binding " i " is not a (name value) pair: "
                          (node-display pair))))))
                  (range 0 (len bindings))))))
          ;; Check: defisland/defcomp body should exist
          (when (and (= (type-of head) "symbol")
                     (or (= (symbol-name head) "defisland")
                         (= (symbol-name head) "defcomp")))
            (when (< (len node) 4)
              (append! errors
                (str "ERROR " (path-str path)
                  ": " (symbol-name head) " has fewer than 3 args (name params body)"))))))
      ;; Recurse
      (for-each (fn (i)
        (validate-node (nth node i) (concat path (list i)) errors))
        (range 0 (len node))))))


;; --------------------------------------------------------------------------
;; navigate — resolve a path to a node
;; --------------------------------------------------------------------------

(define navigate :effects []
  (fn (exprs path)
    (let ((nodes (if (list? exprs) exprs (list exprs))))
      (reduce
        (fn (current idx)
          (if (or (nil? current) (not (list? current))
                  (>= idx (len current)))
            nil
            (nth current idx)))
        nodes
        path))))


;; ==========================================================================
;; Phase 2: Structural edit operations
;; ==========================================================================
;;
;; All operations take a tree (list of top-level expressions), perform the
;; edit, and return {"ok" new-tree} or {"error" message}. Nothing is mutated
;; in place. File writing is a separate concern.
;;
;; Key invariant: fragments are parsed BEFORE the tree is modified. If the
;; fragment is malformed, the operation returns an error immediately.

;; --------------------------------------------------------------------------
;; replace-node — replace the node at a path with new parsed source
;; --------------------------------------------------------------------------

(define replace-node :effects []
  (fn (exprs path new-source)
    (let ((fragment (sx-parse new-source)))
      (if (empty? fragment)
        {"error" (str "Fragment parse error: empty result from: " new-source)}
        (let ((new-node (if (= (len fragment) 1) (first fragment)
                          (cons (make-symbol "begin") fragment)))
              (result (tree-set exprs path new-node)))
          (if (nil? result)
            {"error" (str "Path not found: " (path-str path))}
            {"ok" result}))))))

;; --------------------------------------------------------------------------
;; insert-child — insert a new child at a specific index within a list
;; --------------------------------------------------------------------------

(define insert-child :effects []
  (fn (exprs path index new-source)
    (let ((fragment (sx-parse new-source)))
      (if (empty? fragment)
        {"error" (str "Fragment parse error: empty result from: " new-source)}
        (let ((new-node (if (= (len fragment) 1) (first fragment)
                          (cons (make-symbol "begin") fragment)))
              (parent (navigate exprs path)))
          (if (or (nil? parent) (not (list? parent)))
            {"error" (str "Path not found or not a list: " (path-str path))}
            (if (or (< index 0) (> index (len parent)))
              {"error" (str "Index " index " out of range for list with " (len parent) " children")}
              (let ((new-parent (list-insert parent index new-node))
                    (result (tree-set exprs path new-parent)))
                (if (nil? result)
                  {"error" (str "Failed to update tree at: " (path-str path))}
                  {"ok" result})))))))))

;; --------------------------------------------------------------------------
;; delete-node — remove a node at a path
;; --------------------------------------------------------------------------

(define delete-node :effects []
  (fn (exprs path)
    (if (empty? path)
      {"error" "Cannot delete root"}
      (let ((parent-path (slice path 0 (- (len path) 1)))
            (child-idx (last path))
            (parent (navigate exprs parent-path)))
        (if (or (nil? parent) (not (list? parent)))
          {"error" (str "Parent not found or not a list: " (path-str parent-path))}
          (if (or (< child-idx 0) (>= child-idx (len parent)))
            {"error" (str "Index " child-idx " out of range")}
            (let ((new-parent (list-remove parent child-idx))
                  (result (tree-set exprs parent-path new-parent)))
              (if (nil? result)
                {"error" "Failed to update tree"}
                {"ok" result}))))))))

;; --------------------------------------------------------------------------
;; wrap-node — wrap a node in a new list form
;; --------------------------------------------------------------------------
;;
;; (wrap-node tree [0,2] "(when condition _)")
;; Parses the wrapper, finds the placeholder _, replaces it with the
;; existing node at the path, then replaces the path with the result.

(define wrap-node :effects []
  (fn (exprs path wrapper-source)
    (let ((fragment (sx-parse wrapper-source)))
      (if (empty? fragment)
        {"error" (str "Wrapper parse error: empty result from: " wrapper-source)}
        (let ((wrapper (first fragment))
              (target (navigate exprs path)))
          (if (nil? target)
            {"error" (str "Path not found: " (path-str path))}
            (let ((filled (replace-placeholder wrapper target)))
              (if (nil? filled)
                {"error" "Wrapper must contain _ as placeholder for the wrapped node"}
                (let ((result (tree-set exprs path filled)))
                  (if (nil? result)
                    {"error" "Failed to update tree"}
                    {"ok" result}))))))))))

;; Replace the symbol _ in a tree with a replacement value (first occurrence)
(define replace-placeholder :effects []
  (fn (node replacement)
    (cond
      (and (= (type-of node) "symbol") (= (symbol-name node) "_"))
        replacement
      (list? node)
        (let ((found false)
              (result (map (fn (child)
                (if found
                  child
                  (if (and (= (type-of child) "symbol") (= (symbol-name child) "_"))
                    (do (set! found true) replacement)
                    (if (list? child)
                      (let ((sub (replace-placeholder child replacement)))
                        (if (nil? sub) child (do (set! found true) sub)))
                      child))))
                node)))
          (if found result nil))
      :else nil)))


;; --------------------------------------------------------------------------
;; Tree manipulation helpers
;; --------------------------------------------------------------------------

;; Set a node at a path in the tree, returning a new tree.
;; Immutable — builds new lists along the path.
(define tree-set :effects []
  (fn (exprs path new-node)
    (let ((nodes (if (list? exprs) exprs (list exprs))))
      (if (empty? path)
        (if (list? new-node) new-node (list new-node))
        (tree-set-inner nodes path 0 new-node)))))

(define tree-set-inner :effects []
  (fn (node path depth new-node)
    (if (not (list? node))
      nil
      (let ((idx (nth path depth)))
        (if (or (< idx 0) (>= idx (len node)))
          nil
          (if (= depth (- (len path) 1))
            ;; At target depth — replace the child
            (list-replace node idx new-node)
            ;; Not at target — recurse into child
            (let ((child (nth node idx))
                  (new-child (tree-set-inner child path (+ depth 1) new-node)))
              (if (nil? new-child)
                nil
                (list-replace node idx new-child)))))))))

;; Replace element at index in a list (immutable — returns new list)
(define list-replace :effects []
  (fn (lst idx val)
    (map-indexed (fn (i item) (if (= i idx) val item)) lst)))

;; Insert element at index in a list (immutable — returns new list)
(define list-insert :effects []
  (fn (lst idx val)
    (concat (slice lst 0 idx) (list val) (slice lst idx))))

;; Remove element at index from a list (immutable — returns new list)
(define list-remove :effects []
  (fn (lst idx)
    (concat (slice lst 0 idx) (slice lst (+ idx 1)))))