;; infer.sx — Hindley-Milner Algorithm W for Haskell-on-SX (Phase 4).
;;
;; Types: TVar, TCon, TArr, TApp, TTuple, TScheme
;; Substitution: apply, compose, restrict
;; Unification (with occurs check)
;; Instantiation + generalization (let-polymorphism)
;; Algorithm W for: literals, var, con, lambda, app, let, if, op, tuple, list

;; ─── Type constructors ────────────────────────────────────────────────────────

(define hk-tvar (fn (n) (list "TVar" n)))
(define hk-tcon (fn (s) (list "TCon" s)))
(define hk-tarr (fn (a b) (list "TArr" a b)))
(define hk-tapp (fn (a b) (list "TApp" a b)))
(define hk-ttuple (fn (ts) (list "TTuple" ts)))
(define hk-tscheme (fn (vs t) (list "TScheme" vs t)))

(define hk-tvar? (fn (t) (and (list? t) (not (empty? t)) (= (first t) "TVar"))))
(define hk-tcon? (fn (t) (and (list? t) (not (empty? t)) (= (first t) "TCon"))))
(define hk-tarr? (fn (t) (and (list? t) (not (empty? t)) (= (first t) "TArr"))))
(define hk-tapp? (fn (t) (and (list? t) (not (empty? t)) (= (first t) "TApp"))))
(define hk-ttuple? (fn (t) (and (list? t) (not (empty? t)) (= (first t) "TTuple"))))
(define hk-tscheme? (fn (t) (and (list? t) (not (empty? t)) (= (first t) "TScheme"))))

(define hk-tvar-name (fn (t) (nth t 1)))
(define hk-tcon-name (fn (t) (nth t 1)))
(define hk-tarr-t1 (fn (t) (nth t 1)))
(define hk-tarr-t2 (fn (t) (nth t 2)))
(define hk-tapp-t1 (fn (t) (nth t 1)))
(define hk-tapp-t2 (fn (t) (nth t 2)))
(define hk-ttuple-ts (fn (t) (nth t 1)))
(define hk-tscheme-vs (fn (t) (nth t 1)))
(define hk-tscheme-type (fn (t) (nth t 2)))

(define hk-t-int    (hk-tcon "Int"))
(define hk-t-bool   (hk-tcon "Bool"))
(define hk-t-string (hk-tcon "String"))
(define hk-t-char   (hk-tcon "Char"))
(define hk-t-float  (hk-tcon "Float"))
(define hk-t-list (fn (t) (hk-tapp (hk-tcon "[]") t)))

;; ─── Type formatter ──────────────────────────────────────────────────────────

(define
  hk-type->str
  (fn
    (t)
    (cond
      ((hk-tvar? t)   (hk-tvar-name t))
      ((hk-tcon? t)   (hk-tcon-name t))
      ((hk-tarr? t)
       (let ((s1 (if (hk-tarr? (hk-tarr-t1 t))
                   (str "(" (hk-type->str (hk-tarr-t1 t)) ")")
                   (hk-type->str (hk-tarr-t1 t)))))
         (str s1 " -> " (hk-type->str (hk-tarr-t2 t)))))
      ((hk-tapp? t)
       (let ((h (hk-tapp-t1 t)))
         (cond
           ((and (hk-tcon? h) (= (hk-tcon-name h) "[]"))
            (str "[" (hk-type->str (hk-tapp-t2 t)) "]"))
           (:else
            (str "(" (hk-type->str h) " " (hk-type->str (hk-tapp-t2 t)) ")")))))
      ((hk-ttuple? t)
       (str "(" (join ", " (map hk-type->str (hk-ttuple-ts t))) ")"))
      ((hk-tscheme? t)
       (str "forall " (join " " (hk-tscheme-vs t)) ". " (hk-type->str (hk-tscheme-type t))))
      (:else "<?>"))))

;; ─── Fresh variable counter ───────────────────────────────────────────────────

(define hk-fresh-ctr 0)
(define hk-fresh (fn () (set! hk-fresh-ctr (+ hk-fresh-ctr 1)) (hk-tvar (str "t" hk-fresh-ctr))))
(define hk-reset-fresh (fn () (set! hk-fresh-ctr 0)))

;; ─── Utilities ───────────────────────────────────────────────────────────────

(define hk-infer-member? (fn (x lst) (some (fn (y) (= x y)) lst)))

(define
  hk-nub
  (fn (lst)
    (reduce (fn (acc x) (if (hk-infer-member? x acc) acc (append acc (list x)))) (list) lst)))

;; ─── Free type variables ──────────────────────────────────────────────────────

(define
  hk-ftv
  (fn
    (t)
    (cond
      ((hk-tvar? t)    (list (hk-tvar-name t)))
      ((hk-tcon? t)    (list))
      ((hk-tarr? t)    (append (hk-ftv (hk-tarr-t1 t)) (hk-ftv (hk-tarr-t2 t))))
      ((hk-tapp? t)    (append (hk-ftv (hk-tapp-t1 t)) (hk-ftv (hk-tapp-t2 t))))
      ((hk-ttuple? t)  (reduce append (list) (map hk-ftv (hk-ttuple-ts t))))
      ((hk-tscheme? t)
       (filter
         (fn (v) (not (hk-infer-member? v (hk-tscheme-vs t))))
         (hk-ftv (hk-tscheme-type t))))
      (:else (list)))))

(define
  hk-ftv-env
  (fn (env)
    (reduce (fn (acc k) (append acc (hk-ftv (get env k)))) (list) (keys env))))

;; ─── Substitution ─────────────────────────────────────────────────────────────

(define hk-subst-empty (dict))

(define
  hk-subst-restrict
  (fn
    (s exclude)
    (let ((r (dict)))
      (for-each
        (fn (k)
          (when (not (hk-infer-member? k exclude))
            (dict-set! r k (get s k))))
        (keys s))
      r)))

(define
  hk-subst-apply
  (fn
    (s t)
    (cond
      ((hk-tvar? t)
       (let ((v (get s (hk-tvar-name t))))
         (if (nil? v) t (hk-subst-apply s v))))
      ((hk-tarr? t)
       (hk-tarr (hk-subst-apply s (hk-tarr-t1 t))
                (hk-subst-apply s (hk-tarr-t2 t))))
      ((hk-tapp? t)
       (hk-tapp (hk-subst-apply s (hk-tapp-t1 t))
                (hk-subst-apply s (hk-tapp-t2 t))))
      ((hk-ttuple? t)
       (hk-ttuple (map (fn (u) (hk-subst-apply s u)) (hk-ttuple-ts t))))
      ((hk-tscheme? t)
       (let ((s2 (hk-subst-restrict s (hk-tscheme-vs t))))
         (hk-tscheme (hk-tscheme-vs t)
                     (hk-subst-apply s2 (hk-tscheme-type t)))))
      (:else t))))

(define
  hk-subst-compose
  (fn
    (s2 s1)
    (let ((r (hk-dict-copy s2)))
      (for-each
        (fn (k)
          (when (nil? (get r k))
            (dict-set! r k (hk-subst-apply s2 (get s1 k)))))
        (keys s1))
      r)))

(define
  hk-env-apply-subst
  (fn
    (s env)
    (let ((r (dict)))
      (for-each (fn (k) (dict-set! r k (hk-subst-apply s (get env k)))) (keys env))
      r)))

;; ─── Unification ─────────────────────────────────────────────────────────────

(define
  hk-bind-var
  (fn
    (v t)
    (cond
      ((and (hk-tvar? t) (= (hk-tvar-name t) v))
       hk-subst-empty)
      ((hk-infer-member? v (hk-ftv t))
       (raise (str "Occurs check failed: " v " in " (hk-type->str t))))
      (:else
       (let ((s (dict)))
         (dict-set! s v t)
         s)))))

(define
  hk-zip-unify
  (fn
    (ts1 ts2 acc)
    (if (or (empty? ts1) (empty? ts2))
      acc
      (let ((s (hk-unify (hk-subst-apply acc (first ts1))
                         (hk-subst-apply acc (first ts2)))))
        (hk-zip-unify (rest ts1) (rest ts2) (hk-subst-compose s acc))))))

(define
  hk-unify
  (fn
    (t1 t2)
    (cond
      ((and (hk-tvar? t1) (hk-tvar? t2) (= (hk-tvar-name t1) (hk-tvar-name t2)))
       hk-subst-empty)
      ((hk-tvar? t1) (hk-bind-var (hk-tvar-name t1) t2))
      ((hk-tvar? t2) (hk-bind-var (hk-tvar-name t2) t1))
      ((and (hk-tcon? t1) (hk-tcon? t2) (= (hk-tcon-name t1) (hk-tcon-name t2)))
       hk-subst-empty)
      ((and (hk-tarr? t1) (hk-tarr? t2))
       (let ((s1 (hk-unify (hk-tarr-t1 t1) (hk-tarr-t1 t2))))
         (let ((s2 (hk-unify (hk-subst-apply s1 (hk-tarr-t2 t1))
                             (hk-subst-apply s1 (hk-tarr-t2 t2)))))
           (hk-subst-compose s2 s1))))
      ((and (hk-tapp? t1) (hk-tapp? t2))
       (let ((s1 (hk-unify (hk-tapp-t1 t1) (hk-tapp-t1 t2))))
         (let ((s2 (hk-unify (hk-subst-apply s1 (hk-tapp-t2 t1))
                             (hk-subst-apply s1 (hk-tapp-t2 t2)))))
           (hk-subst-compose s2 s1))))
      ((and (hk-ttuple? t1) (hk-ttuple? t2)
            (= (length (hk-ttuple-ts t1)) (length (hk-ttuple-ts t2))))
       (hk-zip-unify (hk-ttuple-ts t1) (hk-ttuple-ts t2) hk-subst-empty))
      (:else
       (raise (str "Cannot unify " (hk-type->str t1) " with " (hk-type->str t2)))))))

;; ─── Instantiation and generalization ────────────────────────────────────────

(define
  hk-instantiate
  (fn
    (t)
    (if (not (hk-tscheme? t))
      t
      (let ((s (dict)))
        (for-each (fn (v) (dict-set! s v (hk-fresh))) (hk-tscheme-vs t))
        (hk-subst-apply s (hk-tscheme-type t))))))

(define
  hk-generalize
  (fn
    (env t)
    (let ((free-t   (hk-nub (hk-ftv t)))
          (free-env (hk-nub (hk-ftv-env env))))
      (let ((bound (filter (fn (v) (not (hk-infer-member? v free-env))) free-t)))
        (if (empty? bound)
          t
          (hk-tscheme bound t))))))

;; ─── Pattern binding extraction ──────────────────────────────────────────────
;; Returns a dict of name → type bindings introduced by matching pat against tv.

(define
  hk-w-pat
  (fn
    (pat tv)
    (let ((tag (first pat)))
      (cond
        ((= tag "p-var")  (let ((d (dict))) (dict-set! d (nth pat 1) tv) d))
        ((= tag "p-wild") (dict))
        (:else            (dict))))))

;; ─── Algorithm W ─────────────────────────────────────────────────────────────
;; hk-w : env × expr → (list subst type)

(define
  hk-w-let
  (fn
    (env binds body)
    ;; Infer types for each binding in order, generalising at each step.
    (let
      ((env2
        (reduce
          (fn
            (cur-env b)
            (let ((tag (first b)))
              (cond
                ;; Simple pattern binding: let x = expr
                ((or (= tag "bind") (= tag "pat-bind"))
                 (let ((pat  (nth b 1))
                       (rhs  (nth b 2)))
                   (let ((tv (hk-fresh)))
                     (let ((r (hk-w cur-env rhs)))
                       (let ((s1 (first r)) (t1 (nth r 1)))
                         (let ((s2 (hk-unify (hk-subst-apply s1 tv) t1)))
                           (let ((s (hk-subst-compose s2 s1)))
                             (let ((t-gen (hk-generalize (hk-env-apply-subst s cur-env)
                                                         (hk-subst-apply s t1))))
                               (let ((bindings (hk-w-pat pat t-gen)))
                                 (let ((r2 (hk-dict-copy cur-env)))
                                   (for-each
                                     (fn (k) (dict-set! r2 k (get bindings k)))
                                     (keys bindings))
                                   r2))))))))))
                ;; Function clause: let f x y = expr
                ((= tag "fun-clause")
                 (let ((name (nth b 1))
                       (pats (nth b 2))
                       (body2 (nth b 3)))
                   ;; Treat as: let name = lambda pats body2
                   (let ((rhs (if (empty? pats)
                                body2
                                (list "lambda" pats body2))))
                     (let ((tv (hk-fresh)))
                       (let ((env-rec (hk-dict-copy cur-env)))
                         (dict-set! env-rec name tv)
                         (let ((r (hk-w env-rec rhs)))
                           (let ((s1 (first r)) (t1 (nth r 1)))
                             (let ((s2 (hk-unify (hk-subst-apply s1 tv) t1)))
                               (let ((s (hk-subst-compose s2 s1)))
                                 (let ((t-gen (hk-generalize
                                                (hk-env-apply-subst s cur-env)
                                                (hk-subst-apply s t1))))
                                   (let ((r2 (hk-dict-copy cur-env)))
                                     (dict-set! r2 name t-gen)
                                     r2)))))))))))
                (:else cur-env))))
          env
          binds)))
      (hk-w env2 body))))

(define
  hk-w
  (fn
    (env expr)
    (let ((tag (first expr)))
      (cond
        ;; Literals
        ((= tag "int")    (list hk-subst-empty hk-t-int))
        ((= tag "float")  (list hk-subst-empty hk-t-float))
        ((= tag "string") (list hk-subst-empty hk-t-string))
        ((= tag "char")   (list hk-subst-empty hk-t-char))

        ;; Variable
        ((= tag "var")
         (let ((name (nth expr 1)))
           (let ((scheme (get env name)))
             (if (nil? scheme)
               (raise (str "Unbound variable: " name))
               (list hk-subst-empty (hk-instantiate scheme))))))

        ;; Constructor (same lookup as var)
        ((= tag "con")
         (let ((name (nth expr 1)))
           (let ((scheme (get env name)))
             (if (nil? scheme)
               (list hk-subst-empty (hk-fresh))
               (list hk-subst-empty (hk-instantiate scheme))))))

        ;; Unary negation
        ((= tag "neg")
         (let ((r (hk-w env (nth expr 1))))
           (let ((s1 (first r)) (t1 (nth r 1)))
             (let ((s2 (hk-unify t1 hk-t-int)))
               (list (hk-subst-compose s2 s1) hk-t-int)))))

        ;; Lambda: ("lambda" pats body)
        ((= tag "lambda")
         (let ((pats (nth expr 1))
               (body (nth expr 2)))
           (if (empty? pats)
             (hk-w env body)
             (let ((pat  (first pats))
                   (rest (rest pats)))
               (let ((tv (hk-fresh)))
                 (let ((bindings (hk-w-pat pat tv)))
                   (let ((env2 (hk-dict-copy env)))
                     (for-each (fn (k) (dict-set! env2 k (get bindings k))) (keys bindings))
                     (let ((inner (if (empty? rest)
                                    body
                                    (list "lambda" rest body))))
                       (let ((r (hk-w env2 inner)))
                         (let ((s1 (first r)) (t1 (nth r 1)))
                           (list s1 (hk-tarr (hk-subst-apply s1 tv) t1))))))))))))

        ;; Application: ("app" f x)
        ((= tag "app")
         (let ((tv (hk-fresh)))
           (let ((r1 (hk-w env (nth expr 1))))
             (let ((s1 (first r1)) (tf (nth r1 1)))
               (let ((r2 (hk-w (hk-env-apply-subst s1 env) (nth expr 2))))
                 (let ((s2 (first r2)) (tx (nth r2 1)))
                   (let ((s3 (hk-unify (hk-subst-apply s2 tf) (hk-tarr tx tv))))
                     (let ((s (hk-subst-compose s3 (hk-subst-compose s2 s1))))
                       (list s (hk-subst-apply s3 tv))))))))))

        ;; Let: ("let" binds body)
        ((= tag "let")
         (hk-w-let env (nth expr 1) (nth expr 2)))

        ;; If: ("if" cond then else)
        ((= tag "if")
         (let ((r1 (hk-w env (nth expr 1))))
           (let ((s1 (first r1)) (tc (nth r1 1)))
             (let ((s2 (hk-unify tc hk-t-bool)))
               (let ((s12 (hk-subst-compose s2 s1)))
                 (let ((r2 (hk-w (hk-env-apply-subst s12 env) (nth expr 2))))
                   (let ((s3 (first r2)) (tt (nth r2 1)))
                     (let ((s123 (hk-subst-compose s3 s12)))
                       (let ((r3 (hk-w (hk-env-apply-subst s123 env) (nth expr 3))))
                         (let ((s4 (first r3)) (te (nth r3 1)))
                           (let ((s5 (hk-unify (hk-subst-apply s4 tt) te)))
                             (let ((s (hk-subst-compose s5 (hk-subst-compose s4 s123))))
                               (list s (hk-subst-apply s5 te))))))))))))))

        ;; Binary operator: ("op" op-name left right)
        ;; Desugar to double application.
        ((= tag "op")
         (hk-w env
               (list "app"
                     (list "app" (list "var" (nth expr 1)) (nth expr 2))
                     (nth expr 3))))

        ;; Tuple: ("tuple" [e1 e2 ...])
        ((= tag "tuple")
         (let ((elems (nth expr 1)))
           (let ((s-acc hk-subst-empty)
                 (ts    (list)))
             (for-each
               (fn (e)
                 (let ((r (hk-w (hk-env-apply-subst s-acc env) e)))
                   (set! s-acc (hk-subst-compose (first r) s-acc))
                   (set! ts    (append ts (list (nth r 1))))))
               elems)
             (list s-acc (hk-ttuple (map (fn (t) (hk-subst-apply s-acc t)) ts))))))

        ;; List literal: ("list" [e1 e2 ...])
        ((= tag "list")
         (let ((elems (nth expr 1)))
           (if (empty? elems)
             (list hk-subst-empty (hk-t-list (hk-fresh)))
             (let ((tv (hk-fresh)))
               (let ((s-acc hk-subst-empty))
                 (for-each
                   (fn (e)
                     (let ((r (hk-w (hk-env-apply-subst s-acc env) e)))
                       (let ((s2 (first r)) (te (nth r 1)))
                         (let ((s3 (hk-unify (hk-subst-apply s2 tv) te)))
                           (set! s-acc (hk-subst-compose s3 (hk-subst-compose s2 s-acc)))))))
                   elems)
                 (list s-acc (hk-t-list (hk-subst-apply s-acc tv))))))))

        ;; Location annotation: just delegate — position is for outer context.
        ((= tag "loc")
         (hk-w env (nth expr 3)))

        (:else
         (raise (str "hk-w: unhandled tag: " tag)))))))

;; ─── Initial type environment ─────────────────────────────────────────────────
;; Monomorphic numeric ops (no Num typeclass yet — upgraded in Phase 5).

(define
  hk-type-env0
  (fn ()
    (let ((env (dict)))
      ;; Integer arithmetic
      (for-each
        (fn (op)
          (dict-set! env op (hk-tarr hk-t-int (hk-tarr hk-t-int hk-t-int))))
        (list "+" "-" "*" "div" "mod" "quot" "rem"))
      ;; Integer comparison → Bool
      (for-each
        (fn (op)
          (dict-set! env op (hk-tarr hk-t-int (hk-tarr hk-t-int hk-t-bool))))
        (list "==" "/=" "<" "<=" ">" ">="))
      ;; Boolean operators
      (dict-set! env "&&" (hk-tarr hk-t-bool (hk-tarr hk-t-bool hk-t-bool)))
      (dict-set! env "||" (hk-tarr hk-t-bool (hk-tarr hk-t-bool hk-t-bool)))
      (dict-set! env "not" (hk-tarr hk-t-bool hk-t-bool))
      ;; Constructors
      (dict-set! env "True"  hk-t-bool)
      (dict-set! env "False" hk-t-bool)
      ;; Polymorphic list ops (using TScheme)
      (let ((a (hk-tvar "a")))
        (dict-set! env "head"  (hk-tscheme (list "a") (hk-tarr (hk-t-list a) a)))
        (dict-set! env "tail"  (hk-tscheme (list "a") (hk-tarr (hk-t-list a) (hk-t-list a))))
        (dict-set! env "null"  (hk-tscheme (list "a") (hk-tarr (hk-t-list a) hk-t-bool)))
        (dict-set! env "length" (hk-tscheme (list "a") (hk-tarr (hk-t-list a) hk-t-int)))
        (dict-set! env "reverse" (hk-tscheme (list "a") (hk-tarr (hk-t-list a) (hk-t-list a))))
        (dict-set! env ":"
          (hk-tscheme (list "a") (hk-tarr a (hk-tarr (hk-t-list a) (hk-t-list a))))))
      ;; negate
      (dict-set! env "negate" (hk-tarr hk-t-int hk-t-int))
      (dict-set! env "abs"    (hk-tarr hk-t-int hk-t-int))
      env)))

;; ─── Expression brief printer ────────────────────────────────────────────────
;; Produces a short human-readable label for an AST node used in error messages.

(define
  hk-expr->brief
  (fn
    (expr)
    (cond
      ((not (list? expr)) (str expr))
      ((empty? expr)      "()")
      (:else
       (let ((tag (first expr)))
         (cond
           ((= tag "var")    (nth expr 1))
           ((= tag "con")    (nth expr 1))
           ((= tag "int")    (str (nth expr 1)))
           ((= tag "float")  (str (nth expr 1)))
           ((= tag "string") (str "\"" (nth expr 1) "\""))
           ((= tag "char")   (str "'" (nth expr 1) "'"))
           ((= tag "neg")    (str "(-" (hk-expr->brief (nth expr 1)) ")"))
           ((= tag "app")
            (str "(" (hk-expr->brief (nth expr 1))
                 " " (hk-expr->brief (nth expr 2)) ")"))
           ((= tag "op")
            (str "(" (hk-expr->brief (nth expr 2))
                 " " (nth expr 1)
                 " " (hk-expr->brief (nth expr 3)) ")"))
           ((= tag "lambda") "(\\ ...)")
           ((= tag "let")    "(let ...)")
           ((= tag "if")     "(if ...)")
           ((= tag "tuple")  "(tuple ...)")
           ((= tag "list")   "[...]")
           ((= tag "loc")    (hk-expr->brief (nth expr 3)))
           (:else            (str "(" tag " ..."))))))))

;; ─── Loc-annotated inference ──────────────────────────────────────────────────
;; ("loc" LINE COL INNER) node: hk-w catches any error and re-raises with
;; "at LINE:COL: " prepended.  Emitted by the parser or test scaffolding.

;; Extended hk-w handles "loc" — handled inline in the cond below.

;; ─── Program-level inference ─────────────────────────────────────────────────
;; hk-infer-decl : env × decl → ("ok" name type-str) | ("err" msg) | nil
;;   Uses tagged results so callers don't need re-raise.

(define
  hk-infer-decl
  (fn
    (env decl)
    (let
      ((tag (first decl)))
      (cond
        ((= tag "fun-clause")
          (let
            ((name (nth decl 1)) (pats (nth decl 2)) (body (nth decl 3)))
            (let
              ((rhs (if (empty? pats) body (list "lambda" pats body))))
              (guard
                (e (#t (list "err" (str "in '" name "': " e))))
                (begin
                  (hk-reset-fresh)
                  (let
                    ((r (hk-w env rhs)))
                    (let
                      ((final-type (hk-subst-apply (first r) (nth r 1))))
                      (list "ok" name (hk-type->str final-type) final-type))))))))
        ((or (= tag "bind") (= tag "pat-bind"))
          (let
            ((pat (nth decl 1)) (body (nth decl 2)))
            (let
              ((label (if (and (list? pat) (= (first pat) "p-var")) (nth pat 1) "<binding>")))
              (guard
                (e (#t (list "err" (str "in '" label "': " e))))
                (begin
                  (hk-reset-fresh)
                  (let
                    ((r (hk-w env body)))
                    (let
                      ((final-type (hk-subst-apply (first r) (nth r 1))))
                      (list "ok" label (hk-type->str final-type) final-type))))))))
        (:else nil)))))

;; hk-infer-prog : program-ast × env → list of ("ok" name type) | ("err" msg)

(define
  hk-ast-type
  (fn
    (ast)
    (let
      ((tag (first ast)))
      (cond
        ((= tag "t-con") (list "TCon" (nth ast 1)))
        ((= tag "t-var") (list "TVar" (nth ast 1)))
        ((= tag "t-fun")
          (list "TArr" (hk-ast-type (nth ast 1)) (hk-ast-type (nth ast 2))))
        ((= tag "t-app")
          (list "TApp" (hk-ast-type (nth ast 1)) (hk-ast-type (nth ast 2))))
        ((= tag "t-list")
          (list "TApp" (list "TCon" "[]") (hk-ast-type (nth ast 1))))
        ((= tag "t-tuple") (list "TTuple" (map hk-ast-type (nth ast 1))))
        (:else (raise (str "unknown type node: " (first ast))))))))

;; ─── Convenience ─────────────────────────────────────────────────────────────
;; hk-infer-type : Haskell expression source → inferred type string

(define
  hk-collect-tvars
  (fn
    (t acc)
    (cond
      ((= (first t) "TVar")
        (if
          (some (fn (v) (= v (nth t 1))) acc)
          acc
          (begin (append! acc (nth t 1)) acc)))
      ((= (first t) "TArr")
        (hk-collect-tvars (nth t 2) (hk-collect-tvars (nth t 1) acc)))
      ((= (first t) "TApp")
        (hk-collect-tvars (nth t 2) (hk-collect-tvars (nth t 1) acc)))
      ((= (first t) "TTuple")
        (reduce (fn (a elem) (hk-collect-tvars elem a)) acc (nth t 1)))
      (:else acc))))

(define
  hk-check-sig
  (fn
    (declared-ast inferred-type)
    (let
      ((declared (hk-ast-type declared-ast)))
      (let
        ((tvars (hk-collect-tvars declared (list))))
        (let
          ((scheme (if (empty? tvars) declared (list "TScheme" tvars declared))))
          (let
            ((inst (hk-instantiate scheme)))
            (hk-unify inst inferred-type)))))))

(define
  hk-infer-prog
  (fn
    (prog env)
    (let
      ((decls (cond ((and (list? prog) (= (first prog) "program")) (nth prog 1)) ((and (list? prog) (= (first prog) "module")) (nth prog 3)) (:else (list))))
        (results (list))
        (sigs (dict)))
      (for-each
        (fn
          (d)
          (when
            (= (first d) "type-sig")
            (let
              ((names (nth d 1)) (type-ast (nth d 2)))
              (for-each (fn (n) (dict-set! sigs n type-ast)) names))))
        decls)
      (for-each
        (fn
          (d)
          (let
            ((r (hk-infer-decl env d)))
            (when
              (not (nil? r))
              (let
                ((checked (if (and (= (first r) "ok") (has-key? sigs (nth r 1))) (guard (e (true (list "err" (str "in '" (nth r 1) "': declared type mismatch: " e)))) (begin (hk-check-sig (get sigs (nth r 1)) (nth r 3)) r)) r)))
                (append! results checked)
                (when
                  (= (first checked) "ok")
                  (dict-set! env (nth checked 1) (nth checked 3)))))))
        decls)
      results)))

(define
  hk-infer-type
  (fn
    (src)
    (hk-reset-fresh)
    (let
      ((ast (hk-core-expr src)) (env (hk-type-env0)))
      (let
        ((r (hk-w env ast)))
        (hk-type->str (hk-subst-apply (first r) (nth r 1)))))))