;; lib/datalog/api.sx — SX-data embedding API.
;;
;; Where Phase 1's `dl-program` takes a Datalog source string,
;; this module exposes a parser-free API that consumes SX data
;; directly. Two rule shapes are accepted:
;;
;;   - dict:   {:head <literal> :body (<literal> ...)}
;;   - list:   (<head-elements...> <- <body-literal> ...)
;;             — `<-` is an SX symbol used as the rule arrow.
;;
;; Examples:
;;
;;   (dl-program-data
;;     '((parent tom bob) (parent tom liz) (parent bob ann))
;;     '((ancestor X Y <- (parent X Y))
;;       (ancestor X Z <- (parent X Y) (ancestor Y Z))))
;;
;;   (dl-query db '(ancestor tom X))   ; same query API as before
;;
;; Variables follow the parser convention: SX symbols whose first
;; character is uppercase or `_` are variables.

(define
  dl-rule
  (fn (head body) {:head head :body body}))

(define
  dl-rule-arrow?
  (fn
    (x)
    (and (symbol? x) (= (symbol->string x) "<-"))))

(define
  dl-find-arrow
  (fn
    (rl i n)
    (cond
      ((>= i n) nil)
      ((dl-rule-arrow? (nth rl i)) i)
      (else (dl-find-arrow rl (+ i 1) n)))))

;; Given a list of the form (head-elt ... <- body-lit ...) returns
;; {:head (head-elt ...) :body (body-lit ...)}. If no arrow is
;; present, the whole list is treated as the head and the body is
;; empty (i.e. a fact written rule-style).
(define
  dl-rule-from-list
  (fn
    (rl)
    (let ((n (len rl)))
      (let ((idx (dl-find-arrow rl 0 n)))
        (cond
          ((nil? idx) {:head rl :body (list)})
          (else
            (let
              ((head (slice rl 0 idx))
                (body (slice rl (+ idx 1) n)))
              {:head head :body body})))))))

;; Coerce a rule given as either a dict or a list-with-arrow to a dict.
(define
  dl-coerce-rule
  (fn
    (r)
    (cond
      ((dict? r) r)
      ((list? r) (dl-rule-from-list r))
      (else (error (str "dl-coerce-rule: expected dict or list, got " r))))))

;; Build a db from SX data lists.
(define
  dl-program-data
  (fn
    (facts rules)
    (let ((db (dl-make-db)))
      (do
        (for-each (fn (lit) (dl-add-fact! db lit)) facts)
        (for-each
          (fn (r) (dl-add-rule! db (dl-coerce-rule r)))
          rules)
        db))))

;; Add a single fact at runtime, then re-saturate the db so derived
;; tuples reflect the change. Returns the db.
(define
  dl-assert!
  (fn
    (db lit)
    (do
      (dl-add-fact! db lit)
      (dl-saturate! db)
      db)))

;; Remove a fact and re-saturate. Mixed relations (which have BOTH
;; user-asserted facts AND rules) are supported via :edb-keys provenance
;; — explicit facts are marked at dl-add-fact! time, the saturator uses
;; dl-add-derived! which doesn't mark them, so the retract pass can
;; safely wipe IDB-derived tuples while preserving the user's EDB.
;;
;; Effect:
;;   - remove tuples matching `lit` from :facts and :edb-keys
;;   - for every relation that has a rule (i.e. potentially IDB or
;;     mixed), drop the IDB-derived portion (anything not in :edb-keys)
;;     so the saturator can re-derive cleanly
;;   - re-saturate
(define
  dl-retract!
  (fn
    (db lit)
    (let
      ((rel-key (dl-rel-name lit)))
      (do
        ;; Drop the matching tuple from its relation list, its facts-keys,
        ;; its first-arg index, AND from :edb-keys (if present).
        (when
          (has-key? (get db :facts) rel-key)
          (let
            ((existing (get (get db :facts) rel-key))
              (kept (list))
              (kept-keys {})
              (kept-index {})
              (edb-rel (cond
                         ((has-key? (get db :edb-keys) rel-key)
                           (get (get db :edb-keys) rel-key))
                         (else nil)))
              (kept-edb {}))
            (do
              (for-each
                (fn
                  (t)
                  (when
                    (not (dl-tuple-equal? t lit))
                    (do
                      (append! kept t)
                      (let ((tk (dl-tuple-key t)))
                        (do
                          (dict-set! kept-keys tk true)
                          (when
                            (and (not (nil? edb-rel))
                                 (has-key? edb-rel tk))
                            (dict-set! kept-edb tk true))))
                      (when
                        (>= (len t) 2)
                        (let ((k (dl-arg-key (nth t 1))))
                          (do
                            (when
                              (not (has-key? kept-index k))
                              (dict-set! kept-index k (list)))
                            (append! (get kept-index k) t)))))))
                existing)
              (dict-set! (get db :facts) rel-key kept)
              (dict-set! (get db :facts-keys) rel-key kept-keys)
              (dict-set! (get db :facts-index) rel-key kept-index)
              (when
                (not (nil? edb-rel))
                (dict-set! (get db :edb-keys) rel-key kept-edb)))))
        ;; For each rule-head relation, strip the IDB-derived tuples
        ;; (anything not marked in :edb-keys) so the saturator can
        ;; cleanly re-derive without leaving stale tuples that depended
        ;; on the now-removed fact.
        (let ((rule-heads (dl-rule-head-rels db)))
          (for-each
            (fn
              (k)
              (when
                (has-key? (get db :facts) k)
                (let
                  ((existing (get (get db :facts) k))
                    (kept (list))
                    (kept-keys {})
                    (kept-index {})
                    (edb-rel (cond
                               ((has-key? (get db :edb-keys) k)
                                 (get (get db :edb-keys) k))
                               (else {}))))
                  (do
                    (for-each
                      (fn
                        (t)
                        (let ((tk (dl-tuple-key t)))
                          (when
                            (has-key? edb-rel tk)
                            (do
                              (append! kept t)
                              (dict-set! kept-keys tk true)
                              (when
                                (>= (len t) 2)
                                (let ((kk (dl-arg-key (nth t 1))))
                                  (do
                                    (when
                                      (not (has-key? kept-index kk))
                                      (dict-set! kept-index kk (list)))
                                    (append! (get kept-index kk) t))))))))
                      existing)
                    (dict-set! (get db :facts) k kept)
                    (dict-set! (get db :facts-keys) k kept-keys)
                    (dict-set! (get db :facts-index) k kept-index)))))
            rule-heads))
        (dl-saturate! db)
        db))))

;; ── Convenience: single-call source + query ───────────────────
;; (dl-eval source query-source) parses both, builds a db, saturates,
;; runs the query, returns the substitution list. The query source
;; should be `?- goal[, goal ...].` — the parser produces a clause
;; with :query containing a list of literals which is fed straight
;; to dl-query.
(define
  dl-eval
  (fn
    (source query-source)
    (let
      ((db (dl-program source))
        (queries (dl-parse query-source)))
      (cond
        ((= (len queries) 0) (error "dl-eval: query string is empty"))
        ((not (has-key? (first queries) :query))
          (error "dl-eval: second arg must be a `?- ...` query clause"))
        (else
          (dl-query db (get (first queries) :query)))))))

;; (dl-eval-magic source query-source) — like dl-eval but routes a
;; single-positive-literal query through `dl-magic-query` for goal-
;; directed evaluation. Multi-literal query bodies fall back to the
;; standard dl-query path (magic-sets is currently only wired for
;; single-positive goals). The caller's source is parsed afresh
;; each call so successive invocations are independent.
(define
  dl-eval-magic
  (fn
    (source query-source)
    (let
      ((db (dl-program source))
        (queries (dl-parse query-source)))
      (cond
        ((= (len queries) 0) (error "dl-eval-magic: query string is empty"))
        ((not (has-key? (first queries) :query))
          (error
            "dl-eval-magic: second arg must be a `?- ...` query clause"))
        (else
          (let
            ((qbody (get (first queries) :query)))
            (cond
              ((and (= (len qbody) 1)
                    (list? (first qbody))
                    (> (len (first qbody)) 0)
                    (symbol? (first (first qbody))))
                (dl-magic-query db (first qbody)))
              (else (dl-query db qbody)))))))))

;; List rules whose head's relation matches `rel-name`. Useful for
;; inspection ("show me how this relation is derived") without
;; exposing the internal `:rules` list.
(define
  dl-rules-of
  (fn
    (db rel-name)
    (let ((out (list)))
      (do
        (for-each
          (fn
            (rule)
            (when
              (= (dl-rel-name (get rule :head)) rel-name)
              (append! out rule)))
          (dl-rules db))
        out))))

(define
  dl-rule-head-rels
  (fn
    (db)
    (let ((seen (list)))
      (do
        (for-each
          (fn
            (rule)
            (let ((h (dl-rel-name (get rule :head))))
              (when
                (and (not (nil? h)) (not (dl-member-string? h seen)))
                (append! seen h))))
          (dl-rules db))
        seen))))

;; Wipe every relation that has at least one rule (i.e. every IDB
;; relation) — leaves EDB facts and rule definitions intact. Useful
;; before a follow-up `dl-saturate!` if you want a clean restart, or
;; for inspection of the EDB-only baseline.
(define
  dl-clear-idb!
  (fn
    (db)
    (let ((rule-heads (dl-rule-head-rels db)))
      (do
        (for-each
          (fn
            (k)
            (do
              (dict-set! (get db :facts) k (list))
              (dict-set! (get db :facts-keys) k {})
              (dict-set! (get db :facts-index) k {})))
          rule-heads)
        db))))