;; Haskell strict evaluator (Phase 2).
;;
;; Consumes the post-desugar core AST and produces SX values. Strict
;; throughout — laziness and thunks are Phase 3.
;;
;; Value representation:
;;   numbers / strings / chars      → raw SX values
;;   constructor values             → tagged lists (con-name first)
;;   functions: closure / multifun  → {:type "fn" :kind … …}
;;   constructor partials           → {:type "con-partial" …}
;;   built-ins                      → {:type "builtin" …}
;;
;; Multi-clause top-level definitions are bundled into a single
;; multifun keyed by name; arguments are gathered through currying
;; until arity is reached, then each clause's pattern list is matched
;; in order. Recursive let bindings work because the binding env is
;; built mutably so closures captured during evaluation see the
;; eventual full env.

(define
  hk-dict-copy
  (fn
    (d)
    (let ((nd (dict)))
      (for-each
        (fn (k) (dict-set! nd k (get d k)))
        (keys d))
      nd)))

;; ── Thunks (Phase 3 — laziness) ─────────────────────────────
;; A thunk wraps an unevaluated AST plus the env in which it was
;; created. The first call to `hk-force` evaluates the body, replaces
;; the body with the cached value, and flips `forced`. Subsequent
;; forces return the cached value directly.
(define
  hk-mk-thunk
  (fn
    (body env)
    {:type "thunk" :body body :env env :forced false :value nil}))

(define
  hk-is-thunk?
  (fn (v) (and (dict? v) (= (get v "type") "thunk"))))

(define
  hk-force
  (fn
    (v)
    (cond
      ((hk-is-thunk? v)
        (cond
          ((get v "forced") (get v "value"))
          (:else
            (let
              ((res (hk-force (hk-eval (get v "body") (get v "env")))))
              (dict-set! v "forced" true)
              (dict-set! v "value" res)
              res))))
      (:else v))))

;; Recursive force — used at the test/output boundary so test
;; expectations can compare against fully-evaluated structures.
(define
  hk-deep-force
  (fn
    (v)
    (let ((fv (hk-force v)))
      (cond
        ((not (list? fv)) fv)
        ((empty? fv) fv)
        (:else (map hk-deep-force fv))))))

;; ── Function value constructors ──────────────────────────────
(define
  hk-mk-closure
  (fn
    (params body env)
    {:type "fn" :kind "closure" :params params :body body :env env}))

(define
  hk-mk-multifun
  (fn
    (arity clauses env)
    {:type "fn" :kind "multi" :arity arity :clauses clauses :env env :collected (list)}))

(define
  hk-mk-builtin
  (fn
    (name fn arity)
    {:type "builtin" :name name :fn fn :arity arity :lazy false :collected (list)}))

;; A lazy built-in receives its collected args as raw thunks (or
;; values, if those happened to be eager) — the implementation is
;; responsible for forcing exactly what it needs. Used for `seq`
;; and `deepseq`, which are non-strict in their second argument.
(define
  hk-mk-lazy-builtin
  (fn
    (name fn arity)
    {:type "builtin" :name name :fn fn :arity arity :lazy true :collected (list)}))

;; ── Apply a function value to one argument ──────────────────
(define
  hk-apply
  (fn
    (f arg)
    (let ((f (hk-force f)))
      (cond
        ((not (dict? f))
          (raise (str "apply: not a function value: " f)))
        ((= (get f "type") "fn")
          (cond
            ((= (get f "kind") "closure") (hk-apply-closure f arg))
            ((= (get f "kind") "multi") (hk-apply-multi f arg))
            (:else (raise "apply: unknown fn kind"))))
        ((= (get f "type") "con-partial") (hk-apply-con-partial f arg))
        ((= (get f "type") "builtin") (hk-apply-builtin f arg))
        (:else (raise "apply: not a function dict"))))))

(define
  hk-apply-closure
  (fn
    (cl arg)
    (let
      ((params (get cl "params"))
        (body (get cl "body"))
        (env (get cl "env")))
      (cond
        ((empty? params) (raise "apply-closure: no params"))
        (:else
          (let
            ((p1 (first params)) (rest-p (rest params)))
            (let
              ((env-after (hk-match p1 arg env)))
              (cond
                ((nil? env-after)
                  (raise "pattern match failure in lambda"))
                ((empty? rest-p) (hk-eval body env-after))
                (:else
                  (hk-mk-closure rest-p body env-after))))))))))

(define
  hk-apply-multi
  (fn
    (mf arg)
    (let
      ((arity (get mf "arity"))
        (clauses (get mf "clauses"))
        (env (get mf "env"))
        (collected (append (get mf "collected") (list arg))))
      (cond
        ((< (len collected) arity)
          (assoc mf "collected" collected))
        (:else (hk-dispatch-multi clauses collected env))))))

(define
  hk-dispatch-multi
  (fn
    (clauses args env)
    (cond
      ((empty? clauses)
        (raise "non-exhaustive patterns in function definition"))
      (:else
        (let
          ((c (first clauses)))
          (let
            ((pats (first c)) (body (first (rest c))))
            (let
              ((env-after (hk-match-args pats args env)))
              (cond
                ((nil? env-after)
                  (hk-dispatch-multi (rest clauses) args env))
                (:else (hk-eval body env-after))))))))))

(define
  hk-match-args
  (fn
    (pats args env)
    (cond
      ((empty? pats) env)
      (:else
        (let
          ((res (hk-match (first pats) (first args) env)))
          (cond
            ((nil? res) nil)
            (:else
              (hk-match-args (rest pats) (rest args) res))))))))

(define
  hk-apply-con-partial
  (fn
    (cp arg)
    (let
      ((name (get cp "name"))
        (arity (get cp "arity"))
        (args (append (get cp "args") (list arg))))
      (cond
        ((= (len args) arity) (hk-mk-con name args))
        (:else (assoc cp "args" args))))))

(define
  hk-apply-builtin
  (fn
    (b arg)
    (let
      ((arity (get b "arity"))
        (collected (append (get b "collected") (list arg))))
      (cond
        ((< (len collected) arity)
          (assoc b "collected" collected))
        (:else
          ;; Strict built-ins force every collected arg before
          ;; calling. Lazy ones (`seq`, `deepseq`) receive the raw
          ;; thunks so they can choose what to force.
          (cond
            ((get b "lazy") (apply (get b "fn") collected))
            (:else
              (apply
                (get b "fn")
                (map hk-force collected)))))))))

;; ── Bool helpers (Bool values are tagged conses) ────────────
(define
  hk-truthy?
  (fn
    (v)
    (and (list? v) (not (empty? v)) (= (first v) "True"))))

(define hk-true (hk-mk-con "True" (list)))
(define hk-false (hk-mk-con "False" (list)))
(define hk-of-bool (fn (b) (if b hk-true hk-false)))

;; ── Core eval ───────────────────────────────────────────────
(define
  hk-eval
  (fn
    (node env)
    (cond
      ((not (list? node)) (raise (str "eval: not a list: " node)))
      ((empty? node) (raise "eval: empty list node"))
      (:else
        (let
          ((tag (first node)))
          (cond
            ((= tag "int") (nth node 1))
            ((= tag "float") (nth node 1))
            ((= tag "string") (nth node 1))
            ((= tag "char") (nth node 1))
            ((= tag "var") (hk-eval-var (nth node 1) env))
            ((= tag "con") (hk-eval-con-ref (nth node 1)))
            ((= tag "neg")
              (- 0 (hk-force (hk-eval (nth node 1) env))))
            ((= tag "if") (hk-eval-if node env))
            ((= tag "let") (hk-eval-let (nth node 1) (nth node 2) env))
            ((= tag "lambda")
              (hk-mk-closure (nth node 1) (nth node 2) env))
            ((= tag "app")
              (hk-apply
                (hk-eval (nth node 1) env)
                (hk-mk-thunk (nth node 2) env)))
            ((= tag "op")
              (hk-eval-op
                (nth node 1)
                (nth node 2)
                (nth node 3)
                env))
            ((= tag "case")
              (hk-eval-case (nth node 1) (nth node 2) env))
            ((= tag "tuple")
              (hk-mk-tuple
                (map (fn (e) (hk-eval e env)) (nth node 1))))
            ((= tag "list")
              (hk-mk-list
                (map (fn (e) (hk-eval e env)) (nth node 1))))
            ((= tag "range")
              (let
                ((from (hk-force (hk-eval (nth node 1) env)))
                  (to (hk-force (hk-eval (nth node 2) env))))
                (hk-build-range from to 1)))
            ((= tag "range-step")
              (let
                ((from (hk-force (hk-eval (nth node 1) env)))
                  (nxt (hk-force (hk-eval (nth node 2) env)))
                  (to (hk-force (hk-eval (nth node 3) env))))
                (hk-build-range from to (- nxt from))))
            ((= tag "range-from")
              ;; [from..] = iterate (+ 1) from — uses the Prelude.
              (hk-eval
                (list
                  :app
                  (list
                    :app
                    (list :var "iterate")
                    (list
                      :sect-right
                      "+"
                      (list :int 1)))
                  (nth node 1))
                env))
            ((= tag "sect-left")
              (hk-eval-sect-left (nth node 1) (nth node 2) env))
            ((= tag "sect-right")
              (hk-eval-sect-right (nth node 1) (nth node 2) env))
            (:else
              (raise (str "eval: unknown node tag '" tag "'")))))))))

(define
  hk-eval-var
  (fn
    (name env)
    (cond
      ((has-key? env name) (get env name))
      ((hk-is-con? name) (hk-eval-con-ref name))
      (:else (raise (str "unbound variable: " name))))))

(define
  hk-eval-con-ref
  (fn
    (name)
    (let ((arity (hk-con-arity name)))
      (cond
        ((nil? arity) (raise (str "unknown constructor: " name)))
        ((= arity 0) (hk-mk-con name (list)))
        (:else
          {:type "con-partial" :name name :arity arity :args (list)})))))

(define
  hk-eval-if
  (fn
    (node env)
    (let ((cv (hk-force (hk-eval (nth node 1) env))))
      (cond
        ((hk-truthy? cv) (hk-eval (nth node 2) env))
        ((and (list? cv) (= (first cv) "False"))
          (hk-eval (nth node 3) env))
        ((= cv true) (hk-eval (nth node 2) env))
        ((= cv false) (hk-eval (nth node 3) env))
        (:else (raise "if: condition is not Bool"))))))

(define
  hk-extend-env-with-match!
  (fn
    (env match-env)
    (for-each
      (fn (k) (dict-set! env k (get match-env k)))
      (keys match-env))))

(define
  hk-eval-let-bind!
  (fn
    (b env)
    (let ((tag (first b)))
      (cond
        ((= tag "fun-clause")
          (let
            ((name (nth b 1))
              (pats (nth b 2))
              (body (nth b 3)))
            (cond
              ((empty? pats)
                (dict-set! env name (hk-eval body env)))
              (:else
                (dict-set! env name (hk-mk-closure pats body env))))))
        ((or (= tag "bind") (= tag "pat-bind"))
          (let ((pat (nth b 1)) (body (nth b 2)))
            (let ((val (hk-eval body env)))
              (let ((res (hk-match pat val env)))
                (cond
                  ((nil? res)
                    (raise "let: pattern bind failure"))
                  (:else
                    (hk-extend-env-with-match! env res)))))))
        (:else nil)))))

(define
  hk-eval-let
  (fn
    (binds body env)
    (let ((new-env (hk-dict-copy env)))
      ;; Pre-seed names for fn-clauses so closures see themselves
      ;; (mutual recursion across the whole binding group).
      (for-each
        (fn (b)
          (cond
            ((= (first b) "fun-clause")
              (dict-set! new-env (nth b 1) nil))
            ((and
                (= (first b) "bind")
                (list? (nth b 1))
                (= (first (nth b 1)) "p-var"))
              (dict-set! new-env (nth (nth b 1) 1) nil))
            (:else nil)))
        binds)
      (for-each (fn (b) (hk-eval-let-bind! b new-env)) binds)
      (hk-eval body new-env))))

(define
  hk-eval-case
  (fn
    (scrut alts env)
    (let ((sv (hk-force (hk-eval scrut env))))
      (hk-try-alts alts sv env))))

(define
  hk-try-alts
  (fn
    (alts val env)
    (cond
      ((empty? alts) (raise "case: non-exhaustive patterns"))
      (:else
        (let
          ((alt (first alts)))
          (let
            ((pat (nth alt 1)) (body (nth alt 2)))
            (let
              ((res (hk-match pat val env)))
              (cond
                ((nil? res) (hk-try-alts (rest alts) val env))
                (:else (hk-eval body res))))))))))

(define
  hk-eval-op
  (fn
    (op left right env)
    (cond
      ;; Cons is non-strict in both args: build a cons cell whose
      ;; head and tail are deferred. This is what makes `repeat x =
      ;; x : repeat x` and `fibs = 0 : 1 : zipWith (+) fibs (tail
      ;; fibs)` terminate.
      ((= op ":")
        (hk-mk-cons
          (hk-mk-thunk left env)
          (hk-mk-thunk right env)))
      (:else
        (let
          ((lv (hk-force (hk-eval left env)))
            (rv (hk-force (hk-eval right env))))
          (hk-binop op lv rv))))))

(define
  hk-list-append
  (fn
    (a b)
    (cond
      ((and (list? a) (= (first a) "[]")) b)
      ((and (list? a) (= (first a) ":"))
        (hk-mk-cons (nth a 1) (hk-list-append (nth a 2) b)))
      (:else (raise "++: not a list")))))

;; Eager finite-range spine — handles [from..to] and [from,next..to].
;; Step direction is governed by the sign of `step`; when step > 0 we
;; stop at to; when step < 0 we stop at to going down.
(define
  hk-build-range
  (fn
    (from to step)
    (cond
      ((and (> step 0) (> from to)) (hk-mk-nil))
      ((and (< step 0) (< from to)) (hk-mk-nil))
      ((= step 0) (hk-mk-nil))
      (:else
        (hk-mk-cons from (hk-build-range (+ from step) to step))))))

(define
  hk-binop
  (fn
    (op lv rv)
    (cond
      ((= op "+") (+ lv rv))
      ((= op "-") (- lv rv))
      ((= op "*") (* lv rv))
      ((= op "/") (/ lv rv))
      ((= op "==") (hk-of-bool (= lv rv)))
      ((= op "/=") (hk-of-bool (not (= lv rv))))
      ((= op "<") (hk-of-bool (< lv rv)))
      ((= op "<=") (hk-of-bool (<= lv rv)))
      ((= op ">") (hk-of-bool (> lv rv)))
      ((= op ">=") (hk-of-bool (>= lv rv)))
      ((= op "&&") (hk-of-bool (and (hk-truthy? lv) (hk-truthy? rv))))
      ((= op "||") (hk-of-bool (or (hk-truthy? lv) (hk-truthy? rv))))
      ((= op ":") (hk-mk-cons lv rv))
      ((= op "++") (hk-list-append lv rv))
      ((= op "mod") (mod lv rv))
      ((= op "div") (floor (/ lv rv)))
      ((= op "rem") (mod lv rv))
      ((= op "quot") (truncate (/ lv rv)))
      (:else (raise (str "unknown operator: " op))))))

(define
  hk-eval-sect-left
  (fn
    (op e env)
    ;; (e op) = \x -> e op x — bind e once, defer the operator call.
    (let ((ev (hk-eval e env)))
      (let ((cenv (hk-dict-copy env)))
        (dict-set! cenv "__hk-sect-l" ev)
        (hk-mk-closure
          (list (list :p-var "__hk-sect-x"))
          (list
            :op
            op
            (list :var "__hk-sect-l")
            (list :var "__hk-sect-x"))
          cenv)))))

(define
  hk-eval-sect-right
  (fn
    (op e env)
    (let ((ev (hk-eval e env)))
      (let ((cenv (hk-dict-copy env)))
        (dict-set! cenv "__hk-sect-r" ev)
        (hk-mk-closure
          (list (list :p-var "__hk-sect-x"))
          (list
            :op
            op
            (list :var "__hk-sect-x")
            (list :var "__hk-sect-r"))
          cenv)))))

;; ── Top-level program evaluation ────────────────────────────
;; Operator-as-value built-ins — let `(+)`, `(*)`, etc. work as
;; first-class functions for `zipWith (+)` and friends. Strict in
;; both args (built-ins are forced via hk-apply-builtin).
(define
  hk-make-binop-builtin
  (fn
    (name op-name)
    (hk-mk-builtin
      name
      (fn (a b) (hk-binop op-name a b))
      2)))

;; Inline Prelude source — loaded into the initial env so simple
;; programs can use `head`, `take`, `repeat`, etc. without each
;; user file redefining them. The Prelude itself uses lazy `:` for
;; the recursive list-building functions.
(define
  hk-prelude-src
  "head (x:_) = x
tail (_:xs) = xs
fst (a, _) = a
snd (_, b) = b
take 0 _ = []
take _ [] = []
take n (x:xs) = x : take (n - 1) xs
drop 0 xs = xs
drop _ [] = []
drop n (_:xs) = drop (n - 1) xs
repeat x = x : repeat x
iterate f x = x : iterate f (f x)
length [] = 0
length (_:xs) = 1 + length xs
map _ [] = []
map f (x:xs) = f x : map f xs
filter _ [] = []
filter p (x:xs) = if p x then x : filter p xs else filter p xs
zipWith _ [] _ = []
zipWith _ _ [] = []
zipWith f (x:xs) (y:ys) = f x y : zipWith f xs ys
fibs = 0 : 1 : zipWith plus fibs (tail fibs)
plus a b = a + b
")

(define
  hk-load-into!
  (fn
    (env src)
    (let ((ast (hk-core src)))
      (hk-register-program! ast)
      (let
        ((decls
            (cond
              ((= (first ast) "program") (nth ast 1))
              ((= (first ast) "module") (nth ast 4))
              (:else (list)))))
        (hk-bind-decls! env decls)))))

(define
  hk-init-env
  (fn
    ()
    (let ((env (dict)))
      (dict-set! env "otherwise" hk-true)
      (dict-set!
        env
        "error"
        (hk-mk-builtin
          "error"
          (fn (msg) (raise (str "*** Exception: " msg)))
          1))
      (dict-set!
        env
        "not"
        (hk-mk-builtin
          "not"
          (fn (b) (hk-of-bool (not (hk-truthy? b))))
          1))
      (dict-set!
        env
        "id"
        (hk-mk-builtin "id" (fn (x) x) 1))
      ;; `seq a b` — strict in `a`, lazy in `b`. Forces `a` to WHNF
      ;; and returns `b` unchanged (still a thunk if it was one).
      (dict-set!
        env
        "seq"
        (hk-mk-lazy-builtin
          "seq"
          (fn (a b) (do (hk-force a) b))
          2))
      ;; `deepseq a b` — like seq but forces `a` to normal form.
      (dict-set!
        env
        "deepseq"
        (hk-mk-lazy-builtin
          "deepseq"
          (fn (a b) (do (hk-deep-force a) b))
          2))
      ;; ── Stub IO monad ─────────────────────────────────────
      ;; IO actions are tagged values `("IO" payload)`; `>>=` and
      ;; `>>` chain them. Lazy in the action arguments so do-blocks
      ;; can be deeply structured without forcing the whole chain
      ;; up front.
      (dict-set!
        env
        "return"
        (hk-mk-lazy-builtin
          "return"
          (fn (x) (list "IO" x))
          1))
      (dict-set!
        env
        ">>="
        (hk-mk-lazy-builtin
          ">>="
          (fn (m f)
            (let ((io-val (hk-force m)))
              (cond
                ((and
                    (list? io-val)
                    (= (first io-val) "IO"))
                  (hk-apply (hk-force f) (nth io-val 1)))
                (:else
                  (raise "(>>=): left side is not an IO action")))))
          2))
      (dict-set!
        env
        ">>"
        (hk-mk-lazy-builtin
          ">>"
          (fn (m n)
            (let ((io-val (hk-force m)))
              (cond
                ((and
                    (list? io-val)
                    (= (first io-val) "IO"))
                  (hk-force n))
                (:else
                  (raise "(>>): left side is not an IO action")))))
          2))
      ;; Operators as first-class values
      (dict-set! env "+" (hk-make-binop-builtin "+" "+"))
      (dict-set! env "-" (hk-make-binop-builtin "-" "-"))
      (dict-set! env "*" (hk-make-binop-builtin "*" "*"))
      (dict-set! env "/" (hk-make-binop-builtin "/" "/"))
      (dict-set! env "==" (hk-make-binop-builtin "==" "=="))
      (dict-set! env "/=" (hk-make-binop-builtin "/=" "/="))
      (dict-set! env "<" (hk-make-binop-builtin "<" "<"))
      (dict-set! env "<=" (hk-make-binop-builtin "<=" "<="))
      (dict-set! env ">" (hk-make-binop-builtin ">" ">"))
      (dict-set! env ">=" (hk-make-binop-builtin ">=" ">="))
      (dict-set! env "&&" (hk-make-binop-builtin "&&" "&&"))
      (dict-set! env "||" (hk-make-binop-builtin "||" "||"))
      (dict-set! env "++" (hk-make-binop-builtin "++" "++"))
      (dict-set! env "mod" (hk-make-binop-builtin "mod" "mod"))
      (dict-set! env "div" (hk-make-binop-builtin "div" "div"))
      (dict-set! env "rem" (hk-make-binop-builtin "rem" "rem"))
      (dict-set! env "quot" (hk-make-binop-builtin "quot" "quot"))
      (hk-load-into! env hk-prelude-src)
      env)))

(define
  hk-bind-decls!
  (fn
    (env decls)
    (let
      ((groups (dict))
        (group-order (list))
        (pat-binds (list)))
      ;; Pass 1: collect fun-clause groups by name; track first-seen
      ;; order so pass 3 can evaluate 0-arity bodies in source order
      ;; (forward references to other 0-arity definitions still need
      ;; the earlier name to be bound first).
      (for-each
        (fn (d)
          (cond
            ((= (first d) "fun-clause")
              (let
                ((name (nth d 1)))
                (when (not (has-key? groups name))
                  (append! group-order name))
                (dict-set!
                  groups
                  name
                  (append
                    (if
                      (has-key? groups name)
                      (get groups name)
                      (list))
                    (list (list (nth d 2) (nth d 3)))))
                (when
                  (not (has-key? env name))
                  (dict-set! env name nil))))
            ((or (= (first d) "bind") (= (first d) "pat-bind"))
              (append! pat-binds d))
            (:else nil)))
        decls)
      ;; Pass 2: install multifuns (arity > 0) — order doesn't matter
      ;; because they're closures; collect 0-arity names in source
      ;; order for pass 3.
      (let ((zero-arity (list)))
        (for-each
          (fn (name)
            (let ((clauses (get groups name)))
              (let ((arity (len (first (first clauses)))))
                (cond
                  ((> arity 0)
                    (dict-set!
                      env
                      name
                      (hk-mk-multifun arity clauses env)))
                  (:else (append! zero-arity name))))))
          group-order)
        ;; Pass 3: evaluate 0-arity bodies and pat-binds in source
        ;; order — forward references to a later 0-arity name will
        ;; still see its placeholder (nil) and fail noisily, but the
        ;; common case of a top-down program works.
        (for-each
          (fn (name)
            (let ((clauses (get groups name)))
              (dict-set!
                env
                name
                (hk-eval (first (rest (first clauses))) env))))
          zero-arity)
        (for-each
          (fn (d)
            (let ((pat (nth d 1)) (body (nth d 2)))
              (let ((val (hk-eval body env)))
                (let ((res (hk-match pat val env)))
                  (cond
                    ((nil? res)
                      (raise "top-level pattern bind failure"))
                    (:else (hk-extend-env-with-match! env res)))))))
          pat-binds))
      env)))

(define
  hk-eval-program
  (fn
    (ast)
    (cond
      ((nil? ast) (raise "eval-program: nil ast"))
      ((not (list? ast)) (raise "eval-program: not a list"))
      (:else
        (do
          (hk-register-program! ast)
          (let ((env (hk-init-env)))
            (let
              ((decls
                  (cond
                    ((= (first ast) "program") (nth ast 1))
                    ((= (first ast) "module") (nth ast 4))
                    (:else (raise "eval-program: bad shape")))))
              (hk-bind-decls! env decls))))))))

;; ── Source-level convenience ────────────────────────────────
(define
  hk-run
  (fn
    (src)
    (let ((env (hk-eval-program (hk-core src))))
      (cond
        ((has-key? env "main") (get env "main"))
        (:else env)))))

(define
  hk-eval-expr-source
  (fn
    (src)
    (hk-deep-force (hk-eval (hk-core-expr src) (hk-init-env)))))