(ns expand
  (:require [clojure.core.match :refer [match]]
            [clojure.pprint]
            [clojure.set :as set]
            [clojure.spec.alpha :as s]
            [clojure.walk]
            [spec-dict.main :refer [dict]]
            [scm]
            [clojure.java.io :as io]))

(defmacro inherit [& syms]
  (let [f (if (<= 8 (count syms))
            `hash-map
            `array-map)]
    `(~f ~@(mapcat #(list (keyword %) %)
                   syms))))

(s/def ::scope scm/scope?)

(s/def ::binding
  #(match %
     [:lexical sym uniq] (and (symbol? sym) (symbol? uniq))
     [:global sym] (symbol? sym)))

(s/def ::context (s/coll-of ::scope :kind set?))

;; this is an inefficient representation. since we often want to
;; compute the submap of entries with a symbol 'x' but any scope, it
;; would be best to 'curry' the map, as Symbol → Scopes →
;; Binding. however, this is a bit more error-prone and cumbersome for
;; a prototype.
(s/def ::store
  (s/map-of (s/tuple symbol? (s/coll-of ::scope :kind set?))
            ::binding))

(def primitives
  "A set of symbols to be defined as primitive syntax exposed by
  (primitives+builtins)."
  '#{let let-syntax letrec lambda syntax-rules if define
     define-syntax begin})

(def builtins
  "A set of symbols to be defined as variables exposed by
  (primitives+builtins)."
  '#{+ * / - car cdr write display newline cons list})

(s/def ::denotation
  #(match %
     [:var] true
     [:primitive p] (primitives p)
     ;; 나쁜
     [:macro trans] (contains? trans :rules)
     [:syntax datum] true))

(s/def ::environment
  (s/map-of ::binding ::denotation))

(s/def ::syntax scm/syntax?)

(s/def ::identifier scm/identifier?)

(s/fdef expand
  :args (s/cat :stx ::syntax
               :env ::environment
               :store ::store)
  :ret ::syntax)

(s/fdef lookup
  :args (s/cat :ident ::identifier
               :env ::environment
               :store ::store))

(s/fdef nearest
  :args (s/cat :ctx ::context
               :ss (s/coll-of ::context))
  :ret (s/coll-of ::context :kind set?))


;;; prim specs

(s/def ::datum scm/datum?)

(s/def ::unsyntax
  (s/conformer
   (fn [x]
     (match x
       (_ :guard scm/syntax?) (scm/syntax-e x)
       _ ::s/invalid))))

(s/def ::body
  (s/+ ::datum))

(s/def ::lambda-formals
  (s/and ::unsyntax
         #(or (and (list? %) (apply (every-pred scm/identifier?) %))
              scm/improper-list?
              scm/identifier?)))

(s/def ::prim-lambda
  (s/and ::unsyntax
         (s/cat :keyword ::identifier
          :formals ::lambda-formals
          :body ::body)))

(s/def ::prim-let
  (s/and ::unsyntax
         (s/cat :keyword ::identifier
                :bindings
                (s/and ::unsyntax
                       (s/coll-of
                        (s/and ::unsyntax
                               (s/cat :rhs ::identifier
                                      :lhs ::datum))
                        :kind sequential?))
                :body ::body)))

(s/def ::prim-if
  (s/and ::unsyntax
         (s/cat :keyword ::identifier
                :condition ::datum
                :consequent ::datum
                :alternative ::datum)))



#_
(s/fdef binding-form
  :args (s/cat :binds (s/coll-of (s/tuple symbol? ::denotation))
               :store-atom #(instance? clojure.lang.Atom %)
               :datum ::datum))

(defn binding-form [scopeholder env store-atom binds f]
  (let [scopes (conj (scm/syntax-context scopeholder)
                     (scm/new-scope))
        scopeholder* (scm/->Syntax (scm/syntax-e scopeholder) scopes)
        additions (for [[sym denot] binds]
                    [sym
                     [:lexical sym (gensym)]
                     denot])
        env* (into env (for [[_sym bind denot] additions]
                         [bind denot]))]
    (swap! store-atom #(into % (for [[sym bind _denot] additions]
                                 [[sym scopes] bind])))
    (scm/datum->syntax (f scopeholder* env*)
                       scopeholder)))

(comment
  (let [env {}
        store-atom (atom {})]
    (binding-form #scm/stx [lambda #{}]
                  env
                  store-atom
                  {:x [:var]}
                  (fn [ident* env*]
                    (tap> (inherit ident* env*))
                    (list (scm/datum->syntax 'x ident*))))))



(defn- store-scopes
  "Gather the set of scope-sets associated with a given symbol."
  [store sym]
  (into #{} (for [[[x scopes] _] store
                  :when (= sym x)]
              scopes)))

(defn nearest [ctx ss]
  (let [xs (->> ss
                (filter #(set/subset? % ctx))
                (group-by count)
                seq)]
    (set (some->> xs (apply max-key first) second))))

(defn lookup [ident env store]
  (let [sym (scm/syntax-e ident)
        scopes (scm/syntax-context ident)
        ss (store-scopes store sym)
        xs (nearest scopes ss)
        r (compare (count xs) 1)]
    (tap> (inherit ident env store))
    (cond (neg? r) (throw (ex-info "undefined sym" {:sym sym}))
          (zero? r)
          (or (get env (or (get store [sym (first xs)])
                           (throw (ex-info "unreachable" {}))))
              (throw (ex-info "unreachable" {})))
          (pos? r) (throw (ex-info "ambiguous sym" {:sym sym})))))

(declare expand)

(defmulti expand-primitive
  (fn [_datum prim-name _env _store-atom]
    prim-name))

(defmethod expand-primitive 'lambda
  [stx _ env store-atom]
  (let [{:keys [formals body]} (s/conform ::prim-lambda stx)]
    (assert (list? formals))
    (binding-form
     stx env store-atom
     (doall (for [x formals]
              [(scm/syntax-e x) [:var]]))
     (fn [scopeholder env*]
       (concat (list 'lambda
                     (for [x formals]
                       (scm/rescope x scopeholder)))
               (for [x body]
                 (expand (scm/rescope x scopeholder)
                         env* store-atom)))))))

(defmethod expand-primitive 'let
  [stx _ env store-atom]
  (let [{:keys [bindings body]} (s/conform ::prim-let stx)]
    (assert (list? bindings))
    (binding-form
     stx env store-atom
     (doall (for [{:keys [rhs]} bindings]
              [(scm/syntax-e rhs) [:var]]))
     (fn [scopeholder env*]
       (concat (list 'let
                     (for [{:keys [rhs lhs]} bindings]
                       (list (scm/rescope rhs scopeholder) lhs)))
               (for [x body]
                 (expand (scm/rescope x scopeholder)
                         env* store-atom)))))))

(defmethod expand-primitive 'if
  [stx _ env store-atom]
  (let [{:keys [condition consequent alternative]}
        (s/conform ::prim-if stx)]
    (assert ((every-pred scm/datum?) condition consequent alternative))
    (cons 'if
          (for [x [condition consequent alternative]]
            (expand x env store-atom)))))

(defmethod expand-primitive 'let-syntax
  [stx _ env store-atom]
  (let [_ (s/conform ::prim-let-syntax stx)]))

(defmethod expand-primitive 'syntax-rules
  [stx _ env store-atom]
  (let [_ (s/conform ::prim-syntax-rules stx)]))

(defn expand
  ([stx]
   (let [scopeholder (scm/new-scopeholder)
         env {}
         store-atom (atom {})]
     (expand (scm/rescope stx scopeholder)
             env store-atom)))
  ([stx env store-atom]
   (tap> (inherit stx))
   (match (scm/syntax-e stx)
     (_ :guard scm/constant?) stx
     (_ :guard symbol?)
     (do (lookup stx env @store-atom) ; throw if unbound
         stx)

     ([f & xs] :seq)
     (or (when (scm/identifier? f)
           (match (-> f
                      (scm/rescope stx)
                      (lookup env @store-atom))
             [:primitive prim] (expand-primitive stx prim env store-atom)
             #_#_[:macro trans]
             (expand-macro trans datum env symtable-ref scopes)
             [:var] nil))
         (for [x (cons f xs)]
           (expand (scm/rescope x stx) env store-atom))))))

(defn- import-primitives+builtins [scopeholder env-atom store-atom]
  (let [scopes (scm/syntax-context scopeholder)
        additions (concat (for [p primitives]
                            [p
                             [:global p]
                             [:primitive p]])
                          (for [b builtins]
                            [b
                             [:global b]
                             [:var]]))]
    (swap! env-atom #(into % (for [[_sym bind denot] additions]
                               [bind denot])))
    (swap! store-atom #(into % (for [[sym bind _denot] additions]
                                 [[sym scopes] bind])))))

(defn expand-library [lib]
  (let [scopeholder (scm/new-scopeholder)
        env-atom (atom {})
        store-atom (atom {})]
    (update
     lib
     :library/declarations
     (fn [decls]
       (for [decl decls]
         (match decl
           [:import import-name]
           (do (if (= import-name '(primitives+builtins))
                 (import-primitives+builtins scopeholder env-atom store-atom)
                 (throw (Exception. "todo")))
               decl)
           [:begin & xs]
           (->> (for [x xs]
                  (expand (scm/rescope x scopeholder)
                          @env-atom store-atom))
                (cons :begin)
                vec)))))))

(defn library->datum [lib]
  (cons 'define-library
        (cons (:library/name lib)
              (for [d (:library/declarations lib)]
                (match d
                  [:begin & xs] (cons 'begin xs)
                  [:import & xs] (cons 'import xs)
                  [:export & xs] (cons 'export xs))))))

(defn library-syntax->datum [lib]
  (update lib :library/declarations
          (fn [decls]
            (for [decl decls]
              (match decl
                [:begin & xs] (->> xs
                                   (map scm/syntax->datum)
                                   (cons :begin)
                                   vec)
                :else decl)))))

(comment
  (-> "t.scm"
      slurp
      (scm/read-syntax :library? true)
      expand-library
      library-syntax->datum
      library->datum)
  (let [scopeholder (scm/new-scopeholder)
        env {[:global 'lambda] [:primitive 'lambda]}
        store-atom (atom {['lambda #{}] [:global 'lambda]})]
    (expand (-> (slurp "t.scm")
                (scm/read-syntax #_#_:library? true)
                (scm/rescope scopeholder))
            env store-atom)))