;;; Arl Standard Library - Dictionary (Hash Table)

(module dict
  (export dict hash dict? dict-get dict-set dict-remove dict-keys dict-values
          dict-has? dict-merge
          dict-update dict-map dict-filter dict-for-each dict->alist alist->dict)

  (import core :refer (r-call error))
  (import functional :refer (filter))

;;' @section Dictionary (Hash Table)
;;' Dictionaries are mutable key-value stores backed by R environments.
;;' Keys must be strings, symbols, or keywords; values can be anything.

;;; Internal helpers (defined early for use in constructors)

;;' @description Create a new empty dictionary environment.
(define dict-new
  (lambda ()
    (define dict-obj (r-call "new.env" (list :hash #t :parent (r-call "emptyenv" (list)))))
    (r-call "class<-" (list dict-obj (c (list "arl_dict") (r-call "class" (list dict-obj)))))
    (r-call "assign" (list ".__keys" (r-call "character" (list 0)) :envir dict-obj))
    dict-obj))

;;' @description Convert key to string name for dict storage.
;;' @param key String, symbol, or keyword to convert
(define dict-key-to-name
  (lambda (key)
    (if (r-call "is.character" (list key))
      key
      (if (or (r-call "inherits" (list key "arl_keyword")) (r-call "is.symbol" (list key)))
        (r-call "as.character" (list key))
        #nil))))

;;; Dictionary constructor

;;' @description Create a hash-backed dictionary from key/value pairs.
;;' @param args Named key/value pairs (e.g. :x 1 :y 2)
;;' @examples
;;' (dict :x 1 :y 2)           ; => dict with keys "x" and "y"
;;' (dict)                      ; => empty dict
;;' (define d (dict :name "Alice" :age 30))
;;' (dict-get d "name")         ; => "Alice"
;;' @assert
;;' (assert-true (dict? (dict :x 1 :y 2)))
;;' (assert-true (dict? (dict)))
;;' (assert-equal "Alice" (dict-get (dict :name "Alice" :age 30) "name"))
;;' @seealso hash
(define dict
  (lambda (. args)
    (define dict-obj (dict-new))
    (if (= (length args) 0)
      dict-obj
      (begin
        (define arg-names (r-call "names" (list args)))
        (if (or (is.null arg-names) (r-call "any" (list (r-call "!" (list (r-call "nzchar" (list arg-names)))))))
          (error "dict requires named arguments")
          (begin
            (define keys (r-call "as.character" (list arg-names)))
            (define assign-pair
              (lambda (i)
                (begin
                  (define key (r-call "[[" (list keys i)))
                  (define value (r-call "[[" (list args i)))
                  (r-call "assign" (list key value :envir dict-obj)))))
            (r-call "lapply" (list (r-call "seq_len" (list (length keys))) assign-pair))
            (r-call "assign" (list ".__keys" keys :envir dict-obj))
            dict-obj))))))

;;' @description Alias for dict. Create a dictionary from key/value pairs.
;;' @signature (hash . args)
;;' @examples
;;' (hash :x 1 :y 2)           ; => dict with x=1, y=2
;;' @assert
;;' (assert-equal 1 (dict-get (hash :x 1 :y 2) "x"))
;;' @note Alias for `dict`.
;;' @seealso dict
(define hash dict)

;;; Type predicate

;;' @description Return #t if x is a dictionary.
;;' @param x Value to test
;;' @examples
;;' (dict? (dict :x 1))        ; => #t
;;' (dict? (list 1 2))         ; => #f
;;' (dict? 42)                 ; => #f
;;' @assert
;;' (assert-true (dict? (dict :x 1)))
;;' (assert-false (dict? (list 1 2)))
;;' (assert-false (dict? 42))
(define dict?
  (lambda (x)
    (and (r-call "is.environment" (list x)) (r-call "inherits" (list x "arl_dict")))))

;;; Internal helper (defined early for forward references)

;;' @description Return ordered keys vector (internal helper).
;;' @param dict Dictionary to retrieve keys from
(define dict-keys-ordered
  (lambda (dict)
    (if (not (dict? dict))
      (r-call "character" (list 0))
      (begin
        (define keys (r-call "get0" (list ".__keys" :envir dict :inherits #f)))
        (if (is.null keys)
          (r-call "sort" (list (r-call "ls" (list :envir dict :all.names #f))))
          keys)))))

;;; Dictionary operations

;;' @description Get value for key or default if missing.
;;' @param dict Dictionary to look up in
;;' @param key Key to retrieve (string, symbol, or keyword)
;;' @param rest Optional default value if key is missing (defaults to #nil)
;;' @examples
;;' (define d (dict :x 1 :y 2))
;;' (dict-get d "x")           ; => 1
;;' (dict-get d "z")           ; => #nil
;;' (dict-get d "z" 99)        ; => 99  (default value)
;;' @assert
;;' (assert-equal 1 (dict-get (dict :x 1 :y 2) "x"))
;;' (assert-equal #nil (dict-get (dict :x 1 :y 2) "z"))
;;' (assert-equal 99 (dict-get (dict :x 1 :y 2) "z" 99))
;;' @seealso dict-set, dict-has?
(define dict-get
  (lambda (dict key . rest)
    (define default (if (= (length rest) 0) #nil (car rest)))
    (define name (dict-key-to-name key))
    (if (or (is.null name) (not (dict? dict)))
      default
      (if (r-call "exists" (list name :envir dict :inherits #f))
        (r-call "get" (list name :envir dict :inherits #f))
        default))))

;;' @description Set key to value in dict and return dict.
;;' @param dict Dictionary to modify
;;' @param key Key to set (string, symbol, or keyword)
;;' @param value Value to associate with the key
;;' @examples
;;' (define d (dict))
;;' (dict-set d "x" 42)        ; => d (mutated, x=42)
;;' (dict-get d "x")           ; => 42
;;' @assert
;;' (define _test-d (dict))
;;' (dict-set _test-d "x" 42)
;;' (assert-equal 42 (dict-get _test-d "x"))
;;' @seealso dict-get, dict-remove
(define dict-set
  (lambda (dict key value)
    (define name (dict-key-to-name key))
    (if (is.null name)
      (error "dict-set requires a string, symbol, or keyword key")
      (if (not (dict? dict))
        (error "dict-set requires a dict")
        (begin
          (define keys (dict-keys-ordered dict))
          (if (not (r-call "%in%" (list name keys)))
            (r-call "assign" (list ".__keys" (c keys (list name)) :envir dict)))
          (r-call "assign" (list name value :envir dict))
          dict)))))

;;' @description Remove key from dict and return dict.
;;' @param dict Dictionary to modify
;;' @param key Key to remove (string, symbol, or keyword)
;;' @examples
;;' (define d (dict :x 1 :y 2))
;;' (dict-remove d "x")        ; => d (mutated, x removed)
;;' (dict-has? d "x")          ; => #f
;;' @assert
;;' (define _test-dr (dict :x 1 :y 2))
;;' (dict-remove _test-dr "x")
;;' (assert-false (dict-has? _test-dr "x"))
;;' @seealso dict-set
(define dict-remove
  (lambda (dict key)
    (define name (dict-key-to-name key))
    (if (or (is.null name) (not (dict? dict)))
      dict
      (if (r-call "exists" (list name :envir dict :inherits #f))
        (begin
          (r-call "rm" (list :list (c name) :envir dict :inherits #f))
          (define keys (dict-keys-ordered dict))
          (if (> (length keys) 0)
            (r-call "assign" (list ".__keys" (r-call "as.character" (list (filter (lambda (k) (not (= k name))) (as.list keys)))) :envir dict)))
          dict)
        dict))))

;;' @description Return a list of dict keys.
;;' @param dict Dictionary to retrieve keys from
;;' @examples
;;' (define d (dict :x 1 :y 2))
;;' (dict-keys d)              ; => ("x" "y")
;;' (dict-keys (dict))         ; => ()
;;' @assert
;;' (assert-equal (list "x" "y") (dict-keys (dict :x 1 :y 2)))
;;' (assert-equal (list) (dict-keys (dict)))
;;' @seealso dict-values
(define dict-keys
  (lambda (dict)
    (if (not (dict? dict))
      (list)
      (begin
        (define keys (dict-keys-ordered dict))
        (_as-list keys)))))

;;' @description Return a list of dict values.
;;' @param dict Dictionary to retrieve values from
;;' @examples
;;' (define d (dict :x 1 :y 2))
;;' (dict-values d)            ; => (1 2)
;;' @assert
;;' (assert-equal (list 1 2) (dict-values (dict :x 1 :y 2)))
;;' @seealso dict-keys
(define dict-values
  (lambda (dict)
    (if (not (dict? dict))
      (list)
      (begin
        (define keys (dict-keys-ordered dict))
        (if (= (length keys) 0)
          (list)
          (begin
            (define values (r-call "mget" (list keys :envir dict :inherits #f)))
            (_as-list (r-call "unname" (list values)))))))))

;;' @description Return #t if dict contains key.
;;' @param dict Dictionary to check
;;' @param key Key to look for (string, symbol, or keyword)
;;' @examples
;;' (define d (dict :x 1 :y 2))
;;' (dict-has? d "x")          ; => #t
;;' (dict-has? d "z")          ; => #f
;;' @assert
;;' (assert-true (dict-has? (dict :x 1 :y 2) "x"))
;;' (assert-false (dict-has? (dict :x 1 :y 2) "z"))
;;' @seealso dict-get
(define dict-has?
  (lambda (dict key)
    (define name (dict-key-to-name key))
    (if (or (is.null name) (not (dict? dict)))
      #f
      (r-call "exists" (list name :envir dict :inherits #f)))))

;;' @description Merge dicts, later values override earlier.
;;' @param dicts Dictionaries to merge in order (later values win)
;;' @examples
;;' (define a (dict :x 1 :y 2))
;;' (define b (dict :y 99 :z 3))
;;' (define m (dict-merge a b))
;;' (dict-get m "y")           ; => 99  (b overrides a)
;;' (dict-get m "z")           ; => 3
;;' @assert
;;' (define _test-ma (dict :x 1 :y 2))
;;' (define _test-mb (dict :y 99 :z 3))
;;' (define _test-mm (dict-merge _test-ma _test-mb))
;;' (assert-equal 99 (dict-get _test-mm "y"))
;;' (assert-equal 3 (dict-get _test-mm "z"))
;;' @seealso dict, dict-set
(define dict-merge
  (lambda (. dicts)
    (define result (dict-new))
    (define result-keys (list))
    (define merge-dict
      (lambda (d)
        (if (dict? d)
          (begin
            (define keys (dict-keys-ordered d))
            (define merge-key
              (lambda (name)
                (begin
                  (define value (r-call "get" (list name :envir d :inherits #f)))
                  (r-call "assign" (list name value :envir result))
                  (if (not (r-call "%in%" (list name (r-call "as.character" (list result-keys)))))
                    (set! result-keys (c result-keys (list name)))))))
            (if (> (length keys) 0)
              (r-call "lapply" (list (as.list keys) merge-key)))))))
    (define n-dicts (length dicts))
    (define di 1)
    (while (<= di n-dicts)
      (begin
        (merge-dict ([[ dicts di))
        (set! di (+ di 1))))
    (r-call "assign" (list ".__keys" (r-call "as.character" (list result-keys)) :envir result))
    result))

;; ============================================================================
;; Dictionary Transformation
;; ============================================================================

;;' @section Dictionary Transformation
;;' Functions for updating, mapping, filtering, and converting dictionaries.

;;' @description Apply fn to the value at key (or default if missing).
;;' Mutates and returns dict.
;;' @param d Dictionary to update
;;' @param key Key to update
;;' @param fn Function to apply to the current value
;;' @param rest Optional default value if key is missing
;;' @examples
;;' (dict-update (dict :x 1) "x" inc)  ; => dict with x=2
;;' @seealso dict-set, dict-get
(define dict-update
  (lambda (d key fn . rest)
    (define name (dict-key-to-name key))
    (define current
      (if (r-call "exists" (list name :envir d :inherits #f))
        (r-call "get" (list name :envir d :inherits #f))
        (if (> (length rest) 0)
          (car rest)
          (error "dict-update: key not found and no default provided"))))
    (dict-set d key (fn current))))

;;' @description Map fn over dict entries. fn takes key and value, returns new value.
;;' Returns a new dict.
;;' @param fn Function taking key (string) and value, returning new value
;;' @param d Dictionary to map over
;;' @examples
;;' (dict-map (lambda (k v) (* v 10)) (dict :x 1))  ; => dict with x=10
;;' @seealso dict-filter, dict-for-each
(define dict-map
  (lambda (fn d)
    (define result (dict-new))
    (define keys (dict-keys-ordered d))
    (if (> (length keys) 0)
      (begin
        (define process-key
          (lambda (name)
            (begin
              (define value (r-call "get" (list name :envir d :inherits #f)))
              (define new-value (fn name value))
              (r-call "assign" (list name new-value :envir result)))))
        (r-call "lapply" (list (as.list keys) process-key))
        (r-call "assign" (list ".__keys" keys :envir result))))
    result))

;;' @description Filter dict entries. pred takes key and value, returns #t to keep.
;;' Returns a new dict.
;;' @param pred Predicate function taking key (string) and value
;;' @param d Dictionary to filter
;;' @examples
;;' (dict-filter (lambda (k v) (> v 1)) (dict :x 1 :y 2))  ; => dict with y=2
;;' @seealso dict-map, dict-for-each
(define dict-filter
  (lambda (pred d)
    (define result (dict-new))
    (define keys (dict-keys-ordered d))
    (define result-keys (list))
    (if (> (length keys) 0)
      (begin
        (define process-key
          (lambda (name)
            (begin
              (define value (r-call "get" (list name :envir d :inherits #f)))
              (if (isTRUE (pred name value))
                (begin
                  (r-call "assign" (list name value :envir result))
                  (set! result-keys (c result-keys (list name))))))))
        (r-call "lapply" (list (as.list keys) process-key))
        (r-call "assign" (list ".__keys" (r-call "as.character" (list result-keys)) :envir result))))
    result))

;;' @description Iterate over dict entries for side effects. fn takes key and value.
;;' Returns #nil.
;;' @param fn Function taking key (string) and value
;;' @param d Dictionary to iterate over
;;' @examples
;;' (dict-for-each (lambda (k v) (print v)) (dict :x 1))  ; prints 1
;;' @seealso dict-map, dict-filter
(define dict-for-each
  (lambda (fn d)
    (define keys (dict-keys-ordered d))
    (if (> (length keys) 0)
      (begin
        (define process-key
          (lambda (name)
            (begin
              (define value (r-call "get" (list name :envir d :inherits #f)))
              (fn name value))))
        (r-call "lapply" (list (as.list keys) process-key))))
    #nil))

;;' @description Convert dict to association list of (key value) pairs.
;;' @param d Dictionary to convert
;;' @examples
;;' (dict->alist (dict :x 1 :y 2))  ; => (("x" 1) ("y" 2))
;;' @seealso alist->dict
(define dict->alist
  (lambda (d)
    (define keys (dict-keys-ordered d))
    (if (= (length keys) 0)
      (list)
      (begin
        (define make-pair
          (lambda (name)
            (list name (r-call "get" (list name :envir d :inherits #f)))))
        (r-call "lapply" (list (as.list keys) make-pair))))))

;;' @description Convert association list of (key value) pairs to dict.
;;' @param alist List of (key value) pairs
;;' @examples
;;' (alist->dict (list (list "x" 1) (list "y" 2)))  ; => dict with x=1 y=2
;;' @seealso dict->alist
(define alist->dict
  (lambda (alist)
    (define result (dict-new))
    (define items (_as-list alist))
    (define n (length items))
    (define result-keys (list))
    (define i 1)
    (while (<= i n)
      (begin
        (define pair ([[ items i))
        (define key (r-call "as.character" (list (r-call "[[" (list pair 1)))))
        (define value (r-call "[[" (list pair 2)))
        (r-call "assign" (list key value :envir result))
        (set! result-keys (c result-keys (list key)))
        (set! i (+ i 1))))
    (r-call "assign" (list ".__keys" (r-call "as.character" (list result-keys)) :envir result))
    result))
)
