(ns selfhost.types
  (:require [selfhost.rt :as rt]
            [clojure.string :as str]))

(defn ^String char-at [^String s i]
  (if (and (>= i 0) (< i (count s))) (subs s i (+ i 1)) ""))

(defn ^String substring2 [^String s a b]
  (let [n (count s)
   lo (if (< a 0) 0 (if (> a n) n a))
   hi (if (< b lo) lo (if (> b n) n b))]
  (subs s lo hi)))

(defn index-of2 [xs v]
  (let [n (count xs)]
  (loop [i 0]
  (if (>= i n) -1 (if (= (nth xs i) v) i (recur (+ i 1)))))))

(defn obj-set! [obj k v]
  (swap! obj assoc k v))

(def PRIMITIVES ["String" "Int" "Float" "Bool" "Keyword" "Symbol" "Nil" "Any"])

(def CLJ-ALIASES {"Long" "Int" "Double" "Float" "Boolean" "Bool" "Integer" "Int"})

(def PARAMETRIC-CTORS ["Vec" "List" "Set" "Map" "Promise"])

(defn make-prim [^String name]
  {"kind" "prim" "name" name})

(defn make-fn [params rest-type ret]
  {"kind" "fn" "params" params "rest" rest-type "ret" ret})

(defn make-app [^String ctor args]
  {"kind" "app" "name" ctor "args" args})

(defn make-union [members]
  {"kind" "union" "members" members})

(defn make-var [^String name]
  {"kind" "var" "name" name})

(defn make-poly [vars body bounds]
  {"kind" "poly" "vars" vars "body" body "bounds" bounds})

(defn ^Boolean prim? [t]
  (and (not (nil? t)) (not= (get t "kind") nil) (= (get t "kind") "prim")))

(defn ^Boolean fn-type? [t]
  (and (not (nil? t)) (= (get t "kind") "fn")))

(defn ^Boolean app-type? [t]
  (and (not (nil? t)) (= (get t "kind") "app")))

(defn ^Boolean union-type? [t]
  (and (not (nil? t)) (= (get t "kind") "union")))

(defn ^Boolean var-type? [t]
  (and (not (nil? t)) (= (get t "kind") "var")))

(defn ^Boolean poly-type? [t]
  (and (not (nil? t)) (= (get t "kind") "poly")))

(defn ^Boolean any-type? [t]
  (and (prim? t) (= (get t "name") "Any")))

(defn ^String type->string [t]
  (cond
  (nil? t) "?"
  (prim? t) (get t "name")
  (fn-type? t) (let [params (get t "params")
   rest-t (get t "rest")
   ret (get t "ret")
   param-strs (mapv type->string params)]
  (if (not (nil? rest-t)) (str "[" (str/join " " param-strs) " & " (type->string rest-t) " -> " (type->string ret) "]") (str "[" (str/join " " param-strs) " -> " (type->string ret) "]")))
  (app-type? t) (let [ctor (get t "name")
   args (get t "args")
   arg-strs (mapv type->string args)]
  (str "(" ctor " " (str/join " " arg-strs) ")"))
  (union-type? t) (let [members (get t "members")
   all-prim (every? prim? members)
   names (if all-prim (mapv (fn [m] (get m "name")) members) [])]
  (cond
  (and all-prim (= (count members) 2) (>= (index-of2 names "Int") 0) (>= (index-of2 names "Float") 0)) "Number"
  (and (= (count members) 2) (some (fn [m] (and (prim? m) (= (get m "name") "Nil"))) members)) (let [non-nil (first (filter (fn [m] (not (and (prim? m) (= (get m "name") "Nil")))) members))]
  (str (type->string non-nil) "?"))
  :else (str "(U " (str/join " " (mapv type->string members)) ")")))
  (var-type? t) (get t "name")
  (poly-type? t) (let [vars (get t "vars")
   body (get t "body")
   bounds (get t "bounds")
   var-strs (mapv (fn [v] (let [b (if (nil? bounds) nil (get bounds v))]
  (if (nil? b) v (str "(" v " <: " (type->string b) ")")))) vars)]
  (str "(forall [" (str/join " " var-strs) "] " (type->string body) ")"))
  :else "?"))

(defn ^String unqualify-name [^String name]
  (let [idx (str/index-of name "/")]
  (if (nil? idx) name (substring2 name (+ idx 1) (count name)))))

(defn ^Boolean type-compatible? [actual expected]
  (cond
  (or (nil? actual) (nil? expected)) true
  (any-type? actual) true
  (any-type? expected) true
  (var-type? actual) true
  (var-type? expected) true
  (poly-type? expected) (type-compatible? actual (get expected "body"))
  (poly-type? actual) (type-compatible? (get actual "body") expected)
  (and (union-type? actual) (union-type? expected)) (every? (fn [a-alt] (some (fn [e-alt] (type-compatible? a-alt e-alt)) (get expected "members"))) (get actual "members"))
  (union-type? expected) (some (fn [alt] (type-compatible? actual alt)) (get expected "members"))
  (union-type? actual) (every? (fn [alt] (type-compatible? alt expected)) (get actual "members"))
  (and (prim? actual) (prim? expected)) (or (= (get actual "name") (get expected "name")) (= (unqualify-name (get actual "name")) (unqualify-name (get expected "name"))))
  (and (fn-type? actual) (fn-type? expected)) (let [ap (get actual "params")
   ep (get expected "params")
   ar (get actual "rest")
   er (get expected "rest")]
  (and (= (count ap) (count ep)) (every? identity (map-indexed (fn [i p] (type-compatible? p (nth ep i))) ap)) (= (nil? ar) (nil? er)) (or (nil? ar) (type-compatible? ar er)) (type-compatible? (get actual "ret") (get expected "ret"))))
  (and (app-type? actual) (app-type? expected)) (and (= (get actual "name") (get expected "name")) (= (count (get actual "args")) (count (get expected "args"))) (every? identity (map-indexed (fn [i a] (type-compatible? a (nth (get expected "args") i))) (get actual "args"))))
  :else false))

(def user-parametric {})

(declare parse-type)

(defn parse-fn-type [items]
  (let [arrow-idx (index-of2 items "->")]
  (if (= arrow-idx -1) (make-prim "Any") (let [before (subvec items 0 arrow-idx)
   after (subvec items (+ arrow-idx 1))]
  (if (not= (count after) 1) (make-prim "Any") (let [amp-pos (index-of2 before "&")]
  (if (> amp-pos -1) (make-fn (mapv parse-type (subvec before 0 amp-pos)) (parse-type (nth (subvec before (+ amp-pos 1) (count before)) 0)) (parse-type (nth after 0))) (make-fn (mapv parse-type before) nil (parse-type (nth after 0))))))))))

(defn parse-type [t]
  (cond
  (and (vector? t) (> (count t) 0) (= (nth t 0) "#%brackets")) (parse-fn-type (subvec t 1))
  (and (vector? t) (= (count t) 3) (= (nth t 0) "forall")) (let [vars-form (nth t 1)
   raw-vars (if (and (vector? vars-form) (> (count vars-form) 0) (= (nth vars-form 0) "#%brackets")) (subvec vars-form 1) vars-form)
   vars (vec (filter string? raw-vars))]
  (make-poly vars (parse-type (nth t 2)) nil))
  (and (vector? t) (> (count t) 1) (= (nth t 0) "U")) (make-union (mapv parse-type (subvec t 1)))
  (and (vector? t) (> (count t) 0) (string? (nth t 0)) (or (>= (index-of2 PARAMETRIC-CTORS (nth t 0)) 0) (= (get user-parametric (nth t 0)) true))) (make-app (nth t 0) (mapv parse-type (subvec t 1)))
  (and (string? t) (> (count t) 1) (= (char-at t (- (count t) 1)) "?")) (let [base (substring2 t 0 (- (count t) 1))]
  (make-union [(parse-type base) (make-prim "Nil")]))
  (and (string? t) (= t "Number")) (make-union [(make-prim "Int") (make-prim "Float")])
  (and (string? t) (not (nil? (get CLJ-ALIASES t)))) (make-prim (get CLJ-ALIASES t))
  (string? t) (make-prim t)
  :else (make-prim "Any")))

(defn infer-literal-type [e]
  (let [kind (get e "kind")]
  (cond
  (= kind "string") (make-prim "String")
  (= kind "bool") (make-prim "Bool")
  (= kind "number") (make-prim "Int")
  (= kind "float") (make-prim "Float")
  (= kind "nil") (make-prim "Nil")
  (= kind "keyword") (make-prim "Keyword")
  (= kind "symbol") (make-prim "Symbol")
  :else nil)))

(defn infer-type-var-bindings! [expected actual bindings]
  (cond
  (or (nil? expected) (nil? actual)) nil
  (any-type? actual) nil
  (var-type? expected) (do
  (if (nil? (get (deref bindings) (get expected "name"))) (do
  (obj-set! bindings (get expected "name") actual)))
  nil)
  (and (fn-type? expected) (fn-type? actual)) (do
  (if (= (count (get expected "params")) (count (get actual "params"))) (do
  (doseq [i (range (count (get expected "params")))]
  (infer-type-var-bindings! (nth (get expected "params") i) (nth (get actual "params") i) bindings))))
  (if (and (not (nil? (get expected "rest"))) (not (nil? (get actual "rest")))) (do
  (infer-type-var-bindings! (get expected "rest") (get actual "rest") bindings)))
  (infer-type-var-bindings! (get expected "ret") (get actual "ret") bindings)
  nil)
  (and (app-type? expected) (app-type? actual) (= (get expected "name") (get actual "name"))) (do
  (doseq [i (range (count (get expected "args")))]
  (infer-type-var-bindings! (nth (get expected "args") i) (nth (get actual "args") i) bindings))
  nil)
  :else nil))

(defn apply-type-bindings [t bindings]
  (cond
  (nil? t) nil
  (var-type? t) (let [bound (get bindings (get t "name"))]
  (if (nil? bound) (make-prim "Any") bound))
  (prim? t) t
  (fn-type? t) (make-fn (mapv (fn [p] (apply-type-bindings p bindings)) (get t "params")) (if (nil? (get t "rest")) nil (apply-type-bindings (get t "rest") bindings)) (apply-type-bindings (get t "ret") bindings))
  (app-type? t) (make-app (get t "name") (mapv (fn [a] (apply-type-bindings a bindings)) (get t "args")))
  (union-type? t) (make-union (mapv (fn [m] (apply-type-bindings m bindings)) (get t "members")))
  (poly-type? t) t
  :else t))

(def passes (atom []))

(def failures (atom []))

(defn- expect! [^String label ^Boolean result]
  (if result (do
  (swap! passes conj true)
  nil) (do
  (swap! failures conj label)
  nil)))

(defn run-tests! []
  (reset! passes [])
  (reset! failures [])
  (expect! "ts: prim String" (= (type->string (make-prim "String")) "String"))
  (expect! "ts: prim Int" (= (type->string (make-prim "Int")) "Int"))
  (expect! "ts: prim Any" (= (type->string (make-prim "Any")) "Any"))
  (expect! "ts: prim Nil" (= (type->string (make-prim "Nil")) "Nil"))
  (expect! "ts: fn [Int -> String]" (= (type->string (make-fn [(make-prim "Int")] nil (make-prim "String"))) "[Int -> String]"))
  (expect! "ts: fn [Int Bool -> String]" (= (type->string (make-fn [(make-prim "Int") (make-prim "Bool")] nil (make-prim "String"))) "[Int Bool -> String]"))
  (expect! "ts: fn variadic [Int & String -> Bool]" (= (type->string (make-fn [(make-prim "Int")] (make-prim "String") (make-prim "Bool"))) "[Int & String -> Bool]"))
  (expect! "ts: app (Vec Int)" (= (type->string (make-app "Vec" [(make-prim "Int")])) "(Vec Int)"))
  (expect! "ts: app (Map String Int)" (= (type->string (make-app "Map" [(make-prim "String") (make-prim "Int")])) "(Map String Int)"))
  (expect! "ts: union nullable String?" (= (type->string (make-union [(make-prim "String") (make-prim "Nil")])) "String?"))
  (expect! "ts: union Number" (= (type->string (make-union [(make-prim "Int") (make-prim "Float")])) "Number"))
  (expect! "ts: union (U String Int Nil)" (= (type->string (make-union [(make-prim "String") (make-prim "Int") (make-prim "Nil")])) "(U String Int Nil)"))
  (expect! "ts: var T" (= (type->string (make-var "T")) "T"))
  (expect! "ts: poly (forall [T] (Vec T))" (= (type->string (make-poly ["T"] (make-app "Vec" [(make-var "T")]) nil)) "(forall [T] (Vec T))"))
  (expect! "ts: null type" (= (type->string nil) "?"))
  (expect! "tc: same prim" (type-compatible? (make-prim "String") (make-prim "String")))
  (expect! "tc: diff prim" (not (type-compatible? (make-prim "String") (make-prim "Int"))))
  (expect! "tc: actual Any" (type-compatible? (make-prim "Any") (make-prim "Int")))
  (expect! "tc: expected Any" (type-compatible? (make-prim "Int") (make-prim "Any")))
  (expect! "tc: actual var" (type-compatible? (make-var "T") (make-prim "Int")))
  (expect! "tc: expected var" (type-compatible? (make-prim "Int") (make-var "T")))
  (expect! "tc: null actual" (type-compatible? nil (make-prim "Int")))
  (expect! "tc: null expected" (type-compatible? (make-prim "Int") nil))
  (expect! "tc: prim in union" (type-compatible? (make-prim "String") (make-union [(make-prim "String") (make-prim "Nil")])))
  (expect! "tc: prim not in union" (not (type-compatible? (make-prim "Int") (make-union [(make-prim "String") (make-prim "Nil")]))))
  (expect! "tc: union subset" (type-compatible? (make-union [(make-prim "String")]) (make-union [(make-prim "String") (make-prim "Nil")])))
  (expect! "tc: union not subset" (not (type-compatible? (make-union [(make-prim "String") (make-prim "Int")]) (make-union [(make-prim "String") (make-prim "Nil")]))))
  (expect! "tc: same fn" (type-compatible? (make-fn [(make-prim "Int")] nil (make-prim "String")) (make-fn [(make-prim "Int")] nil (make-prim "String"))))
  (expect! "tc: diff fn param" (not (type-compatible? (make-fn [(make-prim "Int")] nil (make-prim "String")) (make-fn [(make-prim "Bool")] nil (make-prim "String")))))
  (expect! "tc: diff fn arity" (not (type-compatible? (make-fn [(make-prim "Int")] nil (make-prim "String")) (make-fn [(make-prim "Int") (make-prim "Bool")] nil (make-prim "String")))))
  (expect! "tc: same app" (type-compatible? (make-app "Vec" [(make-prim "Int")]) (make-app "Vec" [(make-prim "Int")])))
  (expect! "tc: diff app arg" (not (type-compatible? (make-app "Vec" [(make-prim "Int")]) (make-app "Vec" [(make-prim "String")]))))
  (expect! "tc: diff app ctor" (not (type-compatible? (make-app "Vec" [(make-prim "Int")]) (make-app "Set" [(make-prim "Int")]))))
  (expect! "tc: qualified name" (type-compatible? (make-prim "mymod/Type") (make-prim "Type")))
  (expect! "tc: poly unwrap" (type-compatible? (make-prim "String") (make-poly ["T"] (make-prim "String") nil)))
  (expect! "pt: prim" (= (parse-type "String") (make-prim "String")))
  (expect! "pt: app Vec" (= (parse-type ["Vec" "String"]) (make-app "Vec" [(make-prim "String")])))
  (expect! "pt: union" (= (parse-type ["U" "String" "Nil"]) (make-union [(make-prim "String") (make-prim "Nil")])))
  (expect! "pt: nullable sugar" (= (parse-type "String?") (make-union [(make-prim "String") (make-prim "Nil")])))
  (expect! "pt: Number alias" (= (parse-type "Number") (make-union [(make-prim "Int") (make-prim "Float")])))
  (expect! "pt: CLJ alias Long" (= (parse-type "Long") (make-prim "Int")))
  (expect! "pt: fn type" (= (parse-type ["#%brackets" "Int" "->" "String"]) (make-fn [(make-prim "Int")] nil (make-prim "String"))))
  (expect! "pt: variadic fn" (= (parse-type ["#%brackets" "Int" "&" "String" "->" "Bool"]) (make-fn [(make-prim "Int")] (make-prim "String") (make-prim "Bool"))))
  (expect! "pt: nested (Vec (Map String Int))" (= (parse-type ["Vec" ["Map" "String" "Int"]]) (make-app "Vec" [(make-app "Map" [(make-prim "String") (make-prim "Int")])])))
  (expect! "lit: string" (= (infer-literal-type {"kind" "string" "value" "hi"}) (make-prim "String")))
  (expect! "lit: int" (= (infer-literal-type {"kind" "number" "value" 42}) (make-prim "Int")))
  (expect! "lit: float" (= (infer-literal-type {"kind" "float" "value" 3.14}) (make-prim "Float")))
  (expect! "lit: bool" (= (infer-literal-type {"kind" "bool" "value" true}) (make-prim "Bool")))
  (expect! "lit: nil" (= (infer-literal-type {"kind" "nil"}) (make-prim "Nil")))
  (expect! "lit: keyword" (= (infer-literal-type {"kind" "keyword" "value" ":foo"}) (make-prim "Keyword")))
  (doseq [f (deref failures)]
  (selfhost.rt/eprint (str "  FAIL: " f "\n")))
  (println (str "  TYPES: " (count (deref passes)) " passed, " (count (deref failures)) " failed"))
  (count (deref failures)))
